viva_tensor

-module(viva_tensor@bench@rtx).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/bench/rtx.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/rtx.gleam", 176).
?DOC(false).
-spec pad_right(binary(), integer()) -> binary().
pad_right(Value, Width) ->
    Size = string:length(Value),
    case Size >= Width of
        true ->
            Value;

        false ->
            <<Value/binary,
                (gleam@string:repeat(<<" "/utf8>>, Width - Size))/binary>>
    end.

-file("src/viva_tensor/bench/rtx.gleam", 166).
?DOC(false).
-spec float_to_string_1(float()) -> binary().
float_to_string_1(Value) ->
    Rounded = erlang:float(erlang:round(Value * 10.0)) / 10.0,
    gleam_stdlib:float_to_string(Rounded).

-file("src/viva_tensor/bench/rtx.gleam", 184).
?DOC(false).
-spec pad_left(binary(), integer()) -> binary().
pad_left(Value, Width) ->
    Size = string:length(Value),
    case Size >= Width of
        true ->
            Value;

        false ->
            <<(gleam@string:repeat(<<" "/utf8>>, Width - Size))/binary,
                Value/binary>>
    end.

-file("src/viva_tensor/bench/rtx.gleam", 171).
?DOC(false).
-spec float_to_string_3(float()) -> binary().
float_to_string_3(Value) ->
    Rounded = erlang:float(erlang:round(Value * 1000.0)) / 1000.0,
    gleam_stdlib:float_to_string(Rounded).

-file("src/viva_tensor/bench/rtx.gleam", 151).
?DOC(false).
-spec matmul_gflops(integer(), float()) -> float().
matmul_gflops(N, Avg_us) ->
    Nf = erlang:float(N),
    Flops = ((2.0 * Nf) * Nf) * Nf,
    (case Avg_us of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> Flops / Gleam@denominator
    end) / 1000.0.

-file("src/viva_tensor/bench/rtx.gleam", 130).
?DOC(false).
-spec run_matmul_case(binary(), integer(), integer(), fun(() -> nil)) -> nil.
run_matmul_case(Label, N, Iterations, F) ->
    F(),
    Start = viva_tensor@core@ffi:now_microseconds(),
    _pipe = gleam@list:range(1, Iterations),
    gleam@list:each(_pipe, fun(_) -> F() end),
    Stop = viva_tensor@core@ffi:now_microseconds(),
    Avg_us = case erlang:float(Iterations) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> erlang:float(Stop - Start) / Gleam@denominator
    end,
    Gflops = matmul_gflops(N, Avg_us),
    gleam_stdlib:println(
        <<<<<<<<<<"  "/utf8, (pad_right(Label, 32))/binary>>/binary,
                        (pad_left(float_to_string_3(Avg_us / 1000.0), 10))/binary>>/binary,
                    " ms  "/utf8>>/binary,
                (pad_left(float_to_string_1(Gflops), 10))/binary>>/binary,
            " GFLOPS"/utf8>>
    ).

