-module(viva_tensor@native@cuda).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/native/cuda.gleam").
-export([new/2, to_list/1, shape/1, matmul/5, fp16_available/0, new16/2, to_list16/1, shape16/1, matmul16/5, gpu_workspace/0, to_rtx4090_fp32/1, to_rtx4090_fp16/1, workspace_zeros/2, workspace_from_tensor/2, workspace_backend/1, linear_layer/3, linear_layer_fp16/2, linear_layer_output_features/1, linear_output/3, accelerated_shape/1, linear_gelu_accelerated_into/4, linear_relu_accelerated_into/4, linear_relu_forward_into/3, linear_gelu_forward_into/3, linear_layer_backend/1, linear_layer_input_features/1, to_accelerated/1, matmul_auto/2, to_cpu_tensor/1, backend/1, matmul_accelerated/2, matmul_accelerated_into/3, matmul_relu_accelerated_into/3, matmul_gelu_accelerated_into/3, sync/0]).
-export_type([acceleration_backend/0, accelerated_tensor/0, gpu_workspace/0, linear_layer/0]).

-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(false).

-type acceleration_backend() :: rtx4090_fp16 |
    rtx4090_fp32 |
    mkl_native |
    cpu_fallback.

-type accelerated_tensor() :: {cuda_fp16,
        viva_tensor@core@ffi:cuda_tensor16_ref(),
        list(integer()),
        acceleration_backend()} |
    {cuda_fp32,
        viva_tensor@core@ffi:cuda_tensor_ref(),
        list(integer()),
        acceleration_backend()} |
    {cpu, viva_tensor@tensor:tensor(), acceleration_backend()}.

-type gpu_workspace() :: {gpu_workspace, acceleration_backend()}.

-type linear_layer() :: {linear_layer,
        accelerated_tensor(),
        accelerated_tensor(),
        integer(),
        integer(),
        acceleration_backend()}.

-file("src/viva_tensor/native/cuda.gleam", 25).
?DOC(false).
-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/native/cuda.gleam", 30).
?DOC(false).
-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/native/cuda.gleam", 35).
?DOC(false).
-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/native/cuda.gleam", 41).
?DOC(false).
-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/native/cuda.gleam", 60).
?DOC(false).
-spec fp16_available() -> boolean().
fp16_available() ->
    viva_tensor_zig:ct16_available().

-file("src/viva_tensor/native/cuda.gleam", 65).
?DOC(false).
-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/native/cuda.gleam", 73).
?DOC(false).
-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/native/cuda.gleam", 78).
?DOC(false).
-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/native/cuda.gleam", 87).
?DOC(false).
-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_tensor_ref()} | {error, binary()}.
matmul16(A, B, M, N, K) ->
    viva_tensor_zig:ct16_matmul(A, B, M, N, K).

-file("src/viva_tensor/native/cuda.gleam", 141).
?DOC(false).
-spec gpu_workspace() -> {ok, gpu_workspace()} |
    {error, viva_tensor@core@error:tensor_error()}.
