| 303 | // DATA_TYPE: 1 = FP4, 2 = NF4 |
| 304 | template <typename T, int DATA_TYPE> |
| 305 | void dequantizeBlockwise4bitCpu( |
| 306 | unsigned char* A, const float* absmax, T* out, long long blocksize, long long m, long long n |
| 307 | ) { |
| 308 | static_assert(DATA_TYPE == 1 || DATA_TYPE == 2, "dequantizeBlockwise4bitCpu called with non 4-bit DATA_TYPE"); |
| 309 | if (blocksize <= 0 || m < 0 || n <= 0) |
| 310 | return; |
| 311 | |
| 312 | #if defined(_M_ARM64) || defined(__aarch64__) |
| 313 | // n % blocksize == 0: absmax is organized by flat element blocks; row and block |
| 314 | // boundaries must align or the 2D absmax indexing gives wrong scale values. |
| 315 | if (n % blocksize == 0) { |
| 316 | long long dim_0 = m; |
| 317 | long long dim_1 = n; |
| 318 | long long input_dim_1 = dim_1 >> 1; |
| 319 | long long absmax_dim_1 = dim_1 / blocksize; |
| 320 | float32x4_t neon_lut[4]; |
| 321 | if constexpr (DATA_TYPE == 1) { |
| 322 | neon_fp4_lut(neon_lut); |
| 323 | } else { |
| 324 | neon_nf4_lut(neon_lut); |
| 325 | } |
| 326 | constexpr long long k_step = 8; // 8 packed bytes = 16 output values |
| 327 | BNB_OMP_PARALLEL_FOR |
| 328 | for (long long block_idx = 0; block_idx < dim_0; ++block_idx) { |
| 329 | for (long long k = 0; k < input_dim_1; k += k_step) { |
| 330 | long long scale_idx = k * 2 / blocksize; |
| 331 | float scale = absmax[block_idx * absmax_dim_1 + scale_idx]; |
| 332 | const uint8_t* p = &A[block_idx * input_dim_1 + k]; |
| 333 | float tmp_f32[16]; |
| 334 | neon_dequant_4bit_16values(p, scale, neon_lut, tmp_f32); |
| 335 | T* pout = &out[block_idx * dim_1 + k * 2]; |
| 336 | if constexpr (std::is_same<T, float>()) { |
| 337 | std::memcpy(pout, tmp_f32, 16 * sizeof(float)); |
| 338 | } else if constexpr (std::is_same<T, bf16_t>()) { |
| 339 | neon_f32_to_bf16x4(vld1q_f32(tmp_f32), pout); |
| 340 | neon_f32_to_bf16x4(vld1q_f32(tmp_f32 + 4), pout + 4); |
| 341 | neon_f32_to_bf16x4(vld1q_f32(tmp_f32 + 8), pout + 8); |
| 342 | neon_f32_to_bf16x4(vld1q_f32(tmp_f32 + 12), pout + 12); |
| 343 | } else { |
| 344 | neon_f32_to_fp16x4(vld1q_f32(tmp_f32), pout); |
| 345 | neon_f32_to_fp16x4(vld1q_f32(tmp_f32 + 4), pout + 4); |
| 346 | neon_f32_to_fp16x4(vld1q_f32(tmp_f32 + 8), pout + 8); |
| 347 | neon_f32_to_fp16x4(vld1q_f32(tmp_f32 + 12), pout + 12); |
| 348 | } |
| 349 | } |
| 350 | } |
| 351 | return; |
| 352 | } |
| 353 | #endif // _M_ARM64 || __aarch64__ |
| 354 | |
| 355 | #if defined(__AVX512F__) |
| 356 | if (has_avx512f()) { |
| 357 | long long dim_0 = m; |
| 358 | long long dim_1 = n; |
| 359 | long long input_dim_1 = dim_1 >> 1; |
| 360 | long long absmax_dim_1 = dim_1 / blocksize; |
| 361 | using Tcomp = float; |
| 362 | constexpr auto VEC_LEN = sizeof(__m512i) / sizeof(Tcomp); // 16 |
nothing calls this directly
no test coverage detected