Packages
Tensor library for Gleam/BEAM with a pure Gleam API, zero-copy views, and optional native acceleration
Current section
Files
Jump to
Current section
Files
src/viva_tensor@native@cuda.erl
-module(viva_tensor@native@cuda).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/native/cuda.gleam").
-export([new/2, to_list/1, shape/1, matmul/5, fp16_available/0, new16/2, to_list16/1, shape16/1, matmul16/5, gpu_workspace/0, to_rtx4090_fp32/1, to_rtx4090_fp16/1, workspace_zeros/2, workspace_from_tensor/2, workspace_backend/1, linear_layer/3, linear_layer_fp16/2, linear_layer_output_features/1, linear_output/3, accelerated_shape/1, linear_gelu_accelerated_into/4, linear_relu_accelerated_into/4, linear_relu_forward_into/3, linear_gelu_forward_into/3, linear_layer_backend/1, linear_layer_input_features/1, to_accelerated/1, matmul_auto/2, to_cpu_tensor/1, backend/1, matmul_accelerated/2, matmul_accelerated_into/3, matmul_relu_accelerated_into/3, matmul_gelu_accelerated_into/3, sync/0]).
-export_type([acceleration_backend/0, accelerated_tensor/0, gpu_workspace/0, linear_layer/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 acceleration_backend() :: rtx4090_fp16 |
rtx4090_fp32 |
mkl_native |
cpu_fallback.
-type accelerated_tensor() :: {cuda_fp16,
viva_tensor@core@ffi:cuda_tensor16_ref(),
list(integer()),
acceleration_backend()} |
{cuda_fp32,
viva_tensor@core@ffi:cuda_tensor_ref(),
list(integer()),
acceleration_backend()} |
{cpu, viva_tensor@tensor:tensor(), acceleration_backend()}.
-type gpu_workspace() :: {gpu_workspace, acceleration_backend()}.
-type linear_layer() :: {linear_layer,
accelerated_tensor(),
accelerated_tensor(),
integer(),
integer(),
acceleration_backend()}.
-file("src/viva_tensor/native/cuda.gleam", 25).
?DOC(false).
-spec new(list(float()), list(integer())) -> {ok,
viva_tensor@core@ffi:cuda_tensor_ref()} |
{error, binary()}.
new(Data, Shape) ->
viva_tensor_zig:ct_from_list(Data, Shape).
-file("src/viva_tensor/native/cuda.gleam", 30).
?DOC(false).
-spec to_list(viva_tensor@core@ffi:cuda_tensor_ref()) -> {ok, list(float())} |
{error, binary()}.
to_list(Tensor) ->
viva_tensor_zig:ct_to_list(Tensor).
-file("src/viva_tensor/native/cuda.gleam", 35).
?DOC(false).
-spec shape(viva_tensor@core@ffi:cuda_tensor_ref()) -> {ok, list(integer())} |
{error, binary()}.
shape(Tensor) ->
viva_tensor_zig:ct_shape(Tensor).
-file("src/viva_tensor/native/cuda.gleam", 41).
?DOC(false).
-spec matmul(
viva_tensor@core@ffi:cuda_tensor_ref(),
viva_tensor@core@ffi:cuda_tensor_ref(),
integer(),
integer(),
integer()
) -> {ok, viva_tensor@core@ffi:cuda_tensor_ref()} | {error, binary()}.
matmul(A, B, M, N, K) ->
viva_tensor_zig:ct_matmul(A, B, M, N, K).
-file("src/viva_tensor/native/cuda.gleam", 60).
?DOC(false).
-spec fp16_available() -> boolean().
fp16_available() ->
viva_tensor_zig:ct16_available().
-file("src/viva_tensor/native/cuda.gleam", 65).
?DOC(false).
-spec new16(list(float()), list(integer())) -> {ok,
viva_tensor@core@ffi:cuda_tensor16_ref()} |
{error, binary()}.
new16(Data, Shape) ->
viva_tensor_zig:ct16_from_list(Data, Shape).
-file("src/viva_tensor/native/cuda.gleam", 73).
?DOC(false).
-spec to_list16(viva_tensor@core@ffi:cuda_tensor16_ref()) -> {ok, list(float())} |
{error, binary()}.
to_list16(Tensor) ->
viva_tensor_zig:ct16_to_list(Tensor).
-file("src/viva_tensor/native/cuda.gleam", 78).
?DOC(false).
-spec shape16(viva_tensor@core@ffi:cuda_tensor16_ref()) -> {ok, list(integer())} |
{error, binary()}.
shape16(Tensor) ->
viva_tensor_zig:ct16_shape(Tensor).
-file("src/viva_tensor/native/cuda.gleam", 87).
?DOC(false).
-spec matmul16(
viva_tensor@core@ffi:cuda_tensor16_ref(),
viva_tensor@core@ffi:cuda_tensor16_ref(),
integer(),
integer(),
integer()
) -> {ok, viva_tensor@core@ffi:cuda_tensor_ref()} | {error, binary()}.
matmul16(A, B, M, N, K) ->
viva_tensor_zig:ct16_matmul(A, B, M, N, K).
-file("src/viva_tensor/native/cuda.gleam", 141).
?DOC(false).
-spec gpu_workspace() -> {ok, gpu_workspace()} |
{error, viva_tensor@core@error:tensor_error()}.
gpu_workspace() ->
case viva_tensor@core@ffi:cuda_available() andalso viva_tensor_zig:ct16_available(
) of
true ->
{ok, {gpu_workspace, rtx4090_fp16}};
false ->
{error,
{dimension_error, <<"CUDA FP16 backend is not available"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 665).
?DOC(false).
-spec to_mkl(viva_tensor@tensor:tensor()) -> accelerated_tensor().
to_mkl(T) ->
{cpu, T, mkl_native}.
-file("src/viva_tensor/native/cuda.gleam", 304).
?DOC(false).
-spec to_rtx4090_fp32(viva_tensor@tensor:tensor()) -> {ok, accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
to_rtx4090_fp32(T) ->
case viva_tensor@core@ffi:cuda_available() of
false ->
{error,
{dimension_error, <<"CUDA FP32 backend is not available"/utf8>>}};
true ->
_pipe = viva_tensor_zig:ct_from_list(
viva_tensor@tensor:to_list(T),
viva_tensor@tensor:shape(T)
),
_pipe@1 = gleam@result:map(
_pipe,
fun(Ref) ->
{cuda_fp32, Ref, viva_tensor@tensor:shape(T), rtx4090_fp32}
end
),
gleam@result:map_error(
_pipe@1,
fun(Reason) -> {dimension_error, Reason} end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 289).
?DOC(false).
-spec to_rtx4090_fp16(viva_tensor@tensor:tensor()) -> {ok, accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
to_rtx4090_fp16(T) ->
case viva_tensor@core@ffi:cuda_available() andalso viva_tensor_zig:ct16_available(
) of
false ->
{error,
{dimension_error, <<"CUDA FP16 backend is not available"/utf8>>}};
true ->
_pipe = viva_tensor_zig:ct16_from_list(
viva_tensor@tensor:to_list(T),
viva_tensor@tensor:shape(T)
),
_pipe@1 = gleam@result:map(
_pipe,
fun(Ref) ->
{cuda_fp16, Ref, viva_tensor@tensor:shape(T), rtx4090_fp16}
end
),
gleam@result:map_error(
_pipe@1,
fun(Reason) -> {dimension_error, Reason} end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 149).
?DOC(false).
-spec workspace_zeros(gpu_workspace(), list(integer())) -> {ok,
accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
workspace_zeros(Workspace, Shape) ->
case Workspace of
{gpu_workspace, rtx4090_fp16} ->
to_rtx4090_fp16(viva_tensor@tensor:zeros(Shape));
{gpu_workspace, rtx4090_fp32} ->
to_rtx4090_fp32(viva_tensor@tensor:zeros(Shape));
{gpu_workspace, mkl_native} ->
_pipe = viva_tensor@tensor:native_zeros(Shape),
gleam@result:map(_pipe, fun to_mkl/1);
{gpu_workspace, cpu_fallback} ->
{ok, {cpu, viva_tensor@tensor:zeros(Shape), cpu_fallback}}
end.
-file("src/viva_tensor/native/cuda.gleam", 163).
?DOC(false).
-spec workspace_from_tensor(gpu_workspace(), viva_tensor@tensor:tensor()) -> {ok,
accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
workspace_from_tensor(Workspace, T) ->
case Workspace of
{gpu_workspace, rtx4090_fp16} ->
to_rtx4090_fp16(T);
{gpu_workspace, rtx4090_fp32} ->
to_rtx4090_fp32(T);
{gpu_workspace, mkl_native} ->
Shape = viva_tensor@tensor:shape(T),
_pipe = viva_tensor@tensor:native_from_list(
viva_tensor@tensor:to_list(T),
Shape
),
gleam@result:map(_pipe, fun to_mkl/1);
{gpu_workspace, cpu_fallback} ->
{ok, {cpu, T, cpu_fallback}}
end.
-file("src/viva_tensor/native/cuda.gleam", 244).
?DOC(false).
-spec workspace_backend(gpu_workspace()) -> acceleration_backend().
workspace_backend(Workspace) ->
case Workspace of
{gpu_workspace, Backend} ->
Backend
end.
-file("src/viva_tensor/native/cuda.gleam", 189).
?DOC(false).
-spec linear_layer(
gpu_workspace(),
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor()
) -> {ok, linear_layer()} | {error, viva_tensor@core@error:tensor_error()}.
linear_layer(Workspace, Weight, Bias) ->
case {viva_tensor@tensor:shape(Weight), viva_tensor@tensor:shape(Bias)} of
{[Input_features, Output_features], [Bias_features]} when Output_features =:= Bias_features ->
gleam@result:'try'(
workspace_from_tensor(Workspace, Weight),
fun(Weight_acc) ->
gleam@result:'try'(
workspace_from_tensor(Workspace, Bias),
fun(Bias_acc) ->
{ok,
{linear_layer,
Weight_acc,
Bias_acc,
Input_features,
Output_features,
workspace_backend(Workspace)}}
end
)
end
);
{[_, Output_features@1], [Bias_features@1]} ->
{error, {shape_mismatch, [Output_features@1], [Bias_features@1]}};
{_, _} ->
{error,
{dimension_error,
<<"Expected weight matrix and bias vector"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 180).
?DOC(false).
-spec linear_layer_fp16(
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor()
) -> {ok, linear_layer()} | {error, viva_tensor@core@error:tensor_error()}.
linear_layer_fp16(Weight, Bias) ->
gleam@result:'try'(
gpu_workspace(),
fun(Workspace) -> linear_layer(Workspace, Weight, Bias) end
).
-file("src/viva_tensor/native/cuda.gleam", 265).
?DOC(false).
-spec linear_layer_output_features(linear_layer()) -> integer().
linear_layer_output_features(Layer) ->
case Layer of
{linear_layer, _, _, _, Output_features, _} ->
Output_features
end.
-file("src/viva_tensor/native/cuda.gleam", 217).
?DOC(false).
-spec linear_output(gpu_workspace(), linear_layer(), integer()) -> {ok,
accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
linear_output(Workspace, Layer, Batch_size) ->
workspace_zeros(
Workspace,
[Batch_size, linear_layer_output_features(Layer)]
).
-file("src/viva_tensor/native/cuda.gleam", 597).
?DOC(false).
-spec fused_linear_accelerated_into_checked(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
integer(),
integer(),
integer(),
binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_linear_accelerated_into_checked(Out, A, B, Bias, M, N, K, Activation) ->
case {Out, A, B, Bias} of
{{cuda_fp16, Out_ref, _, _},
{cuda_fp16, A_ref, _, _},
{cuda_fp16, B_ref, _, _},
{cuda_fp16, Bias_ref, _, _}} ->
_pipe = case Activation of
<<"relu"/utf8>> ->
viva_tensor_zig:ct16_linear_relu(
A_ref,
B_ref,
Bias_ref,
Out_ref,
M,
N,
K
);
<<"gelu"/utf8>> ->
viva_tensor_zig:ct16_linear_gelu(
A_ref,
B_ref,
Bias_ref,
Out_ref,
M,
N,
K
);
_ ->
{error, <<"unsupported_activation"/utf8>>}
end,
gleam@result:map_error(
_pipe,
fun(Reason) -> {dimension_error, Reason} end
);
{{cpu, Out_tensor, _},
{cpu, A_tensor, _},
{cpu, B_tensor, _},
{cpu, Bias_tensor, _}} ->
case Activation of
<<"relu"/utf8>> ->
viva_tensor@tensor:linear_relu_into(
Out_tensor,
A_tensor,
B_tensor,
Bias_tensor
);
_ ->
{error,
{dimension_error,
<<"CPU fused GELU is not implemented"/utf8>>}}
end;
{_, _, _, _} ->
{error,
{dimension_error,
<<"Fused linear activation requires matching backends"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 437).
?DOC(false).
-spec accelerated_shape(accelerated_tensor()) -> list(integer()).
accelerated_shape(T) ->
case T of
{cuda_fp16, _, Shape, _} ->
Shape;
{cuda_fp32, _, Shape@1, _} ->
Shape@1;
{cpu, Tensor, _} ->
erlang:element(3, Tensor)
end.
-file("src/viva_tensor/native/cuda.gleam", 563).
?DOC(false).
-spec fused_linear_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_linear_accelerated_into(Out, A, B, Bias, Activation) ->
case {accelerated_shape(Out),
accelerated_shape(A),
accelerated_shape(B),
accelerated_shape(Bias)} of
{[M_out, N_out], [M, K], [K2, N], [N_bias]} when (((K =:= K2) andalso (M_out =:= M)) andalso (N_out =:= N)) andalso (N_bias =:= N) ->
fused_linear_accelerated_into_checked(
Out,
A,
B,
Bias,
M,
N,
K,
Activation
);
{[M_out@1, N_out@1], [M@1, _], [_, N@1], [_]} ->
{error, {shape_mismatch, [M@1, N@1], [M_out@1, N_out@1]}};
{_, _, _, _} ->
{error,
{dimension_error,
<<"Expected matrices and a bias vector"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 401).
?DOC(false).
-spec linear_gelu_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_gelu_accelerated_into(Out, A, B, Bias) ->
fused_linear_accelerated_into(Out, A, B, Bias, <<"gelu"/utf8>>).
-file("src/viva_tensor/native/cuda.gleam", 391).
?DOC(false).
-spec linear_relu_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_relu_accelerated_into(Out, A, B, Bias) ->
fused_linear_accelerated_into(Out, A, B, Bias, <<"relu"/utf8>>).
-file("src/viva_tensor/native/cuda.gleam", 635).
?DOC(false).
-spec linear_forward_into(
accelerated_tensor(),
accelerated_tensor(),
linear_layer(),
binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_forward_into(Out, Input, Layer, Activation) ->
case Layer of
{linear_layer, Weight, Bias, Input_features, Output_features, _} ->
case {accelerated_shape(Input), accelerated_shape(Out)} of
{[Batch_size, Got_input], [Out_batch, Got_output]} when ((Got_input =:= Input_features) andalso (Out_batch =:= Batch_size)) andalso (Got_output =:= Output_features) ->
case Activation of
<<"relu"/utf8>> ->
linear_relu_accelerated_into(
Out,
Input,
Weight,
Bias
);
<<"gelu"/utf8>> ->
linear_gelu_accelerated_into(
Out,
Input,
Weight,
Bias
);
_ ->
{error,
{dimension_error,
<<"Unsupported activation"/utf8>>}}
end;
{[_, Got_input@1], [_, _]} ->
{error, {shape_mismatch, [Input_features], [Got_input@1]}};
{_, _} ->
{error,
{dimension_error,
<<"Expected input and output matrices"/utf8>>}}
end
end.
-file("src/viva_tensor/native/cuda.gleam", 226).
?DOC(false).
-spec linear_relu_forward_into(
accelerated_tensor(),
accelerated_tensor(),
linear_layer()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_relu_forward_into(Out, Input, Layer) ->
linear_forward_into(Out, Input, Layer, <<"relu"/utf8>>).
-file("src/viva_tensor/native/cuda.gleam", 235).
?DOC(false).
-spec linear_gelu_forward_into(
accelerated_tensor(),
accelerated_tensor(),
linear_layer()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
linear_gelu_forward_into(Out, Input, Layer) ->
linear_forward_into(Out, Input, Layer, <<"gelu"/utf8>>).
-file("src/viva_tensor/native/cuda.gleam", 251).
?DOC(false).
-spec linear_layer_backend(linear_layer()) -> acceleration_backend().
linear_layer_backend(Layer) ->
case Layer of
{linear_layer, _, _, _, _, Backend} ->
Backend
end.
-file("src/viva_tensor/native/cuda.gleam", 258).
?DOC(false).
-spec linear_layer_input_features(linear_layer()) -> integer().
linear_layer_input_features(Layer) ->
case Layer of
{linear_layer, _, _, Input_features, _, _} ->
Input_features
end.
-file("src/viva_tensor/native/cuda.gleam", 740).
?DOC(false).
-spec to_native(viva_tensor@tensor:tensor(), list(integer())) -> {ok,
viva_tensor@tensor:tensor()} |
{error, nil}.
to_native(T, Shape) ->
case viva_tensor@tensor:native_ref(T) of
{ok, _} ->
{ok, T};
{error, _} ->
_pipe = viva_tensor@tensor:native_from_list(
viva_tensor@tensor:to_list(T),
Shape
),
gleam@result:map_error(_pipe, fun(_) -> nil end)
end.
-file("src/viva_tensor/native/cuda.gleam", 749).
?DOC(false).
-spec to_mkl_or_cpu(viva_tensor@tensor:tensor(), list(integer())) -> accelerated_tensor().
to_mkl_or_cpu(T, Shape) ->
case to_native(T, Shape) of
{ok, Native} ->
{cpu, Native, mkl_native};
{error, _} ->
{cpu, T, cpu_fallback}
end.
-file("src/viva_tensor/native/cuda.gleam", 273).
?DOC(false).
-spec to_accelerated(viva_tensor@tensor:tensor()) -> {ok, accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
to_accelerated(T) ->
Shape = viva_tensor@tensor:shape(T),
case to_rtx4090_fp16(T) of
{ok, Gpu} ->
{ok, Gpu};
{error, _} ->
case to_rtx4090_fp32(T) of
{ok, Gpu@1} ->
{ok, Gpu@1};
{error, _} ->
{ok, to_mkl_or_cpu(T, Shape)}
end
end.
-file("src/viva_tensor/native/cuda.gleam", 723).
?DOC(false).
-spec matmul_mkl_then_cpu(
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
integer(),
integer(),
integer()
) -> {ok, accelerated_tensor()} | {error, viva_tensor@core@error:tensor_error()}.
matmul_mkl_then_cpu(A, B, M, N, K) ->
case {to_native(A, [M, K]), to_native(B, [K, N])} of
{{ok, A_native}, {ok, B_native}} ->
_pipe = viva_tensor@tensor:matmul(A_native, B_native),
gleam@result:map(_pipe, fun(Out) -> {cpu, Out, mkl_native} end);
{_, _} ->
_pipe@1 = viva_tensor@tensor:matmul(A, B),
gleam@result:map(
_pipe@1,
fun(Out@1) -> {cpu, Out@1, cpu_fallback} end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 696).
?DOC(false).
-spec try_rtx4090_fp32(
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
integer(),
integer(),
integer()
) -> {ok, accelerated_tensor()} | {error, nil}.
try_rtx4090_fp32(A, B, M, N, K) ->
case viva_tensor@core@ffi:cuda_available() of
false ->
{error, nil};
true ->
gleam@result:'try'(
begin
_pipe = viva_tensor_zig:ct_from_list(
viva_tensor@tensor:to_list(A),
[M, K]
),
gleam@result:map_error(_pipe, fun(_) -> nil end)
end,
fun(A_gpu) ->
gleam@result:'try'(
begin
_pipe@1 = viva_tensor_zig:ct_from_list(
viva_tensor@tensor:to_list(B),
[K, N]
),
gleam@result:map_error(_pipe@1, fun(_) -> nil end)
end,
fun(B_gpu) ->
_pipe@2 = viva_tensor_zig:ct_matmul(
A_gpu,
B_gpu,
M,
N,
K
),
_pipe@3 = gleam@result:map(
_pipe@2,
fun(Out) ->
{cuda_fp32, Out, [M, N], rtx4090_fp32}
end
),
gleam@result:map_error(_pipe@3, fun(_) -> nil end)
end
)
end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 669).
?DOC(false).
-spec try_rtx4090_fp16(
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
integer(),
integer(),
integer()
) -> {ok, accelerated_tensor()} | {error, nil}.
try_rtx4090_fp16(A, B, M, N, K) ->
case viva_tensor@core@ffi:cuda_available() andalso viva_tensor_zig:ct16_available(
) of
false ->
{error, nil};
true ->
gleam@result:'try'(
begin
_pipe = viva_tensor_zig:ct16_from_list(
viva_tensor@tensor:to_list(A),
[M, K]
),
gleam@result:map_error(_pipe, fun(_) -> nil end)
end,
fun(A_gpu) ->
gleam@result:'try'(
begin
_pipe@1 = viva_tensor_zig:ct16_from_list(
viva_tensor@tensor:to_list(B),
[K, N]
),
gleam@result:map_error(_pipe@1, fun(_) -> nil end)
end,
fun(B_gpu) ->
_pipe@2 = viva_tensor_zig:ct16_matmul(
A_gpu,
B_gpu,
M,
N,
K
),
_pipe@3 = gleam@result:map(
_pipe@2,
fun(Out) ->
{cuda_fp32, Out, [M, N], rtx4090_fp16}
end
),
gleam@result:map_error(_pipe@3, fun(_) -> nil end)
end
)
end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 319).
?DOC(false).
-spec matmul_auto(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> {ok,
accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
matmul_auto(A, B) ->
case {viva_tensor@tensor:shape(A), viva_tensor@tensor:shape(B)} of
{[M, K], [K2, N]} when K =:= K2 ->
case try_rtx4090_fp16(A, B, M, N, K) of
{ok, Result} ->
{ok, Result};
{error, _} ->
case try_rtx4090_fp32(A, B, M, N, K) of
{ok, Result@1} ->
{ok, Result@1};
{error, _} ->
matmul_mkl_then_cpu(A, B, M, N, K)
end
end;
{[_, K@1], [K2@1, _]} ->
{error, {shape_mismatch, [K@1, -1], [K2@1, -1]}};
{_, _} ->
{error, {dimension_error, <<"Expected two matrices"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 411).
?DOC(false).
-spec to_cpu_tensor(accelerated_tensor()) -> {ok, viva_tensor@tensor:tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
to_cpu_tensor(T) ->
case T of
{cpu, Tensor, _} ->
{ok, Tensor};
{cuda_fp16, Ref, Shape, _} ->
_pipe = viva_tensor_zig:ct16_to_list(Ref),
_pipe@1 = gleam@result:map(
_pipe,
fun(Data) -> {tensor, Data, Shape} end
),
gleam@result:map_error(
_pipe@1,
fun(Reason) -> {dimension_error, Reason} end
);
{cuda_fp32, Ref@1, Shape@1, _} ->
_pipe@2 = viva_tensor_zig:ct_to_list(Ref@1),
_pipe@3 = gleam@result:map(
_pipe@2,
fun(Data@1) -> {tensor, Data@1, Shape@1} end
),
gleam@result:map_error(
_pipe@3,
fun(Reason@1) -> {dimension_error, Reason@1} end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 428).
?DOC(false).
-spec backend(accelerated_tensor()) -> acceleration_backend().
backend(T) ->
case T of
{cuda_fp16, _, _, Backend} ->
Backend;
{cuda_fp32, _, _, Backend@1} ->
Backend@1;
{cpu, _, Backend@2} ->
Backend@2
end.
-file("src/viva_tensor/native/cuda.gleam", 451).
?DOC(false).
-spec matmul_accelerated_checked(
accelerated_tensor(),
accelerated_tensor(),
integer(),
integer(),
integer()
) -> {ok, accelerated_tensor()} | {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated_checked(A, B, M, N, K) ->
case {A, B} of
{{cuda_fp16, A_ref, _, _}, {cuda_fp16, B_ref, _, _}} ->
_pipe = viva_tensor_zig:ct16_matmul(A_ref, B_ref, M, N, K),
_pipe@1 = gleam@result:map(
_pipe,
fun(Out) -> {cuda_fp32, Out, [M, N], rtx4090_fp16} end
),
gleam@result:map_error(
_pipe@1,
fun(Reason) -> {dimension_error, Reason} end
);
{{cuda_fp32, A_ref@1, _, _}, {cuda_fp32, B_ref@1, _, _}} ->
_pipe@2 = viva_tensor_zig:ct_matmul(A_ref@1, B_ref@1, M, N, K),
_pipe@3 = gleam@result:map(
_pipe@2,
fun(Out@1) -> {cuda_fp32, Out@1, [M, N], rtx4090_fp32} end
),
gleam@result:map_error(
_pipe@3,
fun(Reason@1) -> {dimension_error, Reason@1} end
);
{{cpu, A_tensor, _}, {cpu, B_tensor, _}} ->
_pipe@4 = viva_tensor@tensor:matmul(A_tensor, B_tensor),
gleam@result:map(
_pipe@4,
fun(Out@2) -> {cpu, Out@2, backend(A)} end
);
{_, _} ->
gleam@result:'try'(
to_cpu_tensor(A),
fun(A_cpu) ->
gleam@result:'try'(
to_cpu_tensor(B),
fun(B_cpu) -> matmul_auto(A_cpu, B_cpu) end
)
end
)
end.
-file("src/viva_tensor/native/cuda.gleam", 344).
?DOC(false).
-spec matmul_accelerated(accelerated_tensor(), accelerated_tensor()) -> {ok,
accelerated_tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated(A, B) ->
case {accelerated_shape(A), accelerated_shape(B)} of
{[M, K], [K2, N]} when K =:= K2 ->
matmul_accelerated_checked(A, B, M, N, K);
{[_, K@1], [K2@1, _]} ->
{error, {shape_mismatch, [K@1, -1], [K2@1, -1]}};
{_, _} ->
{error, {dimension_error, <<"Expected two matrices"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 515).
?DOC(false).
-spec backend_label(accelerated_tensor()) -> binary().
backend_label(T) ->
case T of
{cuda_fp16, _, _, _} ->
<<"cuda/fp16"/utf8>>;
{cuda_fp32, _, _, _} ->
<<"cuda/fp32"/utf8>>;
{cpu, _, _} ->
<<"cpu"/utf8>>
end.
-file("src/viva_tensor/native/cuda.gleam", 485).
?DOC(false).
-spec matmul_accelerated_into_checked(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
integer(),
integer(),
integer()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated_into_checked(Out, A, B, M, N, K) ->
case {Out, A, B} of
{{cuda_fp16, Out_ref, _, _},
{cuda_fp16, A_ref, _, _},
{cuda_fp16, B_ref, _, _}} ->
_pipe = viva_tensor_zig:ct16_matmul_inplace(
A_ref,
B_ref,
Out_ref,
M,
N,
K
),
gleam@result:map_error(
_pipe,
fun(Reason) -> {dimension_error, Reason} end
);
{{cuda_fp32, Out_ref@1, _, _},
{cuda_fp32, A_ref@1, _, _},
{cuda_fp32, B_ref@1, _, _}} ->
_pipe@1 = viva_tensor_zig:ct_matmul_inplace(
A_ref@1,
B_ref@1,
Out_ref@1,
M,
N,
K
),
gleam@result:map_error(
_pipe@1,
fun(Reason@1) -> {dimension_error, Reason@1} end
);
{{cpu, Out_tensor, _}, {cpu, A_tensor, _}, {cpu, B_tensor, _}} ->
viva_tensor@tensor:matmul_into(Out_tensor, A_tensor, B_tensor);
{_, _, _} ->
{error,
{backend_mismatch,
<<"matmul_accelerated_into"/utf8>>,
backend_label(Out),
backend_label(A),
backend_label(B)}}
end.
-file("src/viva_tensor/native/cuda.gleam", 358).
?DOC(false).
-spec matmul_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_accelerated_into(Out, A, B) ->
case {accelerated_shape(Out), accelerated_shape(A), accelerated_shape(B)} of
{[M_out, N_out], [M, K], [K2, N]} when ((K =:= K2) andalso (M_out =:= M)) andalso (N_out =:= N) ->
matmul_accelerated_into_checked(Out, A, B, M, N, K);
{[M_out@1, N_out@1], [M@1, _], [_, N@1]} ->
{error, {shape_mismatch, [M@1, N@1], [M_out@1, N_out@1]}};
{_, _, _} ->
{error, {dimension_error, <<"Expected matrices"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 538).
?DOC(false).
-spec fused_activation_accelerated_into_checked(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
integer(),
integer(),
integer(),
binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_activation_accelerated_into_checked(Out, A, B, M, N, K, Activation) ->
case {Out, A, B} of
{{cuda_fp16, Out_ref, _, _},
{cuda_fp16, A_ref, _, _},
{cuda_fp16, B_ref, _, _}} ->
_pipe = case Activation of
<<"relu"/utf8>> ->
viva_tensor_zig:ct16_matmul_fused_relu(
A_ref,
B_ref,
Out_ref,
M,
N,
K
);
<<"gelu"/utf8>> ->
viva_tensor_zig:ct16_matmul_fused_gelu(
A_ref,
B_ref,
Out_ref,
M,
N,
K
);
_ ->
{error, <<"unsupported_activation"/utf8>>}
end,
gleam@result:map_error(
_pipe,
fun(Reason) -> {dimension_error, Reason} end
);
{_, _, _} ->
{error,
{dimension_error,
<<"Fused CUDA activation requires FP16 accelerated tensors"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 523).
?DOC(false).
-spec fused_activation_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor(),
binary()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
fused_activation_accelerated_into(Out, A, B, Activation) ->
case {accelerated_shape(Out), accelerated_shape(A), accelerated_shape(B)} of
{[M_out, N_out], [M, K], [K2, N]} when ((K =:= K2) andalso (M_out =:= M)) andalso (N_out =:= N) ->
fused_activation_accelerated_into_checked(
Out,
A,
B,
M,
N,
K,
Activation
);
{[M_out@1, N_out@1], [M@1, _], [_, N@1]} ->
{error, {shape_mismatch, [M@1, N@1], [M_out@1, N_out@1]}};
{_, _, _} ->
{error, {dimension_error, <<"Expected matrices"/utf8>>}}
end.
-file("src/viva_tensor/native/cuda.gleam", 373).
?DOC(false).
-spec matmul_relu_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_relu_accelerated_into(Out, A, B) ->
fused_activation_accelerated_into(Out, A, B, <<"relu"/utf8>>).
-file("src/viva_tensor/native/cuda.gleam", 382).
?DOC(false).
-spec matmul_gelu_accelerated_into(
accelerated_tensor(),
accelerated_tensor(),
accelerated_tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
matmul_gelu_accelerated_into(Out, A, B) ->
fused_activation_accelerated_into(Out, A, B, <<"gelu"/utf8>>).
-file("src/viva_tensor/native/cuda.gleam", 446).
?DOC(false).
-spec sync() -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
sync() ->
_pipe = viva_tensor_zig:cuda_sync(),
gleam@result:map_error(_pipe, fun(Reason) -> {dimension_error, Reason} end).