// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Quantization-aware-training kernel: fake-quantize a weight matrix row by
// row (per-output-channel scales). One thread per row, two strided passes
// over the row (reduce scale, then quantize). Rows are weight-matrix sized
// (C = in_features, a few thousand max) and this runs once per layer per
// forward, so a simple thread-per-row layout is plenty.
//
// Backward is the straight-through estimator (identity), so there is no
// backward kernel — autograd routes dout straight into the master weight's
// grad.
#include <metal_stdlib>
using namespace metal;

// p = (C, mode); mode 0 = int8 (absmax/127), mode 1 = ternary BitNet-style
// (scale = mean|w|, values in {-s, 0, +s}).
kernel void fake_quant_f32(device const float* W   [[buffer(0)]],
                           device float*       OUT [[buffer(1)]],
                           constant uint2&     p   [[buffer(2)]],
                           uint row [[thread_position_in_grid]]) {
    const uint C = p.x;
    device const float* w = W + ulong(row) * C;
    device float* out = OUT + ulong(row) * C;

    if (p.y == 0) { // int8: symmetric absmax
        float amax = 0.0f;
        for (uint j = 0; j < C; ++j) amax = max(amax, fabs(w[j]));
        const float s = max(amax / 127.0f, 1e-12f);
        for (uint j = 0; j < C; ++j) out[j] = rint(w[j] / s) * s;
    } else {        // ternary: BitNet b1.58 absmean
        float asum = 0.0f;
        for (uint j = 0; j < C; ++j) asum += fabs(w[j]);
        const float s = max(asum / float(C), 1e-12f);
        for (uint j = 0; j < C; ++j)
            out[j] = clamp(rint(w[j] / s), -1.0f, 1.0f) * s;
    }
}