viva_tensor

-module(viva_tensor@backend@protocol).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/backend/protocol.gleam").
-export([auto_select/0, is_available/1, name/1, info/1, matmul/6, dot/3, sum/2, scale/3, add/3]).
-export_type([backend/0, node_/0, row_shard/0, task_ref/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 backend() :: pure | accelerate | zig | {distributed, list(node_())}.

-type node_() :: {node, binary()}.

-type row_shard() :: {row_shard, list(float()), integer()}.

-type task_ref() :: any().

-file("src/viva_tensor/backend/protocol.gleam", 69).
?DOC(false).
-spec auto_select() -> backend().
auto_select() ->
    case viva_tensor@core@ffi:zig_is_loaded() of
        true ->
            zig;

        false ->
            case viva_tensor@core@ffi:is_nif_loaded() of
                true ->
                    accelerate;

                false ->
                    pure
            end
    end.

-file("src/viva_tensor/backend/protocol.gleam", 83).
?DOC(false).
-spec is_available(backend()) -> boolean().
is_available(Backend) ->
    case Backend of
        pure ->
            true;

        accelerate ->
            viva_tensor@core@ffi:is_nif_loaded();

        zig ->
            viva_tensor@core@ffi:zig_is_loaded();

        {distributed, Nodes} ->
            Nodes /= []
    end.

-file("src/viva_tensor/backend/protocol.gleam", 93).
?DOC(false).
-spec name(backend()) -> binary().
name(Backend) ->
    case Backend of
        pure ->
            <<"Pure Erlang"/utf8>>;

        accelerate ->
            <<"Apple Accelerate"/utf8>>;

        zig ->
            <<"Zig SIMD"/utf8>>;

        {distributed, _} ->
            <<"Distributed BEAM"/utf8>>
    end.

-file("src/viva_tensor/backend/protocol.gleam", 103).
?DOC(false).
-spec info(backend()) -> binary().
info(Backend) ->
    case Backend of
        pure ->
            <<"Pure Erlang with O(1) array access (~100 MFLOPS)"/utf8>>;

        accelerate ->
            viva_tensor@core@ffi:nif_backend_info();

        zig ->
            viva_tensor@core@ffi:zig_backend_info();

        {distributed, Nodes} ->
            <<<<"Distributed across "/utf8,
                    (erlang:integer_to_list(erlang:length(Nodes)))/binary>>/binary,
                " nodes"/utf8>>
    end.

-file("src/viva_tensor/backend/protocol.gleam", 379).
?DOC(false).
-spec collect_results(list(task_ref()), list(list(float()))) -> {ok,
        list(list(float()))} |
    {error, binary()}.
collect_results(Tasks, Acc) ->
    case Tasks of
        [] ->
            {ok, lists:reverse(Acc)};

        [Task | Rest] ->
            case viva_tensor_distributed:await_task(Task) of
                {ok, Result} ->
                    collect_results(Rest, [Result | Acc]);

                {error, E} ->
                    {error, E}
            end
    end.

-file("src/viva_tensor/backend/protocol.gleam", 368).
?DOC(false).
-spec spawn_matmul_task(
    node_(),
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> task_ref().
spawn_matmul_task(Node, A, B, M, N, K) ->
    viva_tensor_distributed:spawn_matmul_task(Node, A, B, M, N, K).

-file("src/viva_tensor/backend/protocol.gleam", 353).
?DOC(false).
-spec list_split_acc(list(AMBN), integer(), list(AMBN)) -> {list(AMBN),
    list(AMBN)}.
list_split_acc(Lst, N, Acc) ->
    case N =< 0 of
        true ->
            {lists:reverse(Acc), Lst};

        false ->
            case Lst of
                [] ->
                    {lists:reverse(Acc), []};

                [Head | Tail] ->
                    list_split_acc(Tail, N - 1, [Head | Acc])
            end
    end.

-file("src/viva_tensor/backend/protocol.gleam", 349).
?DOC(false).
-spec list_split(list(AMBJ), integer()) -> {list(AMBJ), list(AMBJ)}.
list_split(Lst, N) ->
    list_split_acc(Lst, N, []).

-file("src/viva_tensor/backend/protocol.gleam", 313).
?DOC(false).
-spec create_row_shards_acc(
    list(float()),
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    list(row_shard())
) -> list(row_shard()).
create_row_shards_acc(A, K, Rows_per_node, Remainder, Node_count, Current, Acc) ->
    case Current >= Node_count of
        true ->
            lists:reverse(Acc);

        false ->
            Extra = case Current < Remainder of
                true ->
                    1;

                false ->
                    0
            end,
            Rows = Rows_per_node + Extra,
            Elements = Rows * K,
            {Shard_data, Rest} = list_split(A, Elements),
            Shard = {row_shard, Shard_data, Rows},
            create_row_shards_acc(
                Rest,
                K,
                Rows_per_node,
                Remainder,
                Node_count,
                Current + 1,
                [Shard | Acc]
            )
    end.

-file("src/viva_tensor/backend/protocol.gleam", 303).
?DOC(false).
-spec create_row_shards(
    list(float()),
    integer(),
    integer(),
    integer(),
    integer()
) -> list(row_shard()).
create_row_shards(A, K, Rows_per_node, Remainder, Node_count) ->
    create_row_shards_acc(A, K, Rows_per_node, Remainder, Node_count, 0, []).

-file("src/viva_tensor/backend/protocol.gleam", 266).
?DOC(false).
-spec distributed_matmul(
    list(node_()),
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> {ok, list(float())} | {error, binary()}.
distributed_matmul(Nodes, A, B, M, N, K) ->
    Node_count = erlang:length(Nodes),
    case Node_count of
        0 ->
            {error, <<"No nodes available for distributed computation"/utf8>>};

        _ ->
            Rows_per_node = case Node_count of
                0 -> 0;
                Gleam@denominator -> M div Gleam@denominator
            end,
            Remainder = case Node_count of
                0 -> 0;
                Gleam@denominator@1 -> M rem Gleam@denominator@1
            end,
            Shards = create_row_shards(
                A,
                K,
                Rows_per_node,
                Remainder,
                Node_count
            ),
            Tasks = gleam@list:map2(
                Nodes,
                Shards,
                fun(Node, Shard) ->
                    spawn_matmul_task(
                        Node,
                        erlang:element(2, Shard),
                        B,
                        erlang:element(3, Shard),
                        N,
                        K
                    )
                end
            ),
            _pipe = collect_results(Tasks, []),
            gleam@result:map(_pipe, fun lists:append/1)
    end.

-file("src/viva_tensor/backend/protocol.gleam", 233).
?DOC(false).
-spec pure_matmul(list(float()), list(float()), integer(), integer(), integer()) -> {ok,
        list(float())} |
    {error, binary()}.
pure_matmul(A, B, M, N, K) ->
    A_arr = viva_tensor@core@ffi:list_to_array(A),
    B_arr = viva_tensor@core@ffi:list_to_array(B),
    Result_arr = viva_tensor@core@ffi:array_matmul(A_arr, B_arr, M, N, K),
    {ok, viva_tensor@core@ffi:array_to_list(Result_arr)}.

-file("src/viva_tensor/backend/protocol.gleam", 129).
?DOC(false).
-spec matmul(
    backend(),
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> {ok, list(float())} | {error, binary()}.
matmul(Backend, A, B, M, N, K) ->
    case Backend of
        pure ->
            pure_matmul(A, B, M, N, K);

        accelerate ->
            viva_tensor@core@ffi:nif_matmul(A, B, M, N, K);

        zig ->
            viva_tensor@core@ffi:zig_matmul(A, B, M, N, K);

        {distributed, Nodes} ->
            distributed_matmul(Nodes, A, B, M, N, K)
    end.

-file("src/viva_tensor/backend/protocol.gleam", 393).
?DOC(false).
-spec auto_select_local() -> backend().
auto_select_local() ->
    case viva_tensor@core@ffi:zig_is_loaded() of
        true ->
            zig;

        false ->
            case viva_tensor@core@ffi:is_nif_loaded() of
                true ->
                    accelerate;

                false ->
                    pure
            end
    end.

-file("src/viva_tensor/backend/protocol.gleam", 214).
?DOC(false).
-spec pure_dot(list(float()), list(float())) -> float().
pure_dot(A, B) ->
    A_arr = viva_tensor@core@ffi:list_to_array(A),
    B_arr = viva_tensor@core@ffi:list_to_array(B),
    viva_tensor@core@ffi:array_dot(A_arr, B_arr).

-file("src/viva_tensor/backend/protocol.gleam", 150).
?DOC(false).
-spec dot(backend(), list(float()), list(float())) -> {ok, float()} |
    {error, binary()}.
dot(Backend, A, B) ->
    case Backend of
        pure ->
            {ok, pure_dot(A, B)};

        accelerate ->
            viva_tensor@core@ffi:nif_dot(A, B);

        zig ->
            viva_tensor@core@ffi:zig_dot(A, B);

        {distributed, _} ->
            dot(auto_select_local(), A, B)
    end.

-file("src/viva_tensor/backend/protocol.gleam", 220).
?DOC(false).
-spec pure_sum(list(float())) -> float().
pure_sum(Data) ->
    Arr = viva_tensor@core@ffi:list_to_array(Data),
    viva_tensor@core@ffi:array_sum(Arr).

-file("src/viva_tensor/backend/protocol.gleam", 166).
?DOC(false).
-spec sum(backend(), list(float())) -> {ok, float()} | {error, binary()}.
sum(Backend, Data) ->
    case Backend of
        pure ->
            {ok, pure_sum(Data)};

        accelerate ->
            viva_tensor@core@ffi:nif_sum(Data);

        zig ->
            viva_tensor@core@ffi:zig_sum(Data);

        {distributed, _} ->
            sum(auto_select_local(), Data)
    end.

-file("src/viva_tensor/backend/protocol.gleam", 225).
?DOC(false).
-spec pure_scale(list(float()), float()) -> list(float()).
pure_scale(Data, Scalar) ->
    gleam@list:map(Data, fun(X) -> X * Scalar end).

-file("src/viva_tensor/backend/protocol.gleam", 176).
?DOC(false).
-spec scale(backend(), list(float()), float()) -> {ok, list(float())} |
    {error, binary()}.
scale(Backend, Data, Scalar) ->
    case Backend of
        pure ->
            {ok, pure_scale(Data, Scalar)};

        accelerate ->
            viva_tensor@core@ffi:nif_scale(Data, Scalar);

        zig ->
            viva_tensor@core@ffi:zig_scale(Data, Scalar);

        {distributed, _} ->
            scale(auto_select_local(), Data, Scalar)
    end.

-file("src/viva_tensor/backend/protocol.gleam", 229).
?DOC(false).
-spec pure_add(list(float()), list(float())) -> list(float()).
pure_add(A, B) ->
    gleam@list:map2(A, B, fun(X, Y) -> X + Y end).

-file("src/viva_tensor/backend/protocol.gleam", 190).
?DOC(false).
-spec add(backend(), list(float()), list(float())) -> {ok, list(float())} |
    {error, binary()}.
add(Backend, A, B) ->
    case Backend of
        pure ->
            {ok, pure_add(A, B)};

        accelerate ->
            {ok, pure_add(A, B)};

        zig ->
            viva_tensor@core@ffi:zig_add(A, B);

        {distributed, _} ->
            add(auto_select_local(), A, B)
    end.