viva_tensor

-module(viva_tensor@runtime).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/runtime.gleam").
-export([cache_key_for/2, plan_runtime/4, cache_key/1]).
-export_type([runtime_op/0, runtime_backend_set/1, runtime_rejection/1, runtime_plan/1]).

-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 runtime_op() :: runtime_elementwise |
    runtime_broadcast |
    runtime_reduction |
    runtime_softmax |
    {runtime_matmul, integer(), integer(), integer()} |
    {runtime_linear, integer(), integer(), integer()}.

-type runtime_backend_set(AISX) :: {runtime_backend_set,
        AISX,
        AISX,
        AISX,
        AISX,
        AISX,
        AISX,
        AISX}.

-type runtime_rejection(AISY) :: {runtime_rejection, AISY, binary()}.

-type runtime_plan(AISZ) :: {runtime_plan,
        viva_tensor@spec:tensor_spec(),
        runtime_op(),
        AISZ,
        list(AISZ),
        list(runtime_rejection(AISZ)),
        binary(),
        binary()}.

-file("src/viva_tensor/runtime.gleam", 213).
?DOC(false).
-spec op_key(runtime_op()) -> binary().
op_key(Operation) ->
    case Operation of
        runtime_elementwise ->
            <<"elementwise"/utf8>>;

        runtime_broadcast ->
            <<"broadcast"/utf8>>;

        runtime_reduction ->
            <<"reduction"/utf8>>;

        runtime_softmax ->
            <<"softmax"/utf8>>;

        {runtime_matmul, M, N, K} ->
            <<<<<<<<<<"matmul:"/utf8, (erlang:integer_to_binary(M))/binary>>/binary,
                            "x"/utf8>>/binary,
                        (erlang:integer_to_binary(N))/binary>>/binary,
                    "x"/utf8>>/binary,
                (erlang:integer_to_binary(K))/binary>>;

        {runtime_linear, Batch, In_features, Out_features} ->
            <<<<<<<<<<"linear:"/utf8, (erlang:integer_to_binary(Batch))/binary>>/binary,
                            "x"/utf8>>/binary,
                        (erlang:integer_to_binary(In_features))/binary>>/binary,
                    "x"/utf8>>/binary,
                (erlang:integer_to_binary(Out_features))/binary>>
    end.

-file("src/viva_tensor/runtime.gleam", 73).
?DOC(false).
-spec cache_key_for(viva_tensor@spec:tensor_spec(), runtime_op()) -> binary().
cache_key_for(Spec, Operation) ->
    <<<<(viva_tensor@spec:spec_key(Spec))/binary, "|"/utf8>>/binary,
        (op_key(Operation))/binary>>.

-file("src/viva_tensor/runtime.gleam", 199).
?DOC(false).
-spec reason_for(viva_tensor@spec:tensor_spec(), runtime_op()) -> binary().
reason_for(Spec, Operation) ->
    case {Operation, erlang:element(3, Spec), erlang:element(6, Spec)} of
        {{runtime_linear, _, _, _}, int4, packed_sparse24_layout} ->
            <<"INT4 2:4 packed linear prefers the sparse Tensor Core path."/utf8>>;

        {{runtime_linear, _, _, _}, int8, packed_sparse24_layout} ->
            <<"INT8 2:4 packed linear prefers sparse INT8 Tensor Cores."/utf8>>;

        {{runtime_linear, _, _, _}, float8_e4_m3, packed_fp8_layout} ->
            <<"FP8 packed linear prefers the FP8/CUTLASS-capable CUDA path."/utf8>>;

        {{runtime_matmul, _, _, _}, _, _} ->
            <<"Matmul planning is shape-driven and prefers Tensor Core aligned CUDA."/utf8>>;

        {_, _, _} ->
            <<"Runtime planning selected the first available stable backend."/utf8>>
    end.

