Function quantize_cpu_impl

csrc/cpu_ops.cpp:575–665

template <typename T>
void quantize_cpu_impl(float* code, const T* A, float* absmax, unsigned char* out, long long blocksize, long long n) {
    if (blocksize <= 0 || n <= 0)
        return;

    // Get LUT from global cache (built once per codebook, shared by all OMP threads)
    const unsigned char* lut = get_global_lut(code);

    const long long num_blocks = (n + blocksize - 1) / blocksize;

    BNB_OMP_PARALLEL_FOR
    for (long long b = 0; b < num_blocks; ++b) {
        const long long block_start = b * blocksize;
        const long long block_end = std::min(block_start + blocksize, n);
        const long long block_len = block_end - block_start;

        // Compute absmax for this block
        float absmax_block = 0.0f;

#if defined(_M_ARM64) || defined(__aarch64__)
        absmax_block = neon_absmax<T>(A + block_start, block_len);
#else
#pragma omp simd reduction(max : absmax_block)
        for (long long i = block_start; i < block_end; ++i) {
            float val;
            if constexpr (std::is_same<T, float>::value)
                val = A[i];
            else if constexpr (std::is_same<T, bf16_t>::value)
                val = bf16_to_float(A[i].v);
            else
                val = fp16_to_float(A[i].v);
            absmax_block = std::max(absmax_block, std::fabs(val));
        }
#endif

        absmax[b] = absmax_block;

        if (absmax_block == 0.0f) {
            for (long long i = block_start; i < block_end; ++i) {
                out[i] = 0;
            }
            continue;
        }

        const float inv_absmax = 1.0f / absmax_block;

#if defined(_M_ARM64) || defined(__aarch64__)
        {
            long long i = 0;
            float32x4_t vinv = vdupq_n_f32(inv_absmax);
            for (; i + 4 <= block_len; i += 4) {
                float32x4_t v;
                if constexpr (std::is_same<T, float>::value)
                    v = vld1q_f32(reinterpret_cast<const float*>(A + block_start + i));
                else if constexpr (std::is_same<T, bf16_t>::value)
                    v = neon_bf16x4_to_f32(A + block_start + i);
                else
                    v = neon_fp16x4_to_f32(A + block_start + i);
                v = vmulq_f32(v, vinv);