viva_tensor

-module(viva_tensor@bench@tflops).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/bench/tflops.gleam").
-export([main/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).

-file("src/viva_tensor/bench/tflops.gleam", 262).
?DOC(false).
-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/bench/tflops.gleam", 246).
?DOC(false).
-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/bench/tflops.gleam", 239).
?DOC(false).
-spec abs_int(integer()) -> integer().
abs_int(N) ->
    case N < 0 of
        true ->
            0 - N;

        false ->
            N
    end.

-file("src/viva_tensor/bench/tflops.gleam", 232).
?DOC(false).
-spec float_to_str2(float()) -> binary().
float_to_str2(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/bench/tflops.gleam", 214).
?DOC(false).
-spec format_peak(float()) -> binary().
format_peak(Peak) ->
    case Peak < 0.01 of
        true ->
            <<"<0.01"/utf8>>;

        false ->
            float_to_str2(Peak)
    end.

-file("src/viva_tensor/bench/tflops.gleam", 254).
?DOC(false).
-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/bench/tflops.gleam", 177).
?DOC(false).
-spec print_theoretical_peaks(list(viva_tensor@tflops:backend())) -> nil.
print_theoretical_peaks(Backends) ->
    gleam_stdlib:println(<<"  Theoretical peaks:"/utf8>>),
    gleam@list:each(
        Backends,
        fun(B) ->
            Peak = viva_tensor@tflops:theoretical_peak(B),
            gleam_stdlib:println(
                <<<<<<"    "/utf8,
                            (pad_right(viva_tensor@tflops:backend_name(B), 16))/binary>>/binary,
                        (format_peak(Peak))/binary>>/binary,
                    " TFLOPS"/utf8>>
            )
        end
    ).

-file("src/viva_tensor/bench/tflops.gleam", 113).
?DOC(false).
-spec filter_backends_for_size(list(viva_tensor@tflops:backend()), integer()) -> list(viva_tensor@tflops:backend()).
filter_backends_for_size(Backends, Size) ->
    gleam@list:filter(Backends, fun(B) -> case B of
                pure_erlang ->
                    Size =< 512;

                _ ->
                    true
            end end).

-file("src/viva_tensor/bench/tflops.gleam", 197).
?DOC(false).
-spec format_flops(integer()) -> binary().
format_flops(Flops) ->
    case Flops >= 1000000000 of
        true ->
            Gf = erlang:float(Flops) / 1000000000.0,
            <<(float_to_str2(Gf))/binary, "G"/utf8>>;

        false ->
            case Flops >= 1000000 of
                true ->
                    Mf = erlang:float(Flops) / 1000000.0,
                    <<(float_to_str2(Mf))/binary, "M"/utf8>>;

                false ->
                    erlang:integer_to_binary(Flops)
            end
    end.

-file("src/viva_tensor/bench/tflops.gleam", 83).
?DOC(false).
-spec benchmark_size(integer(), list(viva_tensor@tflops:backend()), integer()) -> nil.
benchmark_size(Size, Backends, Iterations) ->
    Flops = ((2 * Size) * Size) * Size,
    Flops_str = format_flops(Flops),
    gleam_stdlib:println(
        <<<<<<<<<<<<<<<<<<<<"  "/utf8, (erlang:integer_to_binary(Size))/binary>>/binary,
                                            "x"/utf8>>/binary,
                                        (erlang:integer_to_binary(Size))/binary>>/binary,
                                    " @ "/utf8>>/binary,
                                (erlang:integer_to_binary(Size))/binary>>/binary,
                            "x"/utf8>>/binary,
                        (erlang:integer_to_binary(Size))/binary>>/binary,
                    " ("/utf8>>/binary,
                Flops_str/binary>>/binary,
            " FLOPs):"/utf8>>
    ),
    Effective_backends = filter_backends_for_size(Backends, Size),
    Results = gleam@list:map(
        Effective_backends,
        fun(Backend) ->
            viva_tensor@tflops:measure_matmul_averaged(
                Backend,
                Size,
                Size,
                Size,
                Iterations
            )
        end
    ),
    gleam_stdlib:println(viva_tensor@tflops:format_table(Results)),
    gleam_stdlib:println(<<""/utf8>>).

-file("src/viva_tensor/bench/tflops.gleam", 190).
?DOC(false).
-spec status_icon(boolean()) -> binary().
status_icon(Available) ->
    case Available of
        true ->
            <<"+"/utf8>>;

        false ->
            <<"-"/utf8>>
    end.

-file("src/viva_tensor/bench/tflops.gleam", 221).
?DOC(false).
-spec cuda_available() -> boolean().
cuda_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/bench/tflops.gleam", 127).
?DOC(false).
-spec print_backends() -> nil.
print_backends() ->
    gleam_stdlib:println(<<"  BACKENDS:"/utf8>>),
    gleam_stdlib:println(
        <<<<"    "/utf8, (status_icon(true))/binary>>/binary,
            " Pure Erlang    (always available)"/utf8>>
    ),
    Zig = viva_tensor@core@ffi:zig_is_loaded(),
    Zig_info = case Zig of
        true ->
            viva_tensor@core@ffi:zig_backend_info();

        false ->
            <<"not compiled"/utf8>>
    end,
    gleam_stdlib:println(
        <<<<<<<<"    "/utf8, (status_icon(Zig))/binary>>/binary,
                    " Zig SIMD       ("/utf8>>/binary,
                Zig_info/binary>>/binary,
            ")"/utf8>>
    ),
    Cuda = cuda_available(),
    gleam_stdlib:println(
        <<<<<<"    "/utf8, (status_icon(Cuda))/binary>>/binary,
                " CUDA FP32      "/utf8>>/binary,
            (case Cuda of
                true ->
                    <<"(GPU detected)"/utf8>>;

                false ->
                    <<"(no GPU)"/utf8>>
            end)/binary>>
    ),
    Fp16 = viva_tensor_zig:ct16_available(),
    gleam_stdlib:println(
        <<<<<<"    "/utf8, (status_icon(Fp16))/binary>>/binary,
                " CUDA FP16      "/utf8>>/binary,
            (case Fp16 of
                true ->
                    <<"(Tensor Cores)"/utf8>>;

                false ->
                    <<"(unavailable)"/utf8>>
            end)/binary>>
    ),
    Sparse = viva_tensor_zig:sparse_available(),
    gleam_stdlib:println(
        <<<<<<"    "/utf8, (status_icon(Sparse))/binary>>/binary,
                " Sparse 2:4     "/utf8>>/binary,
            (case Sparse of
                true ->
                    <<"(cuSPARSELt)"/utf8>>;

                false ->
                    <<"(unavailable)"/utf8>>
            end)/binary>>
    ).

-file("src/viva_tensor/bench/tflops.gleam", 19).
?DOC(false).
-spec main() -> nil.
main() ->
    gleam_stdlib:println(<<""/utf8>>),
    gleam_stdlib:println(
        <<"╔═══════════════════════════════════════════════════════════╗"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║              viva_tensor - TFLOPS BENCHMARK              ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"╚═══════════════════════════════════════════════════════════╝"/utf8>>
    ),
    gleam_stdlib:println(<<""/utf8>>),
    print_backends(),
    gleam_stdlib:println(<<""/utf8>>),
    Backends = viva_tensor@tflops:detect_backends(),
    Iterations = 10,
    Sizes = [256, 512, 1024, 2048, 4096],
    gleam_stdlib:println(
        <<"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"/utf8>>
    ),
    gleam_stdlib:println(
        <<<<"  MATMUL BENCHMARK (warmup 2 + "/utf8,
                (erlang:integer_to_binary(Iterations))/binary>>/binary,
            " iterations)"/utf8>>
    ),
    gleam_stdlib:println(
        <<"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"/utf8>>
    ),
    gleam_stdlib:println(<<""/utf8>>),
    gleam@list:each(
        Sizes,
        fun(Size) -> benchmark_size(Size, Backends, Iterations) end
    ),
    gleam_stdlib:println(
        <<"═══════════════════════════════════════════════════════════"/utf8>>
    ),
    gleam_stdlib:println(<<"  Peak TFLOPS (largest matrix per backend):"/utf8>>),
    gleam_stdlib:println(
        <<"═══════════════════════════════════════════════════════════"/utf8>>
    ),
    gleam_stdlib:println(<<""/utf8>>),
    Fast_backends = filter_backends_for_size(Backends, 4096),
    Peak_results = gleam@list:map(
        Fast_backends,
        fun(Backend) ->
            viva_tensor@tflops:measure_matmul_averaged(
                Backend,
                4096,
                4096,
                4096,
                Iterations
            )
        end
    ),
    gleam@list:each(
        Peak_results,
        fun(R) ->
            gleam_stdlib:println(
                <<"  "/utf8, (viva_tensor@tflops:format_result(R))/binary>>
            )
        end
    ),
    case gleam@list:contains(Backends, pure_erlang) of
        true ->
            Erlang_peak = viva_tensor@tflops:measure_matmul_averaged(
                pure_erlang,
                512,
                512,
                512,
                Iterations
            ),
            gleam_stdlib:println(
                <<"  "/utf8,
                    (viva_tensor@tflops:format_result(Erlang_peak))/binary>>
            );

        false ->
            nil
    end,
    gleam_stdlib:println(<<""/utf8>>),
    print_theoretical_peaks(Backends),
    gleam_stdlib:println(<<""/utf8>>),
    gleam_stdlib:println(<<"  BENCHMARK COMPLETE!"/utf8>>),
    gleam_stdlib:println(<<""/utf8>>).