viva_tensor

-module(viva_tensor@spec).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/spec.gleam").
-export([from_layout/1, tensor_spec/1, spec_from_parts/5, dtype_name/1, device_name/1, storage_name/1, memory_layout_name/1, spec_key/1]).
-export_type([tensor_spec/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 tensor_spec() :: {tensor_spec,
        list(integer()),
        viva_tensor@layout:tensor_dtype(),
        viva_tensor@layout:tensor_device(),
        viva_tensor@layout:tensor_storage(),
        viva_tensor@layout:tensor_memory_layout(),
        integer(),
        integer()}.

-file("src/viva_tensor/spec.gleam", 24).
?DOC(false).
-spec from_layout(viva_tensor@layout:tensor_layout()) -> tensor_spec().
from_layout(Metadata) ->
    {tensor_spec,
        erlang:element(5, Metadata),
        erlang:element(4, Metadata),
        erlang:element(3, Metadata),
        erlang:element(2, Metadata),
        case erlang:element(10, Metadata) of
            true ->
                row_major;

            false ->
                strided_layout
        end,
        erlang:element(9, Metadata),
        erlang:element(8, Metadata)}.

-file("src/viva_tensor/spec.gleam", 20).
?DOC(false).
-spec tensor_spec(viva_tensor@tensor:tensor()) -> tensor_spec().
tensor_spec(T) ->
    from_layout(viva_tensor@tensor:layout(T)).

-file("src/viva_tensor/spec.gleam", 39).
?DOC(false).
-spec spec_from_parts(
    list(integer()),
    viva_tensor@layout:tensor_dtype(),
    viva_tensor@layout:tensor_device(),
    viva_tensor@layout:tensor_storage(),
    viva_tensor@layout:tensor_memory_layout()
) -> tensor_spec().
spec_from_parts(Shape, Dtype, Device, Storage, Memory_layout) ->
    {tensor_spec,
        Shape,
        Dtype,
        Device,
        Storage,
        Memory_layout,
        erlang:length(Shape),
        gleam@list:fold(Shape, 1, fun(Acc, Dim) -> Acc * Dim end)}.

-file("src/viva_tensor/spec.gleam", 57).
?DOC(false).
-spec dtype_name(viva_tensor@layout:tensor_dtype()) -> binary().
dtype_name(Dtype) ->
    case Dtype of
        float64 ->
            <<"f64"/utf8>>;

        float32 ->
            <<"f32"/utf8>>;

        float16 ->
            <<"f16"/utf8>>;

        b_float16 ->
            <<"bf16"/utf8>>;

        float8_e4_m3 ->
            <<"fp8_e4m3"/utf8>>;

        int8 ->
            <<"int8"/utf8>>;

        int4 ->
            <<"int4"/utf8>>;

        sparse_float16 ->
            <<"sparse_f16"/utf8>>
    end.

-file("src/viva_tensor/spec.gleam", 70).
?DOC(false).
-spec device_name(viva_tensor@layout:tensor_device()) -> binary().
device_name(Device) ->
    case Device of
        beam_cpu ->
            <<"beam_cpu"/utf8>>;

        native_cpu ->
            <<"native_cpu"/utf8>>;

        {cuda_device, Index} ->
            <<"cuda:"/utf8, (erlang:integer_to_binary(Index))/binary>>
    end.

-file("src/viva_tensor/spec.gleam", 78).
?DOC(false).
-spec storage_name(viva_tensor@layout:tensor_storage()) -> binary().
storage_name(Storage) ->
    case Storage of
        dense_storage ->
            <<"dense"/utf8>>;

        strided_storage ->
            <<"strided"/utf8>>;

        native_storage ->
            <<"native"/utf8>>
    end.

-file("src/viva_tensor/spec.gleam", 86).
?DOC(false).
-spec memory_layout_name(viva_tensor@layout:tensor_memory_layout()) -> binary().
memory_layout_name(Memory_layout) ->
    case Memory_layout of
        row_major ->
            <<"row_major"/utf8>>;

        column_major ->
            <<"column_major"/utf8>>;

        strided_layout ->
            <<"strided"/utf8>>;

        packed_fp8_layout ->
            <<"packed_fp8"/utf8>>;

        packed_sparse24_layout ->
            <<"packed_sparse24"/utf8>>
    end.

-file("src/viva_tensor/spec.gleam", 108).
?DOC(false).
-spec shape_key(list(integer())) -> binary().
shape_key(Shape) ->
    case Shape of
        [] ->
            <<"scalar"/utf8>>;

        [Dim] ->
            erlang:integer_to_binary(Dim);

        [Dim@1 | Rest] ->
            <<<<(erlang:integer_to_binary(Dim@1))/binary, "x"/utf8>>/binary,
                (shape_key(Rest))/binary>>
    end.

-file("src/viva_tensor/spec.gleam", 96).
?DOC(false).
-spec spec_key(tensor_spec()) -> binary().
spec_key(Spec) ->
    <<<<<<<<<<<<<<<<(device_name(erlang:element(4, Spec)))/binary, "|"/utf8>>/binary,
                                (dtype_name(erlang:element(3, Spec)))/binary>>/binary,
                            "|"/utf8>>/binary,
                        (storage_name(erlang:element(5, Spec)))/binary>>/binary,
                    "|"/utf8>>/binary,
                (memory_layout_name(erlang:element(6, Spec)))/binary>>/binary,
            "|"/utf8>>/binary,
        (shape_key(erlang:element(2, Spec)))/binary>>.