viva_tensor

-module(viva_tensor@backend@dispatch).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/backend/dispatch.gleam").
-export([capabilities/16, available_backends/1, is_available/2, plan_backend/5]).
-export_type([backend_set/1, operation_kind/0, capability/4, rejection/1, plan/2]).

-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 backend_set(MKY) :: {backend_set, MKY, MKY, MKY, MKY, MKY, MKY, MKY}.

-type operation_kind() :: elementwise |
    broadcast |
    reduction |
    softmax |
    {matmul, integer(), integer(), integer()}.

-type capability(MKZ, MLA, MLB, MLC) :: {capability,
        MKZ,
        boolean(),
        MLA,
        list(MLB),
        list(MLC),
        binary()}.

-type rejection(MLD) :: {rejection, MLD, binary()}.

-type plan(MLE, MLF) :: {plan,
        MLE,
        MLF,
        list(MLF),
        list(rejection(MLF)),
        binary()}.

-file("src/viva_tensor/backend/dispatch.gleam", 52).
?DOC(false).
-spec capabilities(
    backend_set(MLG),
    MLI,
    MLI,
    MLI,
    MLJ,
    MLJ,
    MLJ,
    MLJ,
    MLJ,
    MLK,
    MLK,
    MLK,
    MLK,
    MLK,
    boolean(),
    list(viva_tensor@native@tflops:backend())
) -> list(capability(MLG, MLI, MLJ, MLK)).
capabilities(
    Backends,
    Beam_cpu,
    Native_cpu,
    Cuda,
    Float64,
    Float32,
    Float16,
    Int8,
    Sparse_float16,
    Elementwise,
    Broadcast,
    Reduction,
    Softmax,
    Matmul,
    Zig_loaded,
    Detected
) ->
    [{capability,
            erlang:element(2, Backends),
            true,
            Beam_cpu,
            [Float64],
            [Elementwise, Broadcast, Reduction, Softmax, Matmul],
            <<"Always available fallback."/utf8>>},
        {capability,
            erlang:element(3, Backends),
            Zig_loaded,
            Native_cpu,
            [Float64],
            [Elementwise, Reduction, Matmul],
            <<"Portable SIMD NIF for CPU hot paths."/utf8>>},
        {capability,
            erlang:element(4, Backends),
            Zig_loaded,
            Native_cpu,
            [Float64, Float32],
            [Matmul],
            <<"Native BLAS path exposed through the loaded Zig NIF."/utf8>>},
        {capability,
            erlang:element(5, Backends),
            gleam@list:contains(Detected, cuda_f_p32),
            Cuda,
            [Float32],
            [Matmul],
            <<"CUDA FP32/cuBLAS dense matrix multiplication."/utf8>>},
        {capability,
            erlang:element(6, Backends),
            gleam@list:contains(Detected, cuda_f_p16),
            Cuda,
            [Float16],
            [Matmul],
            <<"CUDA FP16 Tensor Core dense matrix multiplication."/utf8>>},
        {capability,
            erlang:element(7, Backends),
            gleam@list:contains(Detected, cuda_i_n_t8),
            Cuda,
            [Int8],
            [Matmul],
            <<"CUDA INT8 IMMA Tensor Core matrix multiplication."/utf8>>},
        {capability,
            erlang:element(8, Backends),
            gleam@list:contains(Detected, cuda_sparse),
            Cuda,
            [Sparse_float16],
            [Matmul],
            <<"CUDA 2:4 sparse Tensor Core matrix multiplication."/utf8>>}].

-file("src/viva_tensor/backend/dispatch.gleam", 130).
?DOC(false).
-spec available_backends(list(capability(MLR, any(), any(), any()))) -> list(MLR).
available_backends(Capabilities) ->
    _pipe = Capabilities,
    _pipe@1 = gleam@list:filter(
        _pipe,
        fun(Capability) -> erlang:element(3, Capability) end
    ),
    gleam@list:map(
        _pipe@1,
        fun(Capability@1) -> erlang:element(2, Capability@1) end
    ).

