viva_tensor

-module(viva_tensor@cuda).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/cuda.gleam").
-export([new/2, to_list/1, shape/1, matmul/5, fp16_available/0, new16/2, to_list16/1, shape16/1, matmul16/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(
    " CudaTensor - Persistent GPU Memory\n"
    "\n"
    " Tensors that live on the GPU. Ideal for weights and heavy compute.\n"
    "\n"
    " - **FP32 (CudaTensor)**: Standard precision. 40+ TFLOPS on RTX 4090.\n"
    " - **FP16 (CudaTensor16)**: Low precision, high throughput using Tensor Cores. 330+ TFLOPS!\n"
    "\n"
    " Data is uploaded once and stays on device.\n"
    " Operations are launched asynchronously (mostly).\n"
).

-file("src/viva_tensor/cuda.gleam", 22).
?DOC(" Upload data to GPU (FP32)\n").
-spec new(list(float()), list(integer())) -> {ok,
        viva_tensor@core@ffi:cuda_tensor_ref()} |
    {error, binary()}.
new(Data, Shape) ->
    viva_tensor_zig:ct_from_list(Data, Shape).

-file("src/viva_tensor/cuda.gleam", 27).
?DOC(" Download data from GPU (FP32)\n").
-spec to_list(viva_tensor@core@ffi:cuda_tensor_ref()) -> {ok, list(float())} |
    {error, binary()}.
to_list(Tensor) ->
    viva_tensor_zig:ct_to_list(Tensor).

-file("src/viva_tensor/cuda.gleam", 32).
?DOC(" Get shape of tensor\n").
-spec shape(viva_tensor@core@ffi:cuda_tensor_ref()) -> {ok, list(integer())} |
    {error, binary()}.
shape(Tensor) ->
    viva_tensor_zig:ct_shape(Tensor).

-file("src/viva_tensor/cuda.gleam", 38).
?DOC(
    " Matrix Multiplication (FP32)\n"
    " C = A @ B\n"
).
-spec matmul(
    viva_tensor@core@ffi:cuda_tensor_ref(),
    viva_tensor@core@ffi:cuda_tensor_ref(),
    integer(),
    integer(),
    integer()
) -> {ok, viva_tensor@core@ffi:cuda_tensor_ref()} | {error, binary()}.
matmul(A, B, M, N, K) ->
    viva_tensor_zig:ct_matmul(A, B, M, N, K).

-file("src/viva_tensor/cuda.gleam", 57).
?DOC(" Check if FP16 Tensor Cores are available\n").
-spec fp16_available() -> boolean().
fp16_available() ->
    viva_tensor_zig:ct16_available().

-file("src/viva_tensor/cuda.gleam", 62).
?DOC(" Upload data to GPU (converts f64 -> f16)\n").
-spec new16(list(float()), list(integer())) -> {ok,
        viva_tensor@core@ffi:cuda_tensor16_ref()} |
    {error, binary()}.
new16(Data, Shape) ->
    viva_tensor_zig:ct16_from_list(Data, Shape).

-file("src/viva_tensor/cuda.gleam", 70).
?DOC(" Download data from GPU (converts f16 -> f64)\n").
-spec to_list16(viva_tensor@core@ffi:cuda_tensor16_ref()) -> {ok, list(float())} |
    {error, binary()}.
to_list16(Tensor) ->
    viva_tensor_zig:ct16_to_list(Tensor).

-file("src/viva_tensor/cuda.gleam", 75).
?DOC(" Get shape of FP16 tensor\n").
-spec shape16(viva_tensor@core@ffi:cuda_tensor16_ref()) -> {ok, list(integer())} |
    {error, binary()}.
shape16(Tensor) ->
    viva_tensor_zig:ct16_shape(Tensor).

-file("src/viva_tensor/cuda.gleam", 84).
?DOC(
    " Matrix Multiplication (FP16 Tensor Cores)\n"
    " C = A @ B\n"
    "\n"
    " Uses HMMA (Half-precision Matrix Multiply Accumulate) instructions.\n"
    " Expect massive speedups (up to 330 TFLOPS) if dimensions align with 16x16.\n"
).
-spec matmul16(
    viva_tensor@core@ffi:cuda_tensor16_ref(),
    viva_tensor@core@ffi:cuda_tensor16_ref(),
    integer(),
    integer(),
    integer()
) -> {ok, viva_tensor@core@ffi:cuda_tensor16_ref()} | {error, binary()}.
matmul16(A, B, M, N, K) ->
    viva_tensor_zig:ct16_matmul(A, B, M, N, K).