-module(viva_tensor@quant).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/quant.gleam").
-export([matmul_nf4/7]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(
    " Quantization - NF4 / INT8 Compression\n"
    "\n"
    " Run larger models by compressing weights.\n"
    "\n"
    " - **NF4 (Normal Float 4)**: 4-bit quantization optimized for normal distributions (neural weights).\n"
    "   Compresses 16-bit floats to 4 bits (4x memory reduction) with negligible accuracy loss.\n"
    "   Used in QLoRA and modern LLM inference.\n"
).

-file("src/viva_tensor/quant.gleam", 21).
?DOC(
    " Perform Matrix Multiplication with NF4 Quantized Weights\n"
    "\n"
    " C = A @ Quantized(B)\n"
    "\n"
    " Dequantizes B on-the-fly during computation. High compute density, low memory bandwidth.\n"
    "\n"
    " - `a`: Input activations (NativeTensorRef)\n"
    " - `b_indices`: Packed 4-bit indices (2 per byte) for weights\n"
    " - `b_scales`: Scales for blocks of weights\n"
    " - `block_size`: Size of quantization block (usually 64 or 128)\n"
).
-spec matmul_nf4(
    viva_tensor@core@ffi:native_tensor_ref(),
    list(integer()),
    list(float()),
    integer(),
    integer(),
    integer(),
    integer()
) -> viva_tensor@core@ffi:native_tensor_ref().
matmul_nf4(A, B_indices, B_scales, M, N, K, Block_size) ->
    Res@1 = case viva_tensor@core@ffi:nt_matmul_nf4(
        A,
        B_indices,
        B_scales,
        M,
        N,
        K,
        Block_size
    ) of
        {ok, Res} -> Res;
        _assert_fail ->
            erlang:error(#{gleam_error => let_assert,
                        message => <<"Pattern match failed, no pattern matched the value."/utf8>>,
                        file => <<?FILEPATH/utf8>>,
                        module => <<"viva_tensor/quant"/utf8>>,
                        function => <<"matmul_nf4"/utf8>>,
                        line => 30,
                        value => _assert_fail,
                        start => 966,
                        'end' => 1053,
                        pattern_start => 977,
                        pattern_end => 984})
    end,
    Res@1.