diff --git a/mlx/backend/cpu/simd/base_simd.h b/mlx/backend/cpu/simd/base_simd.h index 37c3c755..775f5dfd 100644 --- a/mlx/backend/cpu/simd/base_simd.h +++ b/mlx/backend/cpu/simd/base_simd.h @@ -29,6 +29,14 @@ struct Simd { Simd(Simd v) : value(v.value) {} template Simd(U v) : value(v) {} + + T operator[](int) const { + return value; + } + + T& operator[](int) { + return value; + } }; template diff --git a/mlx/backend/cpu/unary_ops.h b/mlx/backend/cpu/unary_ops.h index f207f28a..f441e88b 100644 --- a/mlx/backend/cpu/unary_ops.h +++ b/mlx/backend/cpu/unary_ops.h @@ -155,11 +155,18 @@ struct FromFP8 { template Simd operator()(Simd x) { auto v = Simd(x & 127) << 7; - auto converted = *(Simd*)(&v); - converted = converted * 256.0; + Simd out; + if constexpr (simd::max_size >= N) { + auto converted = *(Simd*)(&v); + out = converted * 256.0; + } else { + for (int i = 0; i < N; ++i) { + auto converted = *(float16_t*)(&v[i]); + out[i] = converted * 256.0; + } + } auto sign = Simd(x & 128); - Simd out = select(sign, -converted, converted); - return out; + return select(sign, -out, out); } float operator()(uint8_t x) { return (*this)(Simd(x)).value;