-file("src/viva_tensor/runtime.gleam", 187).
?DOC(false).
-spec all_backends(runtime_backend_set(AITV)) -> list(AITV).
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/runtime.gleam", 164).
?DOC(false).
-spec rejected_backends(AITP, list(AITP), list(AITP), runtime_backend_set(AITP)) -> list(runtime_rejection(AITP)).
rejected_backends(Selected, Available, Candidates, Backends) ->
    _pipe = all_backends(Backends),
    _pipe@1 = gleam@list:filter(_pipe, fun(Backend) -> Backend /= Selected end),
    gleam@list:map(
        _pipe@1,
        fun(Backend@1) ->
            {runtime_rejection,
                Backend@1,
                case {gleam@list:contains(Candidates, Backend@1),
                    gleam@list:contains(Available, Backend@1)} of
                    {false, _} ->
                        <<"Backend is not part of the runtime plan for this spec."/utf8>>;

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

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

-file("src/viva_tensor/runtime.gleam", 139).
?DOC(false).
-spec matmul_candidates(
    integer(),
    integer(),
    integer(),
    runtime_backend_set(AITM)
) -> list(AITM).
matmul_candidates(M, N, K, Backends) ->
    case (((M rem 16) =:= 0) andalso ((N rem 16) =:= 0)) andalso ((K rem 16) =:= 0) 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.

-file("src/viva_tensor/runtime.gleam", 96).
?DOC(false).
-spec linear_candidates(
    viva_tensor@spec:tensor_spec(),
    integer(),
    integer(),
    runtime_backend_set(AITJ)
) -> list(AITJ).
linear_candidates(Spec, In_features, Out_features, Backends) ->
    Tensor_core_aligned = ((In_features rem 16) =:= 0) andalso ((Out_features
    rem 16)
    =:= 0),
    case {erlang:element(3, Spec), erlang:element(6, Spec), Tensor_core_aligned} of
        {int4, packed_sparse24_layout, true} ->
            [erlang:element(8, Backends),
                erlang:element(7, Backends),
                erlang:element(6, Backends),
                erlang:element(4, Backends),
                erlang:element(2, Backends)];

        {int8, packed_sparse24_layout, true} ->
            [erlang:element(7, Backends),
                erlang:element(8, Backends),
                erlang:element(6, Backends),
                erlang:element(4, Backends),
                erlang:element(2, Backends)];

        {float8_e4_m3, packed_fp8_layout, true} ->
            [erlang:element(6, Backends),
                erlang:element(5, Backends),
                erlang:element(4, Backends),
                erlang:element(2, Backends)];

        {float16, _, true} ->
            [erlang:element(6, Backends),
                erlang:element(5, Backends),
                erlang:element(4, Backends),
                erlang:element(2, Backends)];

        {_, _, _} ->
            [erlang:element(5, Backends),
                erlang:element(4, Backends),
                erlang:element(3, Backends),
                erlang:element(2, Backends)]
    end.

-file("src/viva_tensor/runtime.gleam", 80).
?DOC(false).
-spec candidates_for(
    viva_tensor@spec:tensor_spec(),
    runtime_op(),
    runtime_backend_set(AITG)
) -> list(AITG).
candidates_for(Spec, Operation, Backends) ->
    case Operation of
        {runtime_linear, _, In_features, Out_features} ->
            linear_candidates(Spec, In_features, Out_features, Backends);

        {runtime_matmul, M, N, K} ->
            matmul_candidates(M, N, K, Backends);

        runtime_elementwise ->
            [erlang:element(3, Backends),
                erlang:element(4, Backends),
                erlang:element(2, Backends)];

        runtime_broadcast ->
            [erlang:element(3, Backends), erlang:element(2, Backends)];

        runtime_reduction ->
            [erlang:element(3, Backends),
                erlang:element(4, Backends),
                erlang:element(2, Backends)];

        runtime_softmax ->
            [erlang:element(2, Backends)]
    end.

-file("src/viva_tensor/runtime.gleam", 46).
?DOC(false).
-spec plan_runtime(
    viva_tensor@spec:tensor_spec(),
    runtime_op(),
    list(AITA),
    runtime_backend_set(AITA)
) -> runtime_plan(AITA).
plan_runtime(Spec, Operation, Available, Backends) ->
    Candidates = candidates_for(Spec, Operation, Backends),
    Selected = begin
        _pipe = Candidates,
        _pipe@1 = gleam@list:find(
            _pipe,
            fun(Candidate) -> gleam@list:contains(Available, Candidate) end
        ),
        gleam@result:unwrap(_pipe@1, erlang:element(2, Backends))
    end,
    {runtime_plan,
        Spec,
        Operation,
        Selected,
        Candidates,
        rejected_backends(Selected, Available, Candidates, Backends),
        reason_for(Spec, Operation),
        cache_key_for(Spec, Operation)}.

-file("src/viva_tensor/runtime.gleam", 69).
?DOC(false).
-spec cache_key(runtime_plan(any())) -> binary().
cache_key(Plan) ->
    erlang:element(8, Plan).