-file("src/viva_tensor/backend/dispatch.gleam", 138).
?DOC(false).
-spec is_available(MMB, list(capability(MMB, any(), any(), any()))) -> boolean().
is_available(Backend, Capabilities) ->
    _pipe = Capabilities,
    gleam@list:any(
        _pipe,
        fun(Capability) ->
            (erlang:element(2, Capability) =:= Backend) andalso erlang:element(
                3,
                Capability
            )
        end
    ).

-file("src/viva_tensor/backend/dispatch.gleam", 327).
?DOC(false).
-spec operation_specific_rejection(
    operation_kind(),
    MNQ,
    backend_set(MNQ),
    binary()
) -> binary().
operation_specific_rejection(Kind, Backend, Backends, Fallback_reason) ->
    case Kind of
        {matmul, _, _, _} when Backend =:= erlang:element(8, Backends) ->
            <<"Sparse Tensor Core dispatch requires an explicit sparse tensor."/utf8>>;

        {matmul, _, _, _} when Backend =:= erlang:element(7, Backends) ->
            <<"INT8 Tensor Core dispatch requires explicit quantized tensors."/utf8>>;

        {matmul, M, N, K} when Backend =:= erlang:element(6, Backends) ->
            case (((M rem 16) =:= 0) andalso ((N rem 16) =:= 0)) andalso ((K rem 16)
            =:= 0) of
                true ->
                    Fallback_reason;

                false ->
                    <<"FP16 Tensor Core matmul requires dimensions aligned to 16."/utf8>>
            end;

        _ ->
            Fallback_reason
    end.

-file("src/viva_tensor/backend/dispatch.gleam", 311).
?DOC(false).
-spec rejection_reason(
    operation_kind(),
    MNM,
    list(MNM),
    list(MNM),
    backend_set(MNM),
    binary()
) -> binary().
rejection_reason(
    Kind,
    Backend,
    Available,
    Candidates,
    Backends,
    Unsupported_reason
) ->
    case {gleam@list:contains(Candidates, Backend),
        gleam@list:contains(Available, Backend)} of
        {false, _} ->
            operation_specific_rejection(
                Kind,
                Backend,
                Backends,
                Unsupported_reason
            );

        {true, false} ->
            <<"Backend is not available in this VM."/utf8>>;

        {true, true} ->
            <<"A higher-priority backend was selected."/utf8>>
    end.

-file("src/viva_tensor/backend/dispatch.gleam", 347).
?DOC(false).
-spec all_backends(backend_set(MNS)) -> list(MNS).
all_backends(Backends) ->
    [erlang:element(8, Backends),
        erlang:element(6, Backends),
        erlang:element(7, Backends),
        erlang:element(5, Backends),
        erlang:element(4, Backends),
        erlang:element(3, Backends),
        erlang:element(2, Backends)].

-file("src/viva_tensor/backend/dispatch.gleam", 286).
?DOC(false).
-spec backend_rejections(
    operation_kind(),
    MNG,
    list(MNG),
    list(MNG),
    backend_set(MNG),
    binary()
) -> list(rejection(MNG)).
backend_rejections(
    Kind,
    Selected,
    Available,
    Candidates,
    Backends,
    Unsupported_reason
) ->
    _pipe = all_backends(Backends),
    _pipe@1 = gleam@list:filter(_pipe, fun(Backend) -> Backend /= Selected end),
    gleam@list:map(
        _pipe@1,
        fun(Backend@1) ->
            {rejection,
                Backend@1,
                rejection_reason(
                    Kind,
                    Backend@1,
                    Available,
                    Candidates,
                    Backends,
                    Unsupported_reason
                )}
        end
    ).

-file("src/viva_tensor/backend/dispatch.gleam", 276).
?DOC(false).
-spec select_backend(list(MND), list(MND), MND) -> MND.
select_backend(Available, Candidates, Fallback) ->
    _pipe = Candidates,
    _pipe@1 = gleam@list:find(
        _pipe,
        fun(Candidate) -> gleam@list:contains(Available, Candidate) end
    ),
    gleam@result:unwrap(_pipe@1, Fallback).

