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, dot/3, sum/2, scale/3, add/3, matmul/6, info/1]).
-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(
    " Backend Protocol - Pluggable tensor computation backends\n"
    "\n"
    " The BEAM's actor model makes distributed tensor sharding natural.\n"
    " Each node is just a process - no special distributed runtime needed.\n"
    " This is why Erlang/Elixir ML libraries can scale horizontally with\n"
    " minimal ceremony compared to MPI-based frameworks.\n"
    "\n"
    " Performance reality (measured on M1 MacBook Pro, 1024x1024 matmul):\n"
    " - Pure Erlang: ~100 MFLOPS (lists are not contiguous memory)\n"
    " - Apple Accelerate: ~50 GFLOPS (500x faster - that's BLAS for you)\n"
    " - Zig SIMD: ~40 GFLOPS (portable, nearly as fast as vendor libs)\n"
    "\n"
    " Priority: Zig > Accelerate > Pure\n"
    " Why? SIMD everywhere > Apple-only > slow but portable.\n"
    " Zig NIFs compile to native code with explicit SIMD intrinsics,\n"
    " work on Linux/Windows/macOS, and approach vendor library speed.\n"
    "\n"
    " Distributed overhead: only worth it for matrices > 10K x 10K.\n"
    " Below that, network latency dominates compute time.\n"
    " The BEAM makes it easy, but easy != free.\n"
    "\n"
    " Usage:\n"
    "   let backend = backend.auto_select()\n"
    "   let result = backend.matmul(a, b, m, n, k)\n"
).

-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(
    " Automatically select the best available backend\n"
    "\n"
    " Priority: Zig > Accelerate > Pure\n"
    " Rationale:\n"
    " - Zig: portable SIMD, works everywhere, ~40 GFLOPS\n"
    " - Accelerate: Apple-specific but highly optimized\n"
    " - Pure: fallback, always works, predictable (if slow)\n"
).
-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(
    " Check if a specific backend is available\n"
    "\n"
    " Used for graceful degradation and testing\n"
).
-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(" Get human-readable backend name\n").
-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", 214).
-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", 220).
-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", 225).
-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", 229).
-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", 233).
-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", 353).
-spec list_split_acc(list(HMY), integer(), list(HMY)) -> {list(HMY), list(HMY)}.
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).
-spec list_split(list(HMU), integer()) -> {list(HMU), list(HMU)}.
list_split(Lst, N) ->
    list_split_acc(Lst, N, []).

-file("src/viva_tensor/backend/protocol.gleam", 313).
-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).
-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", 393).
-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", 150).
?DOC(
    " Dot product using selected backend\n"
    "\n"
    " For distributed: falls back to local backend.\n"
    " Why? Communication overhead > compute for O(n) operations.\n"
    " Only parallelize when compute dominates communication.\n"
).
-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", 166).
?DOC(" Sum reduction using selected backend\n").
-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", 176).
?DOC(" Scale (multiply by scalar) using selected backend\n").
-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", 190).
?DOC(" Element-wise addition using selected backend\n").
-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.

-file("src/viva_tensor/backend/protocol.gleam", 368).
-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", 379).
-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", 266).
?DOC(
    " Distributed matrix multiplication with row sharding\n"
    "\n"
    " Splits matrix A by rows across nodes, broadcasts B to all.\n"
    " Simple and works well for square-ish matrices.\n"
).
-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", 129).
?DOC(
    " Matrix multiplication using selected backend\n"
    " A[m,k] @ B[k,n] -> C[m,n]\n"
    "\n"
    " Complexity: O(m*n*k) FLOPs\n"
    " Memory: O(m*n) for result\n"
    "\n"
    " Strassen/Winograd variants not implemented - the constant factors\n"
    " only win for matrices > 1000x1000, and BLAS is already optimized.\n"
).
-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", 103).
?DOC(" Get detailed backend info including version/capability strings\n").
-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.