Importance Quantization¶

pub const BlockIQ1S = extern struct {
    d: f16,
    qs: [32]u8,       // sign bits: 0 = +scale, 1 = -scale
    qh: [8]u16,       // importance bitmaps for 8 sub-blocks

    pub fn dequantize(self: BlockIQ1S, output: *[QK_K]f32) void {
        const scale: f32 = @floatCast(self.d);

        for (0..QK_K) |i| {
            // Extract sign bit
            const sign: u1 = @truncate(self.qs[i / 8] >> @truncate(i % 8));
            const base_val: f32 = if (sign == 0) scale else -scale;

            // Check importance bitmap for this sub-block
            const sub_block = i / 32;
            const sub_idx: u4 = @truncate(i % 32);
            const important: bool = (self.qh[sub_block] >> sub_idx) & 1 == 1;

            // Important values get a scaled magnitude boost
            output[i] = if (important) base_val * 1.5 else base_val;
        }
    }
};

pub const BlockIQ2XS = extern struct {
    d: f16,
    qs: [64]u8,        // packed 2-bit quantized values
    scales: [8]u8,      // sub-block scales + importance bits

    pub fn dequantize(self: BlockIQ2XS, output: *[QK_K]f32) void {
        const d_val: f32 = @floatCast(self.d);

        for (0..8) |j| {
            const sc_raw = self.scales[j];
            const sc: f32 = @floatFromInt(sc_raw & 0x0F);
            const importance_shift: u1 = @truncate(sc_raw >> 4);

            for (0..32) |k| {
                const idx = j * 32 + k;
                const byte_idx = idx / 4;
                const shift: u3 = @truncate((idx % 4) * 2);
                const q2: i32 = @as(i32, (self.qs[byte_idx] >> shift) & 0x03) - 1;

                const imp_factor: f32 = if (importance_shift == 1 and
                    q2 != 0) 1.25 else 1.0;
                output[idx] = d_val * sc * @as(f32, @floatFromInt(q2)) * imp_factor;
            }
        }
    }
};

/// Multi-level importance map for a 256-element super-block.
pub const ImportanceMap = struct {
    /// 2-bit importance level per value, packed 4 per byte.
    levels: [64]u8,  // 256 values * 2 bits / 8 = 64 bytes

    pub fn getLevel(self: ImportanceMap, idx: usize) u2 {
        const byte = self.levels[idx / 4];
        const shift: u3 = @truncate((idx % 4) * 2);
        return @truncate(byte >> shift);
    }
};

/// IQ4_NL dequantization lookup table.
/// 16 non-linearly spaced reconstruction values.
pub const IQ4NL_LUT: [16]f32 = .{
    -1.0000, -0.6962, -0.5251, -0.3949,
    -0.2876, -0.1922, -0.1040, -0.0218,
     0.0536,  0.1290,  0.2093,  0.2972,
     0.3979,  0.5220,  0.6944,  1.0000,
};

pub const BlockIQ4NL = extern struct {
    d: f16,
    qs: [128]u8,        // 4-bit LUT indices, packed

    pub fn dequantize(self: BlockIQ4NL, output: *[QK_K]f32) void {
        const scale: f32 = @floatCast(self.d);

        for (0..QK_K / 2) |j| {
            const byte = self.qs[j];
            const idx_lo: u4 = @truncate(byte);
            const idx_hi: u4 = @truncate(byte >> 4);

            output[2 * j] = scale * IQ4NL_LUT[idx_lo];
            output[2 * j + 1] = scale * IQ4NL_LUT[idx_hi];
        }
    }
};

pub const IQuantizer = struct {
    allocator: std.mem.Allocator,
    /// Pre-computed importance scores (optional; if null, uniform importance).
    importance: ?[]const f32,

    pub fn init(
        allocator: std.mem.Allocator,
        importance: ?[]const f32,
    ) IQuantizer {
        return .{
            .allocator = allocator,
            .importance = importance,
        };
    }

    /// Quantize to IQ4_NL (non-linear 4-bit).
    pub fn quantizeIQ4NL(
        self: IQuantizer,
        data: []const f32,
    ) ![]BlockIQ4NL {
        const n_blocks = (data.len + QK_K - 1) / QK_K;
        const blocks = try self.allocator.alloc(BlockIQ4NL, n_blocks);

        for (0..n_blocks) |bi| {
            const start = bi * QK_K;
            const end = @min(start + QK_K, data.len);
            const imp = if (self.importance) |imp| imp[start..end] else null;
            blocks[bi] = quantizeBlockIQ4NL(data[start..end], imp);
        }

        return blocks;
    }

    /// Quantize to IQ2_XS (2-bit with importance).
    pub fn quantizeIQ2XS(
        self: IQuantizer,
        data: []const f32,
    ) ![]BlockIQ2XS {
        const n_blocks = (data.len + QK_K - 1) / QK_K;
        const blocks = try self.allocator.alloc(BlockIQ2XS, n_blocks);

        for (0..n_blocks) |bi| {
            const start = bi * QK_K;
            const end = @min(start + QK_K, data.len);
            const imp = if (self.importance) |imp| imp[start..end] else null;
            blocks[bi] = quantizeBlockIQ2XS(data[start..end], imp);
        }

        return blocks;
    }

    /// Quantize to IQ1_S (extreme 1-bit compression).
    pub fn quantizeIQ1S(
        self: IQuantizer,
        data: []const f32,
    ) ![]BlockIQ1S {
        const n_blocks = (data.len + QK_K - 1) / QK_K;
        const blocks = try self.allocator.alloc(BlockIQ1S, n_blocks);

        for (0..n_blocks) |bi| {
            const start = bi * QK_K;
            const end = @min(start + QK_K, data.len);
            const imp = if (self.importance) |imp| imp[start..end] else null;
            blocks[bi] = quantizeBlockIQ1S(data[start..end], imp);
        }

        return blocks;
    }

    /// Dequantize IQ4_NL blocks to f32.
    pub fn dequantizeIQ4NL(
        blocks: []const BlockIQ4NL,
        output: []f32,
    ) void {
        for (blocks, 0..) |block, bi| {
            block.dequantize(output[bi * QK_K ..][0..QK_K]);
        }
    }

    /// Compute importance scores from gradient information.
    /// Returns Fisher diagonal approximation.
    pub fn computeImportance(
        self: IQuantizer,
        gradients: []const f32,
        weights: []const f32,
    ) ![]f32 {
        std.debug.assert(gradients.len == weights.len);
        const imp = try self.allocator.alloc(f32, weights.len);

        for (0..weights.len) |i| {
            // Fisher diagonal: (dL/dw)^2
            imp[i] = gradients[i] * gradients[i];
        }

        return imp;
    }

    pub fn deinit(self: *IQuantizer, blocks: anytype) void {
        self.allocator.free(blocks);
    }
};