-file("src/viva_tensor/backend/dispatch.gleam", 249).
?DOC(false).
-spec plan_first_available(
    MMW,
    operation_kind(),
    list(MMX),
    backend_set(MMX),
    list(MMX),
    binary(),
    binary()
) -> plan(MMW, MMX).
plan_first_available(
    Operation,
    Kind,
    Available,
    Backends,
    Candidates,
    Reason,
    Unsupported_reason
) ->
    Selected = select_backend(
        Available,
        Candidates,
        erlang:element(2, Backends)
    ),
    {plan,
        Operation,
        Selected,
        Candidates,
        backend_rejections(
            Kind,
            Selected,
            Available,
            Candidates,
            Backends,
            Unsupported_reason
        ),
        Reason}.

-file("src/viva_tensor/backend/dispatch.gleam", 201).
?DOC(false).
-spec plan_matmul(
    MMQ,
    integer(),
    integer(),
    integer(),
    list(MMR),
    backend_set(MMR),
    boolean()
) -> plan(MMQ, MMR).
plan_matmul(Operation, M, N, K, Available, Backends, Nif_loaded) ->
    Tensor_core_aligned = (((M rem 16) =:= 0) andalso ((N rem 16) =:= 0))
    andalso ((K rem 16) =:= 0),
    Candidates = case Tensor_core_aligned of
        true ->
            [erlang:element(8, Backends),
                erlang:element(6, Backends),
                erlang:element(7, Backends),
                erlang:element(5, Backends),
                erlang:element(4, Backends),
                erlang:element(3, Backends),
                erlang:element(2, Backends)];

        false ->
            [erlang:element(5, Backends),
                erlang:element(4, Backends),
                erlang:element(3, Backends),
                erlang:element(2, Backends)]
    end,
    Reason = case Nif_loaded of
        true ->
            case Tensor_core_aligned of
                true ->
                    <<"Matmul dimensions are Tensor Core aligned; CUDA is preferred."/utf8>>;

                false ->
                    <<"Matmul dimensions are not Tensor Core aligned; dense CUDA/CPU fallback is preferred."/utf8>>
            end;

        false ->
            <<"Native NIF is not loaded; pure Gleam fallback is selected."/utf8>>
    end,
    plan_first_available(
        Operation,
        {matmul, M, N, K},
        Available,
        Backends,
        Candidates,
        Reason,
        <<"Backend is not part of the stable matmul dispatch path for this shape."/utf8>>
    ).

-file("src/viva_tensor/backend/dispatch.gleam", 148).
?DOC(false).
-spec plan_backend(
    MMK,
    operation_kind(),
    list(MML),
    backend_set(MML),
    boolean()
) -> plan(MMK, MML).
plan_backend(Operation, Kind, Available, Backends, Nif_loaded) ->
    case Kind of
        {matmul, M, N, K} ->
            plan_matmul(Operation, M, N, K, Available, Backends, Nif_loaded);

        elementwise ->
            plan_first_available(
                Operation,
                Kind,
                Available,
                Backends,
                [erlang:element(3, Backends),
                    erlang:element(4, Backends),
                    erlang:element(2, Backends)],
                <<"Element-wise ops prefer SIMD, then native CPU, then pure Gleam."/utf8>>,
                <<"Backend does not support stable element-wise dispatch."/utf8>>
            );

        broadcast ->
            plan_first_available(
                Operation,
                Kind,
                Available,
                Backends,
                [erlang:element(3, Backends), erlang:element(2, Backends)],
                <<"Broadcasting preserves views and only needs native compute when materialized."/utf8>>,
                <<"Backend does not support stable broadcast dispatch."/utf8>>
            );

        reduction ->
            plan_first_available(
                Operation,
                Kind,
                Available,
                Backends,
                [erlang:element(3, Backends),
                    erlang:element(4, Backends),
                    erlang:element(2, Backends)],
                <<"Reductions prefer SIMD/native CPU and fall back to pure Gleam."/utf8>>,
                <<"Backend does not support stable reduction dispatch."/utf8>>
            );

        softmax ->
            plan_first_available(
                Operation,
                Kind,
                Available,
                Backends,
                [erlang:element(2, Backends)],
                <<"Softmax currently uses the stable Gleam implementation."/utf8>>,
                <<"Softmax currently only has stable pure Gleam dispatch."/utf8>>
            )
    end.