-module(viva_tensor@optim@sparsity).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/optim/sparsity.gleam").
-export([prune_24_magnitude/1, prune_24_gradient/2, decompress/1, sparse_matmul/2, compute_metrics/2, benchmark_sparsity/0, main/0]).
-export_type([sparse24_block/0, sparse24_tensor/0, prune_metrics/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 sparse24_block() :: {sparse24_block,
        {float(), float()},
        {integer(), integer()}}.

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

-type prune_metrics() :: {prune_metrics,
        integer(),
        integer(),
        float(),
        float(),
        float()}.

-file("src/viva_tensor/optim/sparsity.gleam", 161).
?DOC(false).
-spec pad_group(list(float())) -> list(float()).
pad_group(Group) ->
    Len = erlang:length(Group),
    case Len < 4 of
        true ->
            lists:append(Group, gleam@list:repeat(+0.0, 4 - Len));

        false ->
            gleam@list:take(Group, 4)
    end.

-file("src/viva_tensor/optim/sparsity.gleam", 133).
?DOC(false).
-spec prune_group_magnitude(list(float())) -> sparse24_block().
prune_group_magnitude(Group) ->
    Indexed = gleam@list:index_map(
        Group,
        fun(Val, Idx) -> {Idx, Val, gleam@float:absolute_value(Val)} end
    ),
    Sorted = gleam@list:sort(
        Indexed,
        fun(A, B) ->
            gleam@float:compare(erlang:element(3, B), erlang:element(3, A))
        end
    ),
    case Sorted of
        [First, Second | _] ->
            {Pos1, Val1, _} = First,
            {Pos2, Val2, _} = Second,
            {P1, V1, P2, V2} = case Pos1 < Pos2 of
                true ->
                    {Pos1, Val1, Pos2, Val2};

                false ->
                    {Pos2, Val2, Pos1, Val1}
            end,
            {sparse24_block, {V1, V2}, {P1, P2}};

        _ ->
            {sparse24_block, {+0.0, +0.0}, {0, 1}}
    end.

-file("src/viva_tensor/optim/sparsity.gleam", 569).
?DOC(false).
-spec get_tensor_shape(viva_tensor@tensor:tensor()) -> list(integer()).
get_tensor_shape(T) ->
    case T of
        {tensor, _, Shape} ->
            Shape;

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

        {native_tensor, _, Shape@2} ->
            Shape@2
    end.

-file("src/viva_tensor/optim/sparsity.gleam", 103).
?DOC(false).
-spec prune_24_magnitude(viva_tensor@tensor:tensor()) -> sparse24_tensor().
prune_24_magnitude(T) ->
    Data = viva_tensor@tensor:to_list(T),
    Shape = get_tensor_shape(T),
    Num_elements = erlang:length(Data),
    Groups = gleam@list:sized_chunk(Data, 4),
    Blocks = gleam@list:map(
        Groups,
        fun(Group) -> prune_group_magnitude(pad_group(Group)) end
    ),
    Num_blocks = erlang:length(Blocks),
    Memory = Num_blocks * 5,
    {sparse24_tensor, Blocks, Shape, Num_elements, Memory, 50.0}.

-file("src/viva_tensor/optim/sparsity.gleam", 584).
?DOC(false).
-spec range_loop(integer(), integer(), list(integer())) -> list(integer()).
range_loop(From, To, Acc) ->
    case From > To of
        true ->
            lists:reverse(Acc);

        false ->
            range_loop(From + 1, To, [From | Acc])
    end.

-file("src/viva_tensor/optim/sparsity.gleam", 580).
?DOC(false).
-spec range_int(integer(), integer()) -> list(integer()).
range_int(From, To) ->
    range_loop(From, To, []).

-file("src/viva_tensor/optim/sparsity.gleam", 211).
?DOC(false).
-spec prune_group_by_importance(list(float()), list(float())) -> sparse24_block().
prune_group_by_importance(Weights, Importance) ->
    Indexed = begin
        _pipe = gleam@list:zip(
            range_int(0, 3),
            gleam@list:zip(Weights, Importance)
        ),
        gleam@list:map(
            _pipe,
            fun(X) ->
                {Idx, {W, I}} = X,
                {Idx, W, I}
            end
        )
    end,
    Sorted = gleam@list:sort(
        Indexed,
        fun(A, B) ->
            gleam@float:compare(erlang:element(3, B), erlang:element(3, A))
        end
    ),
    case Sorted of
        [First, Second | _] ->
            {Pos1, Val1, _} = First,
            {Pos2, Val2, _} = Second,
            {P1, V1, P2, V2} = case Pos1 < Pos2 of
                true ->
                    {Pos1, Val1, Pos2, Val2};

                false ->
                    {Pos2, Val2, Pos1, Val1}
            end,
            {sparse24_block, {V1, V2}, {P1, P2}};

        _ ->
            {sparse24_block, {+0.0, +0.0}, {0, 1}}
    end.

-file("src/viva_tensor/optim/sparsity.gleam", 180).
?DOC(false).
-spec prune_24_gradient(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> sparse24_tensor().
prune_24_gradient(Weights, Gradients) ->
    W_data = viva_tensor@tensor:to_list(Weights),
    G_data = viva_tensor@tensor:to_list(Gradients),
    Shape = get_tensor_shape(Weights),
    Num_elements = erlang:length(W_data),
    Importance = gleam@list:map2(
        W_data,
        G_data,
        fun(W, G) -> gleam@float:absolute_value(W * G) end
    ),
    W_groups = gleam@list:sized_chunk(W_data, 4),
    I_groups = gleam@list:sized_chunk(Importance, 4),
    Blocks = gleam@list:map2(
        W_groups,
        I_groups,
        fun(W_group, I_group) ->
            prune_group_by_importance(pad_group(W_group), I_group)
        end
    ),
    Num_blocks = erlang:length(Blocks),
    {sparse24_tensor, Blocks, Shape, Num_elements, Num_blocks * 5, 50.0}.

-file("src/viva_tensor/optim/sparsity.gleam", 245).
?DOC(false).
-spec decompress(sparse24_tensor()) -> viva_tensor@tensor:tensor().
decompress(Sparse) ->
    Data = gleam@list:flat_map(
        erlang:element(2, Sparse),
        fun(Block) ->
            {V1, V2} = erlang:element(2, Block),
            {P1, P2} = erlang:element(3, Block),
            _pipe = range_int(0, 3),
            gleam@list:map(_pipe, fun(I) -> case I =:= P1 of
                        true ->
                            V1;

                        false ->
                            case I =:= P2 of
                                true ->
                                    V2;

                                false ->
                                    +0.0
                            end
                    end end)
        end
    ),
    Truncated = gleam@list:take(Data, erlang:element(4, Sparse)),
    {tensor, Truncated, erlang:element(3, Sparse)}.

-file("src/viva_tensor/optim/sparsity.gleam", 333).
?DOC(false).
-spec transpose_matrix(list(list(float()))) -> list(list(float())).
transpose_matrix(M) ->
    case M of
        [] ->
            [];

        [First | _] ->
            N_cols = erlang:length(First),
            _pipe = range_int(0, N_cols - 1),
            gleam@list:map(
                _pipe,
                fun(Col_idx) ->
                    gleam@list:filter_map(
                        M,
                        fun(Row) -> case gleam@list:drop(Row, Col_idx) of
                                [X | _] ->
                                    {ok, X};

                                [] ->
                                    {error, nil}
                            end end
                    )
                end
            )
    end.

-file("src/viva_tensor/optim/sparsity.gleam", 302).
?DOC(false).
-spec tensor_matmul(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> viva_tensor@tensor:tensor().
tensor_matmul(A, B) ->
    A_data = viva_tensor@tensor:to_list(A),
    B_data = viva_tensor@tensor:to_list(B),
    A_shape = get_tensor_shape(A),
    B_shape = get_tensor_shape(B),
    {M, K} = case A_shape of
        [Rows, Cols] ->
            {Rows, Cols};

        _ ->
            {1, erlang:length(A_data)}
    end,
    {_, N} = case B_shape of
        [Rows@1, Cols@1] ->
            {Rows@1, Cols@1};

        _ ->
            {erlang:length(B_data), 1}
    end,
    A_rows = gleam@list:sized_chunk(A_data, K),
    B_cols = transpose_matrix(gleam@list:sized_chunk(B_data, N)),
    Result = gleam@list:flat_map(
        A_rows,
        fun(A_row) ->
            gleam@list:map(
                B_cols,
                fun(B_col) ->
                    _pipe = gleam@list:map2(
                        A_row,
                        B_col,
                        fun(X, Y) -> X * Y end
                    ),
                    gleam@list:fold(_pipe, +0.0, fun gleam@float:add/2)
                end
            )
        end
    ),
    {tensor, Result, [M, N]}.

-file("src/viva_tensor/optim/sparsity.gleam", 285).
?DOC(false).
-spec sparse_matmul(sparse24_tensor(), viva_tensor@tensor:tensor()) -> {viva_tensor@tensor:tensor(),
    float()}.
sparse_matmul(Sparse_a, Dense_b) ->
    Dense_a = decompress(Sparse_a),
    Result = tensor_matmul(Dense_a, Dense_b),
    Theoretical_speedup = 2.0,
    {Result, Theoretical_speedup}.

-file("src/viva_tensor/optim/sparsity.gleam", 358).
?DOC(false).
-spec compute_metrics(viva_tensor@tensor:tensor(), sparse24_tensor()) -> prune_metrics().
compute_metrics(Original, Sparse) ->
    Orig_data = viva_tensor@tensor:to_list(Original),
    Decomp_data = viva_tensor@tensor:to_list(decompress(Sparse)),
    Errors = gleam@list:map2(
        Orig_data,
        Decomp_data,
        fun(O, D) -> gleam@float:absolute_value(O - D) end
    ),
    Mean_error = case Errors of
        [] ->
            +0.0;

        _ ->
            case erlang:float(erlang:length(Errors)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> gleam@list:fold(
                    Errors,
                    +0.0,
                    fun gleam@float:add/2
                )
                / Gleam@denominator
            end
    end,
    Kept = gleam@list:filter_map(
        gleam@list:zip(Orig_data, Decomp_data),
        fun(Pair) ->
            {O@1, D@1} = Pair,
            case gleam@float:absolute_value(D@1) > +0.0 of
                true ->
                    {ok, gleam@float:absolute_value(O@1)};

                false ->
                    {error, nil}
            end
        end
    ),
    Pruned = gleam@list:filter_map(
        gleam@list:zip(Orig_data, Decomp_data),
        fun(Pair@1) ->
            {O@2, D@2} = Pair@1,
            case gleam@float:absolute_value(D@2) > +0.0 of
                false ->
                    {ok, gleam@float:absolute_value(O@2)};

                true ->
                    {error, nil}
            end
        end
    ),
    Kept_mean = case Kept of
        [] ->
            +0.0;

        _ ->
            case erlang:float(erlang:length(Kept)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> gleam@list:fold(
                    Kept,
                    +0.0,
                    fun gleam@float:add/2
                )
                / Gleam@denominator@1
            end
    end,
    Pruned_mean = case Pruned of
        [] ->
            +0.0;

        _ ->
            case erlang:float(erlang:length(Pruned)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@2 -> gleam@list:fold(
                    Pruned,
                    +0.0,
                    fun gleam@float:add/2
                )
                / Gleam@denominator@2
            end
    end,
    {prune_metrics,
        erlang:length(Pruned),
        erlang:length(Orig_data),
        Mean_error,
        Kept_mean,
        Pruned_mean}.

-file("src/viva_tensor/optim/sparsity.gleam", 564).
?DOC(false).
-spec float_to_string(float()) -> binary().
float_to_string(F) ->
    Rounded = erlang:float(erlang:round(F * 10000.0)) / 10000.0,
    gleam_stdlib:float_to_string(Rounded).

-file("src/viva_tensor/optim/sparsity.gleam", 417).
?DOC(false).
-spec benchmark_sparsity() -> nil.
benchmark_sparsity() ->
    gleam_stdlib:println(
        <<"╔══════════════════════════════════════════════════════════════════╗"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  2:4 STRUCTURED SPARSITY - NVIDIA Tensor Cores                   ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  Ampere+ Architecture (RTX 3000/4000, A100, H100)                ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"╚══════════════════════════════════════════════════════════════════╝\n"/utf8>>
    ),
    gleam_stdlib:println(<<"CONCEPT:"/utf8>>),
    gleam_stdlib:println(
        <<"  - Keep 2 of every 4 elements (50% structured sparsity)"/utf8>>
    ),
    gleam_stdlib:println(
        <<"  - Tensor Cores skip multiplications by zero"/utf8>>
    ),
    gleam_stdlib:println(
        <<"  - Result: 2x throughput with ~1% accuracy loss"/utf8>>
    ),
    gleam_stdlib:println(
        <<"  - Reference: Mishra et al. (2021) arXiv:2104.08378\n"/utf8>>
    ),
    gleam_stdlib:println(<<"STORAGE FORMAT:"/utf8>>),
    gleam_stdlib:println(<<"  - Original: 4 x FP16 = 64 bits"/utf8>>),
    gleam_stdlib:println(<<"  - Sparse: 2 x FP16 + 4-bit mask = 36 bits"/utf8>>),
    gleam_stdlib:println(<<"  - Compression: 1.78x\n"/utf8>>),
    T = viva_tensor@tensor:random_uniform([1024, 512]),
    gleam_stdlib:println(<<"━━━ BENCHMARK: Tensor [1024, 512] ━━━"/utf8>>),
    {Time_prune, Sparse} = timer:tc(fun() -> prune_24_magnitude(T) end),
    Metrics = compute_metrics(T, Sparse),
    Original_bytes = (1024 * 512) * 4,
    Compression = case erlang:float(erlang:element(5, Sparse)) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> erlang:float(Original_bytes) / Gleam@denominator
    end,
    gleam_stdlib:println(
        <<<<"  Prune time:         "/utf8,
                (erlang:integer_to_binary(Time_prune div 1000))/binary>>/binary,
            "ms"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  Original memory:    "/utf8,
                (erlang:integer_to_binary(Original_bytes div 1024))/binary>>/binary,
            " KB"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  Sparse memory:      "/utf8,
                (erlang:integer_to_binary(erlang:element(5, Sparse) div 1024))/binary>>/binary,
            " KB"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  Compression:        "/utf8,
                (float_to_string(Compression))/binary>>/binary,
            "x"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  Sparsity:           "/utf8,
                (float_to_string(erlang:element(6, Sparse)))/binary>>/binary,
            "%"/utf8>>
    ),
    gleam_stdlib:println(<<""/utf8>>),
    gleam_stdlib:println(
        <<<<<<"  Pruned elements:    "/utf8,
                    (erlang:integer_to_binary(erlang:element(2, Metrics)))/binary>>/binary,
                "/"/utf8>>/binary,
            (erlang:integer_to_binary(erlang:element(3, Metrics)))/binary>>
    ),
    gleam_stdlib:println(
        <<"  Approx error:       "/utf8,
            (float_to_string(erlang:element(4, Metrics)))/binary>>
    ),
    gleam_stdlib:println(
        <<"  Kept magnitude:     "/utf8,
            (float_to_string(erlang:element(5, Metrics)))/binary>>
    ),
    gleam_stdlib:println(
        <<"  Pruned magnitude:   "/utf8,
            (float_to_string(erlang:element(6, Metrics)))/binary>>
    ),
    gleam_stdlib:println(<<"\n━━━ SPARSE MATMUL SIMULATION ━━━"/utf8>>),
    B = viva_tensor@tensor:random_uniform([512, 256]),
    {Time_dense, Dense_result} = timer:tc(
        fun() -> tensor_matmul(decompress(Sparse), B) end
    ),
    {Time_sparse, {Sparse_result, Speedup}} = timer:tc(
        fun() -> sparse_matmul(Sparse, B) end
    ),
    gleam_stdlib:println(
        <<<<"  Dense matmul:       "/utf8,
                (erlang:integer_to_binary(Time_dense div 1000))/binary>>/binary,
            "ms"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  Sparse matmul:      "/utf8,
                (erlang:integer_to_binary(Time_sparse div 1000))/binary>>/binary,
            "ms (simulated)"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  Theoretical speedup: "/utf8, (float_to_string(Speedup))/binary>>/binary,
            "x (real hardware)"/utf8>>
    ),
    Dense_data = viva_tensor@tensor:to_list(Dense_result),
    Sparse_data = viva_tensor@tensor:to_list(Sparse_result),
    Diff = begin
        _pipe = gleam@list:map2(
            Dense_data,
            Sparse_data,
            fun(D, S) -> gleam@float:absolute_value(D - S) end
        ),
        gleam@list:fold(_pipe, +0.0, fun gleam@float:max/2)
    end,
    gleam_stdlib:println(
        <<<<"  Max difference:     "/utf8, (float_to_string(Diff))/binary>>/binary,
            " (should be ~0)"/utf8>>
    ),
    gleam_stdlib:println(<<"\n━━━ COMBINING TECHNIQUES ━━━"/utf8>>),
    gleam_stdlib:println(<<"  FP16 alone:         2x compression"/utf8>>),
    gleam_stdlib:println(<<"  INT8 alone:         4x compression"/utf8>>),
    gleam_stdlib:println(
        <<"  2:4 Sparsity:       2x speedup + 1.78x compression"/utf8>>
    ),
    gleam_stdlib:println(<<"  NF4 alone:          8x compression"/utf8>>),
    gleam_stdlib:println(<<""/utf8>>),
    gleam_stdlib:println(
        <<"  INT8 + 2:4:         4x * 1.78x = 7.12x compression, 8x speedup"/utf8>>
    ),
    gleam_stdlib:println(
        <<"  NF4 + 2:4:          8x * 1.78x = 14.24x compression"/utf8>>
    ),
    gleam_stdlib:println(
        <<"\n╔══════════════════════════════════════════════════════════════════╗"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  WHY 2:4 SPARSITY MATTERS:                                       ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║                                                                  ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  1. Native hardware support in RTX 3000/4000/A100/H100           ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  2. 2x throughput with ~1% accuracy loss                         ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  3. Stacks with quantization for 4x+ total speedup               ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  4. Standard in NVIDIA models (Megatron-LM, etc)                 ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║                                                                  ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  viva_tensor + 2:4 = Maximum Tensor Core utilization!            ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"╚══════════════════════════════════════════════════════════════════╝"/utf8>>
    ).

-file("src/viva_tensor/optim/sparsity.gleam", 413).
?DOC(false).
-spec main() -> nil.
main() ->
    benchmark_sparsity().