gpu_workspace() ->
    case viva_tensor@core@ffi:cuda_available() andalso viva_tensor_zig:ct16_available(
        
    ) of
        true ->
            {ok, {gpu_workspace, rtx4090_fp16}};

        false ->
            {error,
                {dimension_error, <<"CUDA FP16 backend is not available"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 665).
?DOC(false).
-spec to_mkl(viva_tensor@tensor:tensor()) -> accelerated_tensor().
to_mkl(T) ->
    {cpu, T, mkl_native}.

-file("src/viva_tensor/native/cuda.gleam", 304).
?DOC(false).
-spec to_rtx4090_fp32(viva_tensor@tensor:tensor()) -> {ok, accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
to_rtx4090_fp32(T) ->
    case viva_tensor@core@ffi:cuda_available() of
        false ->
            {error,
                {dimension_error, <<"CUDA FP32 backend is not available"/utf8>>}};

        true ->
            _pipe = viva_tensor_zig:ct_from_list(
                viva_tensor@tensor:to_list(T),
                viva_tensor@tensor:shape(T)
            ),
            _pipe@1 = gleam@result:map(
                _pipe,
                fun(Ref) ->
                    {cuda_fp32, Ref, viva_tensor@tensor:shape(T), rtx4090_fp32}
                end
            ),
            gleam@result:map_error(
                _pipe@1,
                fun(Reason) -> {dimension_error, Reason} end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 289).
?DOC(false).
-spec to_rtx4090_fp16(viva_tensor@tensor:tensor()) -> {ok, accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
to_rtx4090_fp16(T) ->
    case viva_tensor@core@ffi:cuda_available() andalso viva_tensor_zig:ct16_available(
        
    ) of
        false ->
            {error,
                {dimension_error, <<"CUDA FP16 backend is not available"/utf8>>}};

        true ->
            _pipe = viva_tensor_zig:ct16_from_list(
                viva_tensor@tensor:to_list(T),
                viva_tensor@tensor:shape(T)
            ),
            _pipe@1 = gleam@result:map(
                _pipe,
                fun(Ref) ->
                    {cuda_fp16, Ref, viva_tensor@tensor:shape(T), rtx4090_fp16}
                end
            ),
            gleam@result:map_error(
                _pipe@1,
                fun(Reason) -> {dimension_error, Reason} end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 149).
?DOC(false).
-spec workspace_zeros(gpu_workspace(), list(integer())) -> {ok,
        accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
workspace_zeros(Workspace, Shape) ->
    case Workspace of
        {gpu_workspace, rtx4090_fp16} ->
            to_rtx4090_fp16(viva_tensor@tensor:zeros(Shape));

        {gpu_workspace, rtx4090_fp32} ->
            to_rtx4090_fp32(viva_tensor@tensor:zeros(Shape));

        {gpu_workspace, mkl_native} ->
            _pipe = viva_tensor@tensor:native_zeros(Shape),
            gleam@result:map(_pipe, fun to_mkl/1);

        {gpu_workspace, cpu_fallback} ->
            {ok, {cpu, viva_tensor@tensor:zeros(Shape), cpu_fallback}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 163).
?DOC(false).
-spec workspace_from_tensor(gpu_workspace(), viva_tensor@tensor:tensor()) -> {ok,
        accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
workspace_from_tensor(Workspace, T) ->
    case Workspace of
        {gpu_workspace, rtx4090_fp16} ->
            to_rtx4090_fp16(T);

        {gpu_workspace, rtx4090_fp32} ->
            to_rtx4090_fp32(T);

        {gpu_workspace, mkl_native} ->
            Shape = viva_tensor@tensor:shape(T),
            _pipe = viva_tensor@tensor:native_from_list(
                viva_tensor@tensor:to_list(T),
                Shape
            ),
            gleam@result:map(_pipe, fun to_mkl/1);

        {gpu_workspace, cpu_fallback} ->
            {ok, {cpu, T, cpu_fallback}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 244).
?DOC(false).
-spec workspace_backend(gpu_workspace()) -> acceleration_backend().
workspace_backend(Workspace) ->
    case Workspace of
        {gpu_workspace, Backend} ->
            Backend
    end.

-file("src/viva_tensor/native/cuda.gleam", 189).
?DOC(false).
-spec linear_layer(
    gpu_workspace(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, linear_layer()} | {error, viva_tensor@core@error:tensor_error()}.
linear_layer(Workspace, Weight, Bias) ->
    case {viva_tensor@tensor:shape(Weight), viva_tensor@tensor:shape(Bias)} of
        {[Input_features, Output_features], [Bias_features]} when Output_features =:= Bias_features ->
            gleam@result:'try'(
                workspace_from_tensor(Workspace, Weight),
                fun(Weight_acc) ->
                    gleam@result:'try'(
                        workspace_from_tensor(Workspace, Bias),
                        fun(Bias_acc) ->
                            {ok,
                                {linear_layer,
                                    Weight_acc,
                                    Bias_acc,
                                    Input_features,
                                    Output_features,
                                    workspace_backend(Workspace)}}
                        end
                    )
                end
            );

        {[_, Output_features@1], [Bias_features@1]} ->
            {error, {shape_mismatch, [Output_features@1], [Bias_features@1]}};

        {_, _} ->
            {error,
                {dimension_error,
                    <<"Expected weight matrix and bias vector"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 180).
?DOC(false).
-spec linear_layer_fp16(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, linear_layer()} | {error, viva_tensor@core@error:tensor_error()}.
linear_layer_fp16(Weight, Bias) ->
    gleam@result:'try'(
        gpu_workspace(),
        fun(Workspace) -> linear_layer(Workspace, Weight, Bias) end
    ).

-file("src/viva_tensor/native/cuda.gleam", 265).
?DOC(false).
-spec linear_layer_output_features(linear_layer()) -> integer().
linear_layer_output_features(Layer) ->
    case Layer of
        {linear_layer, _, _, _, Output_features, _} ->
            Output_features
    end.

-file("src/viva_tensor/native/cuda.gleam", 217).
?DOC(false).
-spec linear_output(gpu_workspace(), linear_layer(), integer()) -> {ok,
        accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
linear_output(Workspace, Layer, Batch_size) ->
    workspace_zeros(
        Workspace,
        [Batch_size, linear_layer_output_features(Layer)]
    ).

-file("src/viva_tensor/native/cuda.gleam", 597).
?DOC(false).
-spec fused_linear_accelerated_into_checked(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    integer(),
    integer(),
    integer(),
    binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_linear_accelerated_into_checked(Out, A, B, Bias, M, N, K, Activation) ->
    case {Out, A, B, Bias} of
        {{cuda_fp16, Out_ref, _, _},
            {cuda_fp16, A_ref, _, _},
            {cuda_fp16, B_ref, _, _},
            {cuda_fp16, Bias_ref, _, _}} ->
            _pipe = case Activation of
                <<"relu"/utf8>> ->
                    viva_tensor_zig:ct16_linear_relu(
                        A_ref,
                        B_ref,
                        Bias_ref,
                        Out_ref,
                        M,
                        N,
                        K
                    );

                <<"gelu"/utf8>> ->
                    viva_tensor_zig:ct16_linear_gelu(
                        A_ref,
                        B_ref,
                        Bias_ref,
                        Out_ref,
                        M,
                        N,
                        K
                    );

                _ ->
                    {error, <<"unsupported_activation"/utf8>>}
            end,
            gleam@result:map_error(
                _pipe,
                fun(Reason) -> {dimension_error, Reason} end
            );

        {{cpu, Out_tensor, _},
            {cpu, A_tensor, _},
            {cpu, B_tensor, _},
            {cpu, Bias_tensor, _}} ->
            case Activation of
                <<"relu"/utf8>> ->
                    viva_tensor@tensor:linear_relu_into(
                        Out_tensor,
                        A_tensor,
                        B_tensor,
                        Bias_tensor
                    );

                _ ->
                    {error,
                        {dimension_error,
                            <<"CPU fused GELU is not implemented"/utf8>>}}
            end;

        {_, _, _, _} ->
            {error,
                {dimension_error,
                    <<"Fused linear activation requires matching backends"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 437).
?DOC(false).
-spec accelerated_shape(accelerated_tensor()) -> list(integer()).
accelerated_shape(T) ->
    case T of
        {cuda_fp16, _, Shape, _} ->
            Shape;

        {cuda_fp32, _, Shape@1, _} ->
            Shape@1;

        {cpu, Tensor, _} ->
            erlang:element(3, Tensor)
    end.

-file("src/viva_tensor/native/cuda.gleam", 563).
?DOC(false).
-spec fused_linear_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_linear_accelerated_into(Out, A, B, Bias, Activation) ->
    case {accelerated_shape(Out),
        accelerated_shape(A),
        accelerated_shape(B),
        accelerated_shape(Bias)} of
        {[M_out, N_out], [M, K], [K2, N], [N_bias]} when (((K =:= K2) andalso (M_out =:= M)) andalso (N_out =:= N)) andalso (N_bias =:= N) ->
            fused_linear_accelerated_into_checked(
                Out,
                A,
                B,
                Bias,
                M,
                N,
                K,
                Activation
            );

        {[M_out@1, N_out@1], [M@1, _], [_, N@1], [_]} ->
            {error, {shape_mismatch, [M@1, N@1], [M_out@1, N_out@1]}};

        {_, _, _, _} ->
            {error,
                {dimension_error,
                    <<"Expected matrices and a bias vector"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 401).
?DOC(false).
-spec linear_gelu_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_gelu_accelerated_into(Out, A, B, Bias) ->
    fused_linear_accelerated_into(Out, A, B, Bias, <<"gelu"/utf8>>).

-file("src/viva_tensor/native/cuda.gleam", 391).
?DOC(false).
-spec linear_relu_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_relu_accelerated_into(Out, A, B, Bias) ->
    fused_linear_accelerated_into(Out, A, B, Bias, <<"relu"/utf8>>).

-file("src/viva_tensor/native/cuda.gleam", 635).
?DOC(false).
-spec linear_forward_into(
    accelerated_tensor(),
    accelerated_tensor(),
    linear_layer(),
    binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_forward_into(Out, Input, Layer, Activation) ->
    case Layer of
        {linear_layer, Weight, Bias, Input_features, Output_features, _} ->
            case {accelerated_shape(Input), accelerated_shape(Out)} of
                {[Batch_size, Got_input], [Out_batch, Got_output]} when ((Got_input =:= Input_features) andalso (Out_batch =:= Batch_size)) andalso (Got_output =:= Output_features) ->
                    case Activation of
                        <<"relu"/utf8>> ->
                            linear_relu_accelerated_into(
                                Out,
                                Input,
                                Weight,
                                Bias
                            );

                        <<"gelu"/utf8>> ->
                            linear_gelu_accelerated_into(
                                Out,
                                Input,
                                Weight,
                                Bias
                            );

                        _ ->
                            {error,
                                {dimension_error,
                                    <<"Unsupported activation"/utf8>>}}
                    end;

                {[_, Got_input@1], [_, _]} ->
                    {error, {shape_mismatch, [Input_features], [Got_input@1]}};

                {_, _} ->
                    {error,
                        {dimension_error,
                            <<"Expected input and output matrices"/utf8>>}}
            end
    end.

-file("src/viva_tensor/native/cuda.gleam", 226).
?DOC(false).
-spec linear_relu_forward_into(
    accelerated_tensor(),
    accelerated_tensor(),
    linear_layer()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_relu_forward_into(Out, Input, Layer) ->
    linear_forward_into(Out, Input, Layer, <<"relu"/utf8>>).

-file("src/viva_tensor/native/cuda.gleam", 235).
?DOC(false).
-spec linear_gelu_forward_into(
    accelerated_tensor(),
    accelerated_tensor(),
    linear_layer()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_gelu_forward_into(Out, Input, Layer) ->
    linear_forward_into(Out, Input, Layer, <<"gelu"/utf8>>).

-file("src/viva_tensor/native/cuda.gleam", 251).
?DOC(false).
-spec linear_layer_backend(linear_layer()) -> acceleration_backend().
linear_layer_backend(Layer) ->
    case Layer of
        {linear_layer, _, _, _, _, Backend} ->
            Backend
    end.

-file("src/viva_tensor/native/cuda.gleam", 258).
?DOC(false).
-spec linear_layer_input_features(linear_layer()) -> integer().
linear_layer_input_features(Layer) ->
    case Layer of
        {linear_layer, _, _, Input_features, _, _} ->
            Input_features
    end.

-file("src/viva_tensor/native/cuda.gleam", 740).
?DOC(false).
-spec to_native(viva_tensor@tensor:tensor(), list(integer())) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, nil}.
to_native(T, Shape) ->
    case viva_tensor@tensor:native_ref(T) of
        {ok, _} ->
            {ok, T};

        {error, _} ->
            _pipe = viva_tensor@tensor:native_from_list(
                viva_tensor@tensor:to_list(T),
                Shape
            ),
            gleam@result:map_error(_pipe, fun(_) -> nil end)
    end.

-file("src/viva_tensor/native/cuda.gleam", 749).
?DOC(false).
-spec to_mkl_or_cpu(viva_tensor@tensor:tensor(), list(integer())) -> accelerated_tensor().
to_mkl_or_cpu(T, Shape) ->
    case to_native(T, Shape) of
        {ok, Native} ->
            {cpu, Native, mkl_native};

        {error, _} ->
            {cpu, T, cpu_fallback}
    end.

-file("src/viva_tensor/native/cuda.gleam", 273).
?DOC(false).
-spec to_accelerated(viva_tensor@tensor:tensor()) -> {ok, accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
to_accelerated(T) ->
    Shape = viva_tensor@tensor:shape(T),
    case to_rtx4090_fp16(T) of
        {ok, Gpu} ->
            {ok, Gpu};

        {error, _} ->
            case to_rtx4090_fp32(T) of
                {ok, Gpu@1} ->
                    {ok, Gpu@1};

                {error, _} ->
                    {ok, to_mkl_or_cpu(T, Shape)}
            end
    end.

-file("src/viva_tensor/native/cuda.gleam", 723).
?DOC(false).
-spec matmul_mkl_then_cpu(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    integer(),
    integer(),
    integer()
) -> {ok, accelerated_tensor()} | {error, viva_tensor@core@error:tensor_error()}.
matmul_mkl_then_cpu(A, B, M, N, K) ->
    case {to_native(A, [M, K]), to_native(B, [K, N])} of
        {{ok, A_native}, {ok, B_native}} ->
            _pipe = viva_tensor@tensor:matmul(A_native, B_native),
            gleam@result:map(_pipe, fun(Out) -> {cpu, Out, mkl_native} end);

        {_, _} ->
            _pipe@1 = viva_tensor@tensor:matmul(A, B),
            gleam@result:map(
                _pipe@1,
                fun(Out@1) -> {cpu, Out@1, cpu_fallback} end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 696).
?DOC(false).
-spec try_rtx4090_fp32(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    integer(),
    integer(),
    integer()
) -> {ok, accelerated_tensor()} | {error, nil}.
try_rtx4090_fp32(A, B, M, N, K) ->
    case viva_tensor@core@ffi:cuda_available() of
        false ->
            {error, nil};

        true ->
            gleam@result:'try'(
                begin
                    _pipe = viva_tensor_zig:ct_from_list(
                        viva_tensor@tensor:to_list(A),
                        [M, K]
                    ),
                    gleam@result:map_error(_pipe, fun(_) -> nil end)
                end,
                fun(A_gpu) ->
                    gleam@result:'try'(
                        begin
                            _pipe@1 = viva_tensor_zig:ct_from_list(
                                viva_tensor@tensor:to_list(B),
                                [K, N]
                            ),
                            gleam@result:map_error(_pipe@1, fun(_) -> nil end)
                        end,
                        fun(B_gpu) ->
                            _pipe@2 = viva_tensor_zig:ct_matmul(
                                A_gpu,
                                B_gpu,
                                M,
                                N,
                                K
                            ),
                            _pipe@3 = gleam@result:map(
                                _pipe@2,
                                fun(Out) ->
                                    {cuda_fp32, Out, [M, N], rtx4090_fp32}
                                end
                            ),
                            gleam@result:map_error(_pipe@3, fun(_) -> nil end)
                        end
                    )
                end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 669).
?DOC(false).
-spec try_rtx4090_fp16(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    integer(),
    integer(),
    integer()
) -> {ok, accelerated_tensor()} | {error, nil}.
try_rtx4090_fp16(A, B, M, N, K) ->
    case viva_tensor@core@ffi:cuda_available() andalso viva_tensor_zig:ct16_available(
        
    ) of
        false ->
            {error, nil};

        true ->
            gleam@result:'try'(
                begin
                    _pipe = viva_tensor_zig:ct16_from_list(
                        viva_tensor@tensor:to_list(A),
                        [M, K]
                    ),
                    gleam@result:map_error(_pipe, fun(_) -> nil end)
                end,
                fun(A_gpu) ->
                    gleam@result:'try'(
                        begin
                            _pipe@1 = viva_tensor_zig:ct16_from_list(
                                viva_tensor@tensor:to_list(B),
                                [K, N]
                            ),
                            gleam@result:map_error(_pipe@1, fun(_) -> nil end)
                        end,
                        fun(B_gpu) ->
                            _pipe@2 = viva_tensor_zig:ct16_matmul(
                                A_gpu,
                                B_gpu,
                                M,
                                N,
                                K
                            ),
                            _pipe@3 = gleam@result:map(
                                _pipe@2,
                                fun(Out) ->
                                    {cuda_fp32, Out, [M, N], rtx4090_fp16}
                                end
                            ),
                            gleam@result:map_error(_pipe@3, fun(_) -> nil end)
                        end
                    )
                end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 319).
?DOC(false).
-spec matmul_auto(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> {ok,
        accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
matmul_auto(A, B) ->
    case {viva_tensor@tensor:shape(A), viva_tensor@tensor:shape(B)} of
        {[M, K], [K2, N]} when K =:= K2 ->
            case try_rtx4090_fp16(A, B, M, N, K) of
                {ok, Result} ->
                    {ok, Result};

                {error, _} ->
                    case try_rtx4090_fp32(A, B, M, N, K) of
                        {ok, Result@1} ->
                            {ok, Result@1};

                        {error, _} ->
                            matmul_mkl_then_cpu(A, B, M, N, K)
                    end
            end;

        {[_, K@1], [K2@1, _]} ->
            {error, {shape_mismatch, [K@1, -1], [K2@1, -1]}};

        {_, _} ->
            {error, {dimension_error, <<"Expected two matrices"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 411).
?DOC(false).
-spec to_cpu_tensor(accelerated_tensor()) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
to_cpu_tensor(T) ->
    case T of
        {cpu, Tensor, _} ->
            {ok, Tensor};

        {cuda_fp16, Ref, Shape, _} ->
            _pipe = viva_tensor_zig:ct16_to_list(Ref),
            _pipe@1 = gleam@result:map(
                _pipe,
                fun(Data) -> {tensor, Data, Shape} end
            ),
            gleam@result:map_error(
                _pipe@1,
                fun(Reason) -> {dimension_error, Reason} end
            );

        {cuda_fp32, Ref@1, Shape@1, _} ->
            _pipe@2 = viva_tensor_zig:ct_to_list(Ref@1),
            _pipe@3 = gleam@result:map(
                _pipe@2,
                fun(Data@1) -> {tensor, Data@1, Shape@1} end
            ),
            gleam@result:map_error(
                _pipe@3,
                fun(Reason@1) -> {dimension_error, Reason@1} end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 428).
?DOC(false).
-spec backend(accelerated_tensor()) -> acceleration_backend().
backend(T) ->
    case T of
        {cuda_fp16, _, _, Backend} ->
            Backend;

        {cuda_fp32, _, _, Backend@1} ->
            Backend@1;

        {cpu, _, Backend@2} ->
            Backend@2
    end.

-file("src/viva_tensor/native/cuda.gleam", 451).
?DOC(false).
-spec matmul_accelerated_checked(
    accelerated_tensor(),
    accelerated_tensor(),
    integer(),
    integer(),
    integer()
) -> {ok, accelerated_tensor()} | {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated_checked(A, B, M, N, K) ->
    case {A, B} of
        {{cuda_fp16, A_ref, _, _}, {cuda_fp16, B_ref, _, _}} ->
            _pipe = viva_tensor_zig:ct16_matmul(A_ref, B_ref, M, N, K),
            _pipe@1 = gleam@result:map(
                _pipe,
                fun(Out) -> {cuda_fp32, Out, [M, N], rtx4090_fp16} end
            ),
            gleam@result:map_error(
                _pipe@1,
                fun(Reason) -> {dimension_error, Reason} end
            );

        {{cuda_fp32, A_ref@1, _, _}, {cuda_fp32, B_ref@1, _, _}} ->
            _pipe@2 = viva_tensor_zig:ct_matmul(A_ref@1, B_ref@1, M, N, K),
            _pipe@3 = gleam@result:map(
                _pipe@2,
                fun(Out@1) -> {cuda_fp32, Out@1, [M, N], rtx4090_fp32} end
            ),
            gleam@result:map_error(
                _pipe@3,
                fun(Reason@1) -> {dimension_error, Reason@1} end
            );

        {{cpu, A_tensor, _}, {cpu, B_tensor, _}} ->
            _pipe@4 = viva_tensor@tensor:matmul(A_tensor, B_tensor),
            gleam@result:map(
                _pipe@4,
                fun(Out@2) -> {cpu, Out@2, backend(A)} end
            );

        {_, _} ->
            gleam@result:'try'(
                to_cpu_tensor(A),
                fun(A_cpu) ->
                    gleam@result:'try'(
                        to_cpu_tensor(B),
                        fun(B_cpu) -> matmul_auto(A_cpu, B_cpu) end
                    )
                end
            )
    end.

-file("src/viva_tensor/native/cuda.gleam", 344).
?DOC(false).
-spec matmul_accelerated(accelerated_tensor(), accelerated_tensor()) -> {ok,
        accelerated_tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated(A, B) ->
    case {accelerated_shape(A), accelerated_shape(B)} of
        {[M, K], [K2, N]} when K =:= K2 ->
            matmul_accelerated_checked(A, B, M, N, K);

        {[_, K@1], [K2@1, _]} ->
            {error, {shape_mismatch, [K@1, -1], [K2@1, -1]}};

        {_, _} ->
            {error, {dimension_error, <<"Expected two matrices"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 515).
?DOC(false).
-spec backend_label(accelerated_tensor()) -> binary().
backend_label(T) ->
    case T of
        {cuda_fp16, _, _, _} ->
            <<"cuda/fp16"/utf8>>;

        {cuda_fp32, _, _, _} ->
            <<"cuda/fp32"/utf8>>;

        {cpu, _, _} ->
            <<"cpu"/utf8>>
    end.

-file("src/viva_tensor/native/cuda.gleam", 485).
?DOC(false).
-spec matmul_accelerated_into_checked(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    integer(),
    integer(),
    integer()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated_into_checked(Out, A, B, M, N, K) ->
    case {Out, A, B} of
        {{cuda_fp16, Out_ref, _, _},
            {cuda_fp16, A_ref, _, _},
            {cuda_fp16, B_ref, _, _}} ->
            _pipe = viva_tensor_zig:ct16_matmul_inplace(
                A_ref,
                B_ref,
                Out_ref,
                M,
                N,
                K
            ),
            gleam@result:map_error(
                _pipe,
                fun(Reason) -> {dimension_error, Reason} end
            );

        {{cuda_fp32, Out_ref@1, _, _},
            {cuda_fp32, A_ref@1, _, _},
            {cuda_fp32, B_ref@1, _, _}} ->
            _pipe@1 = viva_tensor_zig:ct_matmul_inplace(
                A_ref@1,
                B_ref@1,
                Out_ref@1,
                M,
                N,
                K
            ),
            gleam@result:map_error(
                _pipe@1,
                fun(Reason@1) -> {dimension_error, Reason@1} end
            );

        {{cpu, Out_tensor, _}, {cpu, A_tensor, _}, {cpu, B_tensor, _}} ->
            viva_tensor@tensor:matmul_into(Out_tensor, A_tensor, B_tensor);

        {_, _, _} ->
            {error,
                {backend_mismatch,
                    <<"matmul_accelerated_into"/utf8>>,
                    backend_label(Out),
                    backend_label(A),
                    backend_label(B)}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 358).
?DOC(false).
-spec matmul_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated_into(Out, A, B) ->
    case {accelerated_shape(Out), accelerated_shape(A), accelerated_shape(B)} of
        {[M_out, N_out], [M, K], [K2, N]} when ((K =:= K2) andalso (M_out =:= M)) andalso (N_out =:= N) ->
            matmul_accelerated_into_checked(Out, A, B, M, N, K);

        {[M_out@1, N_out@1], [M@1, _], [_, N@1]} ->
            {error, {shape_mismatch, [M@1, N@1], [M_out@1, N_out@1]}};

        {_, _, _} ->
            {error, {dimension_error, <<"Expected matrices"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 538).
?DOC(false).
-spec fused_activation_accelerated_into_checked(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    integer(),
    integer(),
    integer(),
    binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_activation_accelerated_into_checked(Out, A, B, M, N, K, Activation) ->
    case {Out, A, B} of
        {{cuda_fp16, Out_ref, _, _},
            {cuda_fp16, A_ref, _, _},
            {cuda_fp16, B_ref, _, _}} ->
            _pipe = case Activation of
                <<"relu"/utf8>> ->
                    viva_tensor_zig:ct16_matmul_fused_relu(
                        A_ref,
                        B_ref,
                        Out_ref,
                        M,
                        N,
                        K
                    );

                <<"gelu"/utf8>> ->
                    viva_tensor_zig:ct16_matmul_fused_gelu(
                        A_ref,
                        B_ref,
                        Out_ref,
                        M,
                        N,
                        K
                    );

                _ ->
                    {error, <<"unsupported_activation"/utf8>>}
            end,
            gleam@result:map_error(
                _pipe,
                fun(Reason) -> {dimension_error, Reason} end
            );

        {_, _, _} ->
            {error,
                {dimension_error,
                    <<"Fused CUDA activation requires FP16 accelerated tensors"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 523).
?DOC(false).
-spec fused_activation_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor(),
    binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_activation_accelerated_into(Out, A, B, Activation) ->
    case {accelerated_shape(Out), accelerated_shape(A), accelerated_shape(B)} of
        {[M_out, N_out], [M, K], [K2, N]} when ((K =:= K2) andalso (M_out =:= M)) andalso (N_out =:= N) ->
            fused_activation_accelerated_into_checked(
                Out,
                A,
                B,
                M,
                N,
                K,
                Activation
            );

        {[M_out@1, N_out@1], [M@1, _], [_, N@1]} ->
            {error, {shape_mismatch, [M@1, N@1], [M_out@1, N_out@1]}};

        {_, _, _} ->
            {error, {dimension_error, <<"Expected matrices"/utf8>>}}
    end.

-file("src/viva_tensor/native/cuda.gleam", 373).
?DOC(false).
-spec matmul_relu_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_relu_accelerated_into(Out, A, B) ->
    fused_activation_accelerated_into(Out, A, B, <<"relu"/utf8>>).

-file("src/viva_tensor/native/cuda.gleam", 382).
?DOC(false).
-spec matmul_gelu_accelerated_into(
    accelerated_tensor(),
    accelerated_tensor(),
    accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_gelu_accelerated_into(Out, A, B) ->
    fused_activation_accelerated_into(Out, A, B, <<"gelu"/utf8>>).

-file("src/viva_tensor/native/cuda.gleam", 446).
?DOC(false).
-spec sync() -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
sync() ->
    _pipe = viva_tensor_zig:cuda_sync(),
    gleam@result:map_error(_pipe, fun(Reason) -> {dimension_error, Reason} end).