-file("src/viva_tensor/bench/rtx.gleam", 72).
?DOC(false).
-spec bench_rtx(
    integer(),
    integer(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> nil.
bench_rtx(N, Iterations, A, B, Bias) ->
    case {viva_tensor:to_rtx4090_fp16(A),
        viva_tensor:to_rtx4090_fp16(B),
        viva_tensor:to_rtx4090_fp16(Bias),
        viva_tensor:to_rtx4090_fp16(viva_tensor:zeros([N, N])),
        viva_tensor:gpu_workspace()} of
        {{ok, A_gpu},
            {ok, B_gpu},
            {ok, Bias_gpu},
            {ok, Out_gpu},
            {ok, Workspace}} ->
            run_matmul_case(
                <<"rtx matmul_auto upload/call"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:matmul_auto(A, B),
                    _ = viva_tensor:accelerated_sync(),
                    nil
                end
            ),
            run_matmul_case(
                <<"rtx fp16 persistent alloc"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:matmul_accelerated(A_gpu, B_gpu),
                    _ = viva_tensor:accelerated_sync(),
                    nil
                end
            ),
            run_matmul_case(
                <<"rtx fp16 matmul_into"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:matmul_accelerated_into(
                        Out_gpu,
                        A_gpu,
                        B_gpu
                    ),
                    _ = viva_tensor:accelerated_sync(),
                    nil
                end
            ),
            run_matmul_case(
                <<"rtx fp16 linear_relu_into"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:linear_relu_accelerated_into(
                        Out_gpu,
                        A_gpu,
                        B_gpu,
                        Bias_gpu
                    ),
                    _ = viva_tensor:accelerated_sync(),
                    nil
                end
            ),
            case {viva_tensor:workspace_from_tensor(Workspace, A),
                viva_tensor:linear_layer(Workspace, B, Bias)} of
                {{ok, Input}, {ok, Layer}} ->
                    case viva_tensor:linear_output(Workspace, Layer, N) of
                        {ok, Layer_out} ->
                            run_matmul_case(
                                <<"rtx workspace linear layer"/utf8>>,
                                N,
                                Iterations,
                                fun() ->
                                    _ = viva_tensor:linear_relu_forward_into(
                                        Layer_out,
                                        Input,
                                        Layer
                                    ),
                                    _ = viva_tensor:accelerated_sync(),
                                    nil
                                end
                            );

                        {error, _} ->
                            nil
                    end;

                {_, _} ->
                    nil
            end;

        {_, _, _, _, _} ->
            gleam_stdlib:println(
                <<<<"  "/utf8, (pad_right(<<"rtx fp16"/utf8>>, 32))/binary>>/binary,
                    "unavailable"/utf8>>
            )
    end.

-file("src/viva_tensor/bench/rtx.gleam", 43).
?DOC(false).
-spec bench_mkl(
    integer(),
    integer(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> nil.
bench_mkl(N, Iterations, A, B, Bias) ->
    case {viva_tensor:native_from_list(viva_tensor:to_list(A), [N, N]),
        viva_tensor:native_from_list(viva_tensor:to_list(B), [N, N]),
        viva_tensor:native_from_list(viva_tensor:to_list(Bias), [N]),
        viva_tensor:native_zeros([N, N])} of
        {{ok, A_native}, {ok, B_native}, {ok, Bias_native}, {ok, Out_native}} ->
            run_matmul_case(
                <<"mkl native matmul alloc"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:matmul(A_native, B_native),
                    nil
                end
            ),
            run_matmul_case(
                <<"mkl native matmul_into"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:matmul_into(Out_native, A_native, B_native),
                    nil
                end
            ),
            run_matmul_case(
                <<"mkl native linear_relu_into"/utf8>>,
                N,
                Iterations,
                fun() ->
                    _ = viva_tensor:linear_relu_into(
                        Out_native,
                        A_native,
                        B_native,
                        Bias_native
                    ),
                    nil
                end
            );

        {_, _, _, _} ->
            gleam_stdlib:println(
                <<<<"  "/utf8, (pad_right(<<"mkl native"/utf8>>, 32))/binary>>/binary,
                    "unavailable"/utf8>>
            )
    end.

-file("src/viva_tensor/bench/rtx.gleam", 157).
?DOC(false).
-spec iterations_for(integer()) -> integer().
iterations_for(N) ->
    case N of
        128 ->
            40;

        256 ->
            25;

        512 ->
            10;

        _ ->
            5
    end.

-file("src/viva_tensor/bench/rtx.gleam", 26).
?DOC(false).
-spec bench_square_size(integer()) -> nil.
bench_square_size(N) ->
    Iterations = iterations_for(N),
    Shape = [N, N],
    Bias_shape = [N],
    A = viva_tensor:ones(Shape),
    B = viva_tensor:ones(Shape),
    Bias = viva_tensor:ones(Bias_shape),
    gleam_stdlib:println(
        <<<<<<<<"━━━ "/utf8, (erlang:integer_to_binary(N))/binary>>/binary,
                    "x"/utf8>>/binary,
                (erlang:integer_to_binary(N))/binary>>/binary,
            " ━━━"/utf8>>
    ),
    gleam_stdlib:println(
        <<"iterations: "/utf8, (erlang:integer_to_binary(Iterations))/binary>>
    ),
    bench_mkl(N, Iterations, A, B, Bias),
    bench_rtx(N, Iterations, A, B, Bias),
    gleam_stdlib:println(<<""/utf8>>).

-file("src/viva_tensor/bench/rtx.gleam", 14).
?DOC(false).
-spec main() -> nil.
main() ->
    gleam_stdlib:println(
        <<"╔════════════════════════════════════════════════════════════╗"/utf8>>
    ),
    gleam_stdlib:println(
        <<"║  viva_tensor RTX 4090 vs MKL benchmark                    ║"/utf8>>
    ),
    gleam_stdlib:println(
        <<"╚════════════════════════════════════════════════════════════╝"/utf8>>
    ),
    gleam_stdlib:println(<<""/utf8>>),
    gleam_stdlib:println(
        <<"GPU timings call accelerated_sync() after every operation."/utf8>>
    ),
    gleam_stdlib:println(
        <<"That measures completed work, not just CUDA enqueue latency."/utf8>>
    ),
    gleam_stdlib:println(<<""/utf8>>),
    gleam@list:each([128, 256, 512, 1024], fun bench_square_size/1).