viva_tensor

-module(viva_tensor@sparse).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/sparse.gleam").
-export([available/0, from_cuda16/1, shape/1, compression_ratio/1, matmul/5]).

-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(
    " SparseTensor - 2:4 Sparsity\n"
    "\n"
    " Use cuSPARSELt to prune and compress weight matrices.\n"
    "\n"
    " - **2:4 Structure**: For every block of 4 elements, 2 must be zero.\n"
    " - **Compression**: Reduces memory usage by ~50% (1.78x practical).\n"
    " - **Speedup**: Up to 2x theoretical (660 TFLOPS), 61% measured speedup vs dense.\n"
    "\n"
    " Ideal for Large Language Model (LLM) weights.\n"
).

-file("src/viva_tensor/sparse.gleam", 18).
?DOC(" Check if cuSPARSELt is available\n").
-spec available() -> boolean().
available() ->
    viva_tensor_zig:sparse_available().

-file("src/viva_tensor/sparse.gleam", 26).
?DOC(
    " Create SparseTensor from CudaTensor16 (Prune + Compress)\n"
    "\n"
    " This operation is destructive: it prunes the smallest 2 values in every 4-element block.\n"
    " The resulting sparse tensor is stored in a compressed format on the GPU.\n"
).
-spec from_cuda16(viva_tensor@core@ffi:cuda_tensor16_ref()) -> {ok,
        viva_tensor@core@ffi:sparse_tensor_ref()} |
    {error, binary()}.
from_cuda16(Tensor) ->
    viva_tensor_zig:sparse_from_ct16(Tensor).

-file("src/viva_tensor/sparse.gleam", 31).
?DOC(" Get shape of the original dense tensor [Rows, Cols]\n").
-spec shape(viva_tensor@core@ffi:sparse_tensor_ref()) -> {ok, list(integer())} |
    {error, binary()}.
shape(Tensor) ->
    viva_tensor_zig:sparse_shape(Tensor).

-file("src/viva_tensor/sparse.gleam", 36).
?DOC(" Get actual compression ratio (DenseBytes / SparseBytes)\n").
-spec compression_ratio(viva_tensor@core@ffi:sparse_tensor_ref()) -> {ok,
        float()} |
    {error, binary()}.
compression_ratio(Tensor) ->
    viva_tensor_zig:sparse_compression_ratio(Tensor).

-file("src/viva_tensor/sparse.gleam", 48).
?DOC(
    " Sparse Matrix Multiplication (SpMM)\n"
    "\n"
    " C = Sparse(A) @ Dense(B)\n"
    "\n"
    " - `a_sparse`: Compressed weight matrix (2:4 sparse)\n"
    " - `b_dense`: Dense activation matrix (FP16)\n"
    "\n"
    " Returns dense FP16 result.\n"
).
-spec matmul(
    viva_tensor@core@ffi:sparse_tensor_ref(),
    viva_tensor@core@ffi:cuda_tensor16_ref(),
    integer(),
    integer(),
    integer()
) -> {ok, viva_tensor@core@ffi:cuda_tensor16_ref()} | {error, binary()}.
matmul(A_sparse, B_dense, M, N, K) ->
    viva_tensor_zig:sparse_matmul(A_sparse, B_dense, M, N, K).