@@ -37,6 +37,168 @@ inline unsigned char lookup_code_index(const float* codebook, float value) {
3737
3838} // namespace
3939
40+ #if defined(_M_ARM64) || defined(__aarch64__)
41+ #include < arm_neon.h>
42+
43+ // ARM NEON NF4 lookup table (16 float values indexed by 4-bit code)
44+ static inline void neon_nf4_lut (float32x4_t lut[4 ]) {
45+ // Indices 0-15 map to NF4 dequantized values
46+ // Order: index 0 = -1.0, index 1 = -0.6962, ..., index 7 = 0.0, ..., index 15 = 1.0
47+ static const float nf4_values[16 ] = {
48+ -1 .0f , -0 .6961928009986877f , -0 .5250730514526367f , -0 .39491748809814453f ,
49+ -0 .28444138169288635f , -0 .18477343022823334f , -0 .09105003625154495f , 0 .0f ,
50+ 0 .07958029955625534f , 0 .16093020141124725f , 0 .24611230194568634f , 0 .33791524171829224f ,
51+ 0 .44070982933044434f , 0 .5626170039176941f , 0 .7229568362236023f , 1 .0f
52+ };
53+ lut[0 ] = vld1q_f32 (nf4_values);
54+ lut[1 ] = vld1q_f32 (nf4_values + 4 );
55+ lut[2 ] = vld1q_f32 (nf4_values + 8 );
56+ lut[3 ] = vld1q_f32 (nf4_values + 12 );
57+ }
58+
59+ // ARM NEON FP4 lookup table
60+ static inline void neon_fp4_lut (float32x4_t lut[4 ]) {
61+ static const float fp4_values[16 ] = {
62+ 0 .0f , 5 .208333333e-03f , 0 .66666667f , 1 .0f ,
63+ 0 .33333333f , 0 .5f , 0 .16666667f , 0 .25f ,
64+ 0 .0f , -5 .208333333e-03f , -0 .66666667f , -1 .0f ,
65+ -0 .33333333f , -0 .5f , -0 .16666667f , -0 .25f
66+ };
67+ lut[0 ] = vld1q_f32 (fp4_values);
68+ lut[1 ] = vld1q_f32 (fp4_values + 4 );
69+ lut[2 ] = vld1q_f32 (fp4_values + 8 );
70+ lut[3 ] = vld1q_f32 (fp4_values + 12 );
71+ }
72+
73+ // Efficient NEON float LUT lookup using indexed lane extraction
74+ // The LUT has 16 float entries stored as a flat array for direct indexing
75+ static inline float neon_lut_lookup_flat (const float * flat_lut, uint8_t idx) {
76+ return flat_lut[idx];
77+ }
78+
79+ // Vectorized 4-bit dequantization: process 8 packed bytes = 16 output values
80+ // Each byte contains two 4-bit values: high nibble first, low nibble second
81+ static inline void neon_dequant_4bit_16values (
82+ const uint8_t * packed, float scale, const float32x4_t lut[4 ], float * out
83+ ) {
84+ // Load 8 bytes = 16 x 4-bit values
85+ uint8x8_t raw = vld1_u8 (packed);
86+
87+ // Extract high and low nibbles
88+ uint8x8_t mask4 = vdup_n_u8 (0x0F );
89+ uint8x8_t lo_nibbles = vand_u8 (raw, mask4); // low nibble (second value)
90+ uint8x8_t hi_nibbles = vshr_n_u8 (raw, 4 ); // high nibble (first value)
91+
92+ // Interleave hi/lo into 16-element index array for output ordering
93+ // output[2*i] = hi_nibble[i], output[2*i+1] = lo_nibble[i]
94+ uint8x8x2_t interleaved = vzip_u8 (hi_nibbles, lo_nibbles);
95+ // interleaved.val[0] has elements 0-7, interleaved.val[1] has elements 8-15
96+ uint8x16_t indices = vcombine_u8 (interleaved.val [0 ], interleaved.val [1 ]);
97+
98+ // Use flat LUT for fast indexed access
99+ // Store LUT as flat float array on stack (likely in L1 cache)
100+ float flat_lut[16 ];
101+ vst1q_f32 (flat_lut, lut[0 ]);
102+ vst1q_f32 (flat_lut + 4 , lut[1 ]);
103+ vst1q_f32 (flat_lut + 8 , lut[2 ]);
104+ vst1q_f32 (flat_lut + 12 , lut[3 ]);
105+
106+ // Extract indices and do lookups in groups of 4 for NEON multiply
107+ uint8_t idx_arr[16 ];
108+ vst1q_u8 (idx_arr, indices);
109+
110+ float32x4_t vscale = vdupq_n_f32 (scale);
111+
112+ // Process 4 values at a time with NEON - load from temp buffer
113+ float tmp_vals[16 ];
114+ tmp_vals[0 ] = flat_lut[idx_arr[0 ]]; tmp_vals[1 ] = flat_lut[idx_arr[1 ]];
115+ tmp_vals[2 ] = flat_lut[idx_arr[2 ]]; tmp_vals[3 ] = flat_lut[idx_arr[3 ]];
116+ tmp_vals[4 ] = flat_lut[idx_arr[4 ]]; tmp_vals[5 ] = flat_lut[idx_arr[5 ]];
117+ tmp_vals[6 ] = flat_lut[idx_arr[6 ]]; tmp_vals[7 ] = flat_lut[idx_arr[7 ]];
118+ tmp_vals[8 ] = flat_lut[idx_arr[8 ]]; tmp_vals[9 ] = flat_lut[idx_arr[9 ]];
119+ tmp_vals[10 ] = flat_lut[idx_arr[10 ]]; tmp_vals[11 ] = flat_lut[idx_arr[11 ]];
120+ tmp_vals[12 ] = flat_lut[idx_arr[12 ]]; tmp_vals[13 ] = flat_lut[idx_arr[13 ]];
121+ tmp_vals[14 ] = flat_lut[idx_arr[14 ]]; tmp_vals[15 ] = flat_lut[idx_arr[15 ]];
122+ float32x4_t v0 = vld1q_f32 (tmp_vals);
123+ float32x4_t v1 = vld1q_f32 (tmp_vals + 4 );
124+ float32x4_t v2 = vld1q_f32 (tmp_vals + 8 );
125+ float32x4_t v3 = vld1q_f32 (tmp_vals + 12 );
126+
127+ vst1q_f32 (out, vmulq_f32 (v0, vscale));
128+ vst1q_f32 (out + 4 , vmulq_f32 (v1, vscale));
129+ vst1q_f32 (out + 8 , vmulq_f32 (v2, vscale));
130+ vst1q_f32 (out + 12 , vmulq_f32 (v3, vscale));
131+ }
132+
133+ // NEON-optimized BF16 to float conversion (4 values at a time)
134+ static inline float32x4_t neon_bf16x4_to_f32 (const bf16_t * src) {
135+ // BF16 is upper 16 bits of float32, so shift left by 16
136+ uint16x4_t raw = vld1_u16 (reinterpret_cast <const uint16_t *>(src));
137+ uint32x4_t wide = vshll_n_u16 (raw, 16 );
138+ return vreinterpretq_f32_u32 (wide);
139+ }
140+
141+ // NEON-optimized float to BF16 conversion (4 values at a time, with rounding)
142+ static inline void neon_f32_to_bf16x4 (const float32x4_t src, bf16_t * dst) {
143+ uint32x4_t bits = vreinterpretq_u32_f32 (src);
144+ // Round to nearest even: add 0x7FFF + ((bits >> 16) & 1)
145+ uint32x4_t lsb = vshrq_n_u32 (bits, 16 );
146+ lsb = vandq_u32 (lsb, vdupq_n_u32 (1 ));
147+ uint32x4_t rounding = vaddq_u32 (vdupq_n_u32 (0x7FFF ), lsb);
148+ bits = vaddq_u32 (bits, rounding);
149+ // Extract upper 16 bits
150+ uint16x4_t result = vshrn_n_u32 (bits, 16 );
151+ vst1_u16 (reinterpret_cast <uint16_t *>(dst), result);
152+ }
153+
154+ // NEON-optimized float to FP16 conversion (4 values at a time)
155+ static inline void neon_f32_to_fp16x4 (const float32x4_t src, fp16_t * dst) {
156+ // ARM64 has native FP16 conversion
157+ float16x4_t half = vcvt_f16_f32 (src);
158+ vst1_u16 (reinterpret_cast <uint16_t *>(dst), vreinterpret_u16_f16 (half));
159+ }
160+
161+ // NEON-optimized absmax computation for a block of float32
162+ static inline float neon_absmax_f32 (const float * data, long long n) {
163+ float32x4_t vmax = vdupq_n_f32 (0 .0f );
164+ long long i = 0 ;
165+ // Process 16 elements per iteration for better throughput
166+ for (; i + 16 <= n; i += 16 ) {
167+ float32x4_t v0 = vabsq_f32 (vld1q_f32 (data + i));
168+ float32x4_t v1 = vabsq_f32 (vld1q_f32 (data + i + 4 ));
169+ float32x4_t v2 = vabsq_f32 (vld1q_f32 (data + i + 8 ));
170+ float32x4_t v3 = vabsq_f32 (vld1q_f32 (data + i + 12 ));
171+ vmax = vmaxq_f32 (vmax, vmaxq_f32 (vmaxq_f32 (v0, v1), vmaxq_f32 (v2, v3)));
172+ }
173+ for (; i + 4 <= n; i += 4 ) {
174+ float32x4_t v = vld1q_f32 (data + i);
175+ vmax = vmaxq_f32 (vmax, vabsq_f32 (v));
176+ }
177+ // Horizontal max
178+ float result = vmaxvq_f32 (vmax);
179+ // Handle remainder
180+ for (; i < n; ++i) {
181+ result = std::max (result, std::fabs (data[i]));
182+ }
183+ return result;
184+ }
185+
186+ // NEON-optimized norm_to_lut_index for 4 float values at a time
187+ // Maps [-1, 1] → [0, 65535]
188+ static inline uint16x4_t neon_norm_to_lut_index_x4 (float32x4_t vals) {
189+ // clamp to [-1, 1]
190+ vals = vmaxq_f32 (vals, vdupq_n_f32 (-1 .0f ));
191+ vals = vminq_f32 (vals, vdupq_n_f32 (1 .0f ));
192+ // (val + 1.0) * 0.5 * 65535 + 0.5
193+ float32x4_t result = vmlaq_f32 (vdupq_n_f32 (0 .5f ),
194+ vaddq_f32 (vals, vdupq_n_f32 (1 .0f )),
195+ vdupq_n_f32 (0 .5f * 65535 .0f ));
196+ uint32x4_t u32 = vcvtq_u32_f32 (result);
197+ return vmovn_u32 (u32 );
198+ }
199+
200+ #endif // _M_ARM64 || __aarch64__
201+
40202#if defined(__AVX512F__)
41203#include < immintrin.h>
42204
@@ -94,6 +256,57 @@ void dequantizeBlockwise4bitCpu(
94256 if (blocksize <= 0 || m < 0 || n <= 0 )
95257 return ;
96258
259+ #if defined(_M_ARM64) || defined(__aarch64__)
260+ {
261+ long long dim_0 = m;
262+ long long dim_1 = n;
263+ long long input_dim_1 = dim_1 >> 1 ;
264+ long long absmax_dim_1 = dim_1 / blocksize;
265+ // NEON path: process 16 output values at a time (8 packed bytes)
266+ // Only use when blocksize evenly divides dim_1 to ensure correct scale indexing
267+ constexpr long long VEC_LEN = 16 ;
268+ if (dim_1 % VEC_LEN == 0 && blocksize >= VEC_LEN && (dim_1 % blocksize == 0 )) {
269+ float32x4_t lut[4 ];
270+ if constexpr (DATA_TYPE == 1 ) {
271+ neon_fp4_lut (lut);
272+ } else {
273+ neon_nf4_lut (lut);
274+ }
275+ constexpr long long k_step = VEC_LEN / 2 ; // 8 bytes per iteration
276+ BNB_OMP_PARALLEL_FOR
277+ for (long long block_idx = 0 ; block_idx < dim_0; ++block_idx) {
278+ for (long long k = 0 ; k < input_dim_1; k += k_step) {
279+ long long scale_idx = k * 2 / blocksize;
280+ float scale = absmax[block_idx * absmax_dim_1 + scale_idx];
281+ const uint8_t * p = &A[block_idx * input_dim_1 + k];
282+
283+ // Dequantize 16 values into a temp float buffer
284+ float tmp_f32[16 ];
285+ neon_dequant_4bit_16values (p, scale, lut, tmp_f32);
286+
287+ // Store results (convert to output type using NEON)
288+ T* pout = &out[block_idx * dim_1 + k * 2 ];
289+ if constexpr (std::is_same<T, float >()) {
290+ // Direct copy - already float
291+ std::memcpy (pout, tmp_f32, 16 * sizeof (float ));
292+ } else if constexpr (std::is_same<T, bf16_t >()) {
293+ neon_f32_to_bf16x4 (vld1q_f32 (tmp_f32), pout);
294+ neon_f32_to_bf16x4 (vld1q_f32 (tmp_f32 + 4 ), pout + 4 );
295+ neon_f32_to_bf16x4 (vld1q_f32 (tmp_f32 + 8 ), pout + 8 );
296+ neon_f32_to_bf16x4 (vld1q_f32 (tmp_f32 + 12 ), pout + 12 );
297+ } else if constexpr (std::is_same<T, fp16_t >()) {
298+ neon_f32_to_fp16x4 (vld1q_f32 (tmp_f32), pout);
299+ neon_f32_to_fp16x4 (vld1q_f32 (tmp_f32 + 4 ), pout + 4 );
300+ neon_f32_to_fp16x4 (vld1q_f32 (tmp_f32 + 8 ), pout + 8 );
301+ neon_f32_to_fp16x4 (vld1q_f32 (tmp_f32 + 12 ), pout + 12 );
302+ }
303+ }
304+ }
305+ return ;
306+ }
307+ }
308+ #endif // _M_ARM64 || __aarch64__
309+
97310#if defined(__AVX512F__)
98311 if (has_avx512f ()) {
99312 long long dim_0 = m;
@@ -305,19 +518,29 @@ void quantize_cpu_impl(float* code, const T* A, float* absmax, unsigned char* ou
305518 for (long long b = 0 ; b < num_blocks; ++b) {
306519 const long long block_start = b * blocksize;
307520 const long long block_end = std::min (block_start + blocksize, n);
521+ const long long block_len = block_end - block_start;
308522
309523 // Compute absmax for this block
310524 float absmax_block = 0 .0f ;
311- for (long long i = block_start; i < block_end; ++i) {
312- float val;
313- if constexpr (std::is_same<T, float >::value) {
314- val = A[i];
315- } else if constexpr (std::is_same<T, bf16_t >::value) {
316- val = bf16_to_float (A[i].v );
317- } else if constexpr (std::is_same<T, fp16_t >::value) {
318- val = fp16_to_float (A[i].v );
525+
526+ #if defined(_M_ARM64) || defined(__aarch64__)
527+ if constexpr (std::is_same<T, float >::value) {
528+ // Use NEON-optimized absmax for float32
529+ absmax_block = neon_absmax_f32 (reinterpret_cast <const float *>(A + block_start), block_len);
530+ } else
531+ #endif
532+ {
533+ for (long long i = block_start; i < block_end; ++i) {
534+ float val;
535+ if constexpr (std::is_same<T, float >::value) {
536+ val = A[i];
537+ } else if constexpr (std::is_same<T, bf16_t >::value) {
538+ val = bf16_to_float (A[i].v );
539+ } else if constexpr (std::is_same<T, fp16_t >::value) {
540+ val = fp16_to_float (A[i].v );
541+ }
542+ absmax_block = std::max (absmax_block, std::fabs (val));
319543 }
320- absmax_block = std::max (absmax_block, std::fabs (val));
321544 }
322545
323546 absmax[b] = absmax_block;
@@ -328,17 +551,42 @@ void quantize_cpu_impl(float* code, const T* A, float* absmax, unsigned char* ou
328551 }
329552
330553 const float inv_absmax = 1 .0f / absmax_block;
331- for (long long i = block_start; i < block_end; ++i) {
332- float val;
333- if constexpr (std::is_same<T, float >::value) {
334- val = A[i];
335- } else if constexpr (std::is_same<T, bf16_t >::value) {
336- val = bf16_to_float (A[i].v );
337- } else if constexpr (std::is_same<T, fp16_t >::value) {
338- val = fp16_to_float (A[i].v );
554+
555+ #if defined(_M_ARM64) || defined(__aarch64__)
556+ if constexpr (std::is_same<T, float >::value) {
557+ // NEON-optimized normalize + LUT index for float32
558+ const float * src = A + block_start;
559+ long long i = 0 ;
560+ float32x4_t vinv = vdupq_n_f32 (inv_absmax);
561+ for (; i + 4 <= block_len; i += 4 ) {
562+ float32x4_t v = vmulq_f32 (vld1q_f32 (src + i), vinv);
563+ uint16x4_t indices = neon_norm_to_lut_index_x4 (v);
564+ uint16_t idx_arr[4 ];
565+ vst1_u16 (idx_arr, indices);
566+ out[block_start + i] = lut[idx_arr[0 ]];
567+ out[block_start + i + 1 ] = lut[idx_arr[1 ]];
568+ out[block_start + i + 2 ] = lut[idx_arr[2 ]];
569+ out[block_start + i + 3 ] = lut[idx_arr[3 ]];
570+ }
571+ for (; i < block_len; ++i) {
572+ float normed_value = src[i] * inv_absmax;
573+ out[block_start + i] = lut[norm_to_lut_index (normed_value)];
574+ }
575+ } else
576+ #endif
577+ {
578+ for (long long i = block_start; i < block_end; ++i) {
579+ float val;
580+ if constexpr (std::is_same<T, float >::value) {
581+ val = A[i];
582+ } else if constexpr (std::is_same<T, bf16_t >::value) {
583+ val = bf16_to_float (A[i].v );
584+ } else if constexpr (std::is_same<T, fp16_t >::value) {
585+ val = fp16_to_float (A[i].v );
586+ }
587+ float normed_value = val * inv_absmax;
588+ out[i] = lut[norm_to_lut_index (normed_value)];
339589 }
340- float normed_value = val * inv_absmax;
341- out[i] = lut[norm_to_lut_index (normed_value)];
342590 }
343591 }
344592}
0 commit comments