-module(viva_tensor@tflops).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/tflops.gleam").
-export([backend_name/1, best_backend/0, detect_backends/0, theoretical_peak/1, measure_matmul/4, measure_matmul_averaged/5, format_result/1, format_table/1]).
-export_type([backend/0, tflops_result/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(
    " TFLOPS - Tera Floating Point Operations Per Second\n"
    "\n"
    " Multi-platform computational throughput measurement and auto-dispatch.\n"
    " From Pure Erlang (~0.001 TFLOPS) to CUDA Sparse 2:4 (~660 TFLOPS).\n"
    "\n"
    " The `Auto` backend automatically selects the fastest available compute:\n"
    " GPU Sparse > GPU FP16 > GPU INT8 > GPU FP32 > CPU MKL > CPU SIMD > Erlang\n"
    "\n"
    " ```gleam\n"
    " import viva_tensor/tflops\n"
    "\n"
    " // Auto-select fastest backend\n"
    " let result = tflops.measure_matmul(tflops.Auto, 2048, 2048, 2048)\n"
    " io.println(tflops.format_result(result))\n"
    "\n"
    " // Benchmark all available backends\n"
    " let backends = tflops.detect_backends()\n"
    " let results = list.map(backends, fn(b) { tflops.measure_matmul(b, 1024, 1024, 1024) })\n"
    " io.println(tflops.format_table(results))\n"
    " ```\n"
).

-type backend() :: pure_erlang |
    zig_s_i_m_d |
    mkl_b_l_a_s |
    cuda_f_p32 |
    cuda_f_p16 |
    cuda_i_n_t8 |
    cuda_sparse |
    auto.

-type tflops_result() :: {tflops_result,
        backend(),
        integer(),
        integer(),
        integer(),
        float(),
        float(),
        float()}.

-file("src/viva_tensor/tflops.gleam", 204).
?DOC(" Backend name as string\n").
-spec backend_name(backend()) -> binary().
backend_name(Backend) ->
    case Backend of
        pure_erlang ->
            <<"Pure Erlang"/utf8>>;

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

        mkl_b_l_a_s ->
            <<"MKL BLAS"/utf8>>;

        cuda_f_p32 ->
            <<"CUDA FP32"/utf8>>;

        cuda_f_p16 ->
            <<"CUDA FP16"/utf8>>;

        cuda_i_n_t8 ->
            <<"CUDA INT8"/utf8>>;

        cuda_sparse ->
            <<"Sparse 2:4"/utf8>>;

        auto ->
            <<"Auto"/utf8>>
    end.

-file("src/viva_tensor/tflops.gleam", 295).
-spec run_pure_erlang(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_pure_erlang(Data_a, Data_b, M, N, K) ->
    A = viva_tensor@core@ffi:list_to_array(Data_a),
    B = viva_tensor@core@ffi:list_to_array(Data_b),
    Start = viva_tensor@core@ffi:now_microseconds(),
    _ = viva_tensor@core@ffi:array_matmul(A, B, M, N, K),
    End = viva_tensor@core@ffi:now_microseconds(),
    End - Start.

-file("src/viva_tensor/tflops.gleam", 310).
-spec run_zig_simd(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_zig_simd(Data_a, Data_b, M, N, K) ->
    Start = viva_tensor@core@ffi:now_microseconds(),
    _ = viva_tensor@core@ffi:zig_matmul(Data_a, Data_b, M, N, K),
    End = viva_tensor@core@ffi:now_microseconds(),
    End - Start.

-file("src/viva_tensor/tflops.gleam", 323).
-spec run_mkl_blas(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_mkl_blas(Data_a, Data_b, M, N, K) ->
    case viva_tensor@core@ffi:nt_from_list(Data_a, [M, K]) of
        {ok, A_nt} ->
            case viva_tensor@core@ffi:nt_from_list(Data_b, [K, N]) of
                {ok, B_nt} ->
                    Start = viva_tensor@core@ffi:now_microseconds(),
                    _ = viva_tensor@core@ffi:nt_matmul(A_nt, B_nt, M, N, K),
                    End = viva_tensor@core@ffi:now_microseconds(),
                    End - Start;

                {error, _} ->
                    0
            end;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/tflops.gleam", 346).
-spec run_cuda_fp32(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_cuda_fp32(Data_a, Data_b, M, N, K) ->
    case viva_tensor_zig:ct_from_list(Data_a, [M, K]) of
        {ok, A_ct} ->
            case viva_tensor_zig:ct_from_list(Data_b, [K, N]) of
                {ok, B_ct} ->
                    Start = viva_tensor@core@ffi:now_microseconds(),
                    _ = viva_tensor_zig:ct_matmul(A_ct, B_ct, M, N, K),
                    End = viva_tensor@core@ffi:now_microseconds(),
                    End - Start;

                {error, _} ->
                    0
            end;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/tflops.gleam", 368).
-spec run_cuda_fp16(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_cuda_fp16(Data_a, Data_b, M, N, K) ->
    case viva_tensor_zig:ct16_from_list(Data_a, [M, K]) of
        {ok, A_ct16} ->
            case viva_tensor_zig:ct16_from_list(Data_b, [K, N]) of
                {ok, B_ct16} ->
                    Start = viva_tensor@core@ffi:now_microseconds(),
                    _ = viva_tensor_zig:ct16_matmul(A_ct16, B_ct16, M, N, K),
                    End = viva_tensor@core@ffi:now_microseconds(),
                    End - Start;

                {error, _} ->
                    0
            end;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/tflops.gleam", 390).
-spec run_cuda_int8(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_cuda_int8(Data_a, Data_b, M, N, K) ->
    case viva_tensor_zig:ct_int8_from_list(Data_a, [M, K]) of
        {ok, A_int8} ->
            case viva_tensor_zig:ct_int8_from_list(Data_b, [K, N]) of
                {ok, B_int8} ->
                    Start = viva_tensor@core@ffi:now_microseconds(),
                    _ = viva_tensor_zig:ct_int8_matmul(A_int8, B_int8, M, N, K),
                    End = viva_tensor@core@ffi:now_microseconds(),
                    End - Start;

                {error, _} ->
                    0
            end;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/tflops.gleam", 412).
-spec run_cuda_sparse(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_cuda_sparse(Data_a, Data_b, M, N, K) ->
    case viva_tensor_zig:ct16_from_list(Data_a, [M, K]) of
        {ok, A_ct16} ->
            case viva_tensor_zig:sparse_from_ct16(A_ct16) of
                {ok, A_sparse} ->
                    case viva_tensor_zig:ct16_from_list(Data_b, [K, N]) of
                        {ok, B_ct16} ->
                            Start = viva_tensor@core@ffi:now_microseconds(),
                            _ = viva_tensor_zig:sparse_matmul(
                                A_sparse,
                                B_ct16,
                                M,
                                N,
                                K
                            ),
                            End = viva_tensor@core@ffi:now_microseconds(),
                            End - Start;

                        {error, _} ->
                            0
                    end;

                {error, _} ->
                    0
            end;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/tflops.gleam", 460).
-spec cuda_fp32_available() -> boolean().
cuda_fp32_available() ->
    case viva_tensor_zig:ct_from_list([1.0, +0.0, +0.0, 1.0], [2, 2]) of
        {ok, _} ->
            true;

        {error, _} ->
            false
    end.

-file("src/viva_tensor/tflops.gleam", 70).
?DOC(" Detect the fastest available backend\n").
-spec best_backend() -> backend().
best_backend() ->
    case viva_tensor_zig:sparse_available() of
        true ->
            cuda_sparse;

        false ->
            case viva_tensor_zig:ct16_available() of
                true ->
                    cuda_f_p16;

                false ->
                    case viva_tensor_zig:ct_int8_available() of
                        true ->
                            cuda_i_n_t8;

                        false ->
                            case cuda_fp32_available() of
                                true ->
                                    cuda_f_p32;

                                false ->
                                    case viva_tensor@core@ffi:zig_is_loaded() of
                                        true ->
                                            mkl_b_l_a_s;

                                        false ->
                                            pure_erlang
                                    end
                            end
                    end
            end
    end.

-file("src/viva_tensor/tflops.gleam", 160).
?DOC(" Detect all available backends (ordered slowest to fastest)\n").
-spec detect_backends() -> list(backend()).
detect_backends() ->
    Base = [pure_erlang],
    With_zig = case viva_tensor@core@ffi:zig_is_loaded() of
        true ->
            lists:append(Base, [zig_s_i_m_d, mkl_b_l_a_s]);

        false ->
            Base
    end,
    With_cuda = case cuda_fp32_available() of
        true ->
            lists:append(With_zig, [cuda_f_p32]);

        false ->
            With_zig
    end,
    With_fp16 = case viva_tensor_zig:ct16_available() of
        true ->
            lists:append(With_cuda, [cuda_f_p16]);

        false ->
            With_cuda
    end,
    With_int8 = case viva_tensor_zig:ct_int8_available() of
        true ->
            lists:append(With_fp16, [cuda_i_n_t8]);

        false ->
            With_fp16
    end,
    case viva_tensor_zig:sparse_available() of
        true ->
            lists:append(With_int8, [cuda_sparse]);

        false ->
            With_int8
    end.

-file("src/viva_tensor/tflops.gleam", 190).
?DOC(" Theoretical peak TFLOPS for a backend (RTX 4090 / i9-13900K)\n").
-spec theoretical_peak(backend()) -> float().
theoretical_peak(Backend) ->
    case Backend of
        pure_erlang ->
            0.001;

        zig_s_i_m_d ->
            1.5;

        mkl_b_l_a_s ->
            2.0;

        cuda_f_p32 ->
            82.6;

        cuda_f_p16 ->
            330.3;

        cuda_i_n_t8 ->
            660.0;

        cuda_sparse ->
            660.6;

        auto ->
            theoretical_peak(best_backend())
    end.

-file("src/viva_tensor/tflops.gleam", 264).
-spec resolve_backend(backend()) -> backend().
resolve_backend(Backend) ->
    case Backend of
        auto ->
            best_backend();

        Other ->
            Other
    end.

-file("src/viva_tensor/tflops.gleam", 275).
-spec run_matmul(
    backend(),
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> integer().
run_matmul(Backend, Data_a, Data_b, M, N, K) ->
    case Backend of
        pure_erlang ->
            run_pure_erlang(Data_a, Data_b, M, N, K);

        zig_s_i_m_d ->
            run_zig_simd(Data_a, Data_b, M, N, K);

        mkl_b_l_a_s ->
            run_mkl_blas(Data_a, Data_b, M, N, K);

        cuda_f_p32 ->
            run_cuda_fp32(Data_a, Data_b, M, N, K);

        cuda_f_p16 ->
            run_cuda_fp16(Data_a, Data_b, M, N, K);

        cuda_i_n_t8 ->
            run_cuda_int8(Data_a, Data_b, M, N, K);

        cuda_sparse ->
            run_cuda_sparse(Data_a, Data_b, M, N, K);

        auto ->
            run_matmul(best_backend(), Data_a, Data_b, M, N, K)
    end.

-file("src/viva_tensor/tflops.gleam", 438).
-spec measure_n(
    backend(),
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer(),
    integer(),
    list(integer())
) -> list(integer()).
measure_n(Backend, Data_a, Data_b, M, N, K, Remaining, Acc) ->
    case Remaining =< 0 of
        true ->
            lists:reverse(Acc);

        false ->
            Time_us = run_matmul(Backend, Data_a, Data_b, M, N, K),
            measure_n(
                Backend,
                Data_a,
                Data_b,
                M,
                N,
                K,
                Remaining - 1,
                [Time_us | Acc]
            )
    end.

-file("src/viva_tensor/tflops.gleam", 471).
-spec compute_tflops(integer(), integer()) -> float().
compute_tflops(Flops, Time_us) ->
    case Time_us > 0 of
        true ->
            case (erlang:float(Time_us) * 1000000.0) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> erlang:float(Flops) / Gleam@denominator
            end;

        false ->
            +0.0
    end.

-file("src/viva_tensor/tflops.gleam", 482).
-spec random_floats_acc(integer(), list(float())) -> list(float()).
random_floats_acc(Remaining, Acc) ->
    case Remaining =< 0 of
        true ->
            Acc;

        false ->
            random_floats_acc(
                Remaining - 1,
                [viva_tensor@core@ffi:random_uniform() | Acc]
            )
    end.

-file("src/viva_tensor/tflops.gleam", 478).
-spec random_floats(integer()) -> list(float()).
random_floats(N) ->
    random_floats_acc(N, []).

-file("src/viva_tensor/tflops.gleam", 94).
?DOC(" Measure single matmul TFLOPS for a backend\n").
-spec measure_matmul(backend(), integer(), integer(), integer()) -> tflops_result().
measure_matmul(Backend, M, N, K) ->
    Actual = resolve_backend(Backend),
    Flops = ((2 * M) * N) * K,
    Data_a = random_floats(M * K),
    Data_b = random_floats(K * N),
    Time_us = run_matmul(Actual, Data_a, Data_b, M, N, K),
    Tflops = compute_tflops(Flops, Time_us),
    Peak = theoretical_peak(Actual),
    Eff = case Peak > +0.0 of
        true ->
            (case Peak of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> Tflops / Gleam@denominator
            end) * 100.0;

        false ->
            +0.0
    end,
    {tflops_result, Actual, M, Flops, Time_us, Tflops, Tflops * 1000.0, Eff}.

-file("src/viva_tensor/tflops.gleam", 120).
?DOC(" Measure averaged TFLOPS (warmup + iterations)\n").
-spec measure_matmul_averaged(
    backend(),
    integer(),
    integer(),
    integer(),
    integer()
) -> tflops_result().
measure_matmul_averaged(Backend, M, N, K, Iterations) ->
    Actual = resolve_backend(Backend),
    Flops = ((2 * M) * N) * K,
    Data_a = random_floats(M * K),
    Data_b = random_floats(K * N),
    _ = run_matmul(Actual, Data_a, Data_b, M, N, K),
    _ = run_matmul(Actual, Data_a, Data_b, M, N, K),
    Times = measure_n(Actual, Data_a, Data_b, M, N, K, Iterations, []),
    Total_us = gleam@list:fold(Times, 0, fun(Acc, T) -> Acc + T end),
    Avg_us = case Iterations of
        0 -> 0;
        Gleam@denominator -> Total_us div Gleam@denominator
    end,
    Tflops = compute_tflops(Flops, Avg_us),
    Peak = theoretical_peak(Actual),
    Eff = case Peak > +0.0 of
        true ->
            (case Peak of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> Tflops / Gleam@denominator@1
            end) * 100.0;

        false ->
            +0.0
    end,
    {tflops_result, Actual, M, Flops, Avg_us, Tflops, Tflops * 1000.0, Eff}.

-file("src/viva_tensor/tflops.gleam", 526).
-spec abs_int(integer()) -> integer().
abs_int(N) ->
    case N < 0 of
        true ->
            0 - N;

        false ->
            N
    end.

-file("src/viva_tensor/tflops.gleam", 505).
-spec format_float1(float()) -> binary().
format_float1(F) ->
    Rounded = erlang:round(F * 10.0),
    Whole = Rounded div 10,
    Frac = abs_int(Rounded - (Whole * 10)),
    <<<<(erlang:integer_to_binary(Whole))/binary, "."/utf8>>/binary,
        (erlang:integer_to_binary(Frac))/binary>>.

-file("src/viva_tensor/tflops.gleam", 549).
-spec repeat_string(binary(), integer()) -> binary().
repeat_string(S, N) ->
    case N =< 0 of
        true ->
            <<""/utf8>>;

        false ->
            <<S/binary, (repeat_string(S, N - 1))/binary>>
    end.

-file("src/viva_tensor/tflops.gleam", 533).
-spec pad_left_zero(binary(), integer()) -> binary().
pad_left_zero(S, Width) ->
    Len = string:length(S),
    case Len < Width of
        true ->
            <<(repeat_string(<<"0"/utf8>>, Width - Len))/binary, S/binary>>;

        false ->
            S
    end.

-file("src/viva_tensor/tflops.gleam", 512).
-spec format_float2(float()) -> binary().
format_float2(F) ->
    Rounded = erlang:round(F * 100.0),
    Whole = Rounded div 100,
    Frac = abs_int(Rounded - (Whole * 100)),
    <<<<(erlang:integer_to_binary(Whole))/binary, "."/utf8>>/binary,
        (pad_left_zero(erlang:integer_to_binary(Frac), 2))/binary>>.

-file("src/viva_tensor/tflops.gleam", 500).
-spec format_time_ms(integer()) -> binary().
format_time_ms(Time_us) ->
    Ms = erlang:float(Time_us) / 1000.0,
    format_float2(Ms).

-file("src/viva_tensor/tflops.gleam", 519).
-spec format_float3(float()) -> binary().
format_float3(F) ->
    Rounded = erlang:round(F * 1000.0),
    Whole = Rounded div 1000,
    Frac = abs_int(Rounded - (Whole * 1000)),
    <<<<(erlang:integer_to_binary(Whole))/binary, "."/utf8>>/binary,
        (pad_left_zero(erlang:integer_to_binary(Frac), 3))/binary>>.

-file("src/viva_tensor/tflops.gleam", 493).
-spec format_tflops(float()) -> binary().
format_tflops(T) ->
    case T < 0.001 of
        true ->
            <<"<0.001"/utf8>>;

        false ->
            format_float3(T)
    end.

-file("src/viva_tensor/tflops.gleam", 218).
?DOC(" Format single result as a one-line string\n").
-spec format_result(tflops_result()) -> binary().
format_result(Result) ->
    <<<<<<<<<<<<<<<<<<<<<<<<<<(backend_name(erlang:element(2, Result)))/binary,
                                                        " "/utf8>>/binary,
                                                    (erlang:integer_to_binary(
                                                        erlang:element(
                                                            3,
                                                            Result
                                                        )
                                                    ))/binary>>/binary,
                                                "x"/utf8>>/binary,
                                            (erlang:integer_to_binary(
                                                erlang:element(3, Result)
                                            ))/binary>>/binary,
                                        ": "/utf8>>/binary,
                                    (format_tflops(erlang:element(6, Result)))/binary>>/binary,
                                " TFLOPS ("/utf8>>/binary,
                            (format_float2(erlang:element(7, Result)))/binary>>/binary,
                        " GFLOPS, "/utf8>>/binary,
                    (format_float1(erlang:element(8, Result)))/binary>>/binary,
                "% eff, "/utf8>>/binary,
            (format_time_ms(erlang:element(5, Result)))/binary>>/binary,
        " ms)"/utf8>>.

-file("src/viva_tensor/tflops.gleam", 541).
-spec pad_right(binary(), integer()) -> binary().
pad_right(S, Width) ->
    Len = string:length(S),
    case Len < Width of
        true ->
            <<S/binary, (repeat_string(<<" "/utf8>>, Width - Len))/binary>>;

        false ->
            S
    end.

-file("src/viva_tensor/tflops.gleam", 236).
?DOC(" Format list of results as a table\n").
-spec format_table(list(tflops_result())) -> binary().
format_table(Results) ->
    Header = <<<<"  ┌──────────────────┬────────────┬──────────┬──────────┐\n"/utf8,
            "  │ Backend          │ Time (ms)  │ TFLOPS   │ Eff %    │\n"/utf8>>/binary,
        "  ├──────────────────┼────────────┼──────────┼──────────┤"/utf8>>,
    Rows = gleam@list:map(
        Results,
        fun(R) ->
            <<<<<<<<<<<<<<<<"  │ "/utf8,
                                            (pad_right(
                                                backend_name(
                                                    erlang:element(2, R)
                                                ),
                                                16
                                            ))/binary>>/binary,
                                        " │ "/utf8>>/binary,
                                    (pad_right(
                                        format_time_ms(erlang:element(5, R)),
                                        10
                                    ))/binary>>/binary,
                                " │ "/utf8>>/binary,
                            (pad_right(format_tflops(erlang:element(6, R)), 8))/binary>>/binary,
                        " │ "/utf8>>/binary,
                    (pad_right(
                        <<(format_float1(erlang:element(8, R)))/binary,
                            "%"/utf8>>,
                        8
                    ))/binary>>/binary,
                " │"/utf8>>
        end
    ),
    Footer = <<"  └──────────────────┴────────────┴──────────┴──────────┘"/utf8>>,
    gleam@string:join([Header | lists:append(Rows, [Footer])], <<"\n"/utf8>>).