module

ML::NN::GPUOps

GPU operations module

Constants

TILE_SIZE = 16

Instance methods

add(a : Tensor, b : Tensor, output : Tensor) : Nil

Elementwise add: out = a + b

Source
add_bias(input : Tensor, bias : Tensor, output : Tensor) : Nil

Add bias: output = input + bias (broadcast over batch)

Source
available?

Check if GPU is available for NN ops

Source
batched_matmul(a : Tensor, b : Tensor, c : Tensor) : Nil

Batched matrix multiply: C[b] = A[b] @ B[b] A: [batch, M, K], B: [batch, K, N], C: [batch, M, N]

Source
batched_matmul_tn(a : Tensor, b : Tensor, c : Tensor, scale : Float32 = 1.0_f32) : Nil

Batched matmul with transpose: C[b] = A[b] @ B[b]^T A: [batch, M, K], B: [batch, N, K], C: [batch, M, N]

Source
conv2d_forward(input : Tensor, weight : Tensor, bias : Tensor | Nil, output : Tensor, stride : Int32 = 1, padding : Int32 = 1, fuse_relu : Bool = false) : Nil

Conv2D forward: output = conv2d(input, weight) + bias input: [batch, H, W, in_channels] (NHWC format) weight: [out_channels, in_channels, kH, kW] (OIHW format) bias: [out_channels] or nil output: [batch, H_out, W_out, out_channels]

Source
ensure_initialized

Initialize GPU pipelines (lazy)

Source
fused_attention(q : Tensor, k : Tensor, v : Tensor, output : Tensor, scale : Float32) : Nil

Fused scaled dot-product attention: output = softmax(Q @ K^T / sqrt(d)) @ V Q, K, V: [batch_heads, seq_len, head_dim] Output: [batch_heads, seq_len, head_dim]

Source
fused_attention_tiled(q : Tensor, k : Tensor, v : Tensor, output : Tensor, scale : Float32) : Nil

Fused attention with tiled optimization (for larger sequences)

Source
fused_layernorm_linear(input : Tensor, ln_gamma : Tensor, ln_beta : Tensor, weight : Tensor, bias : Tensor | Nil, output : Tensor, eps : Float32 = 1e-5_f32) : Nil

Fused LayerNorm + Linear: output = Linear(LayerNorm(input))

Source
fused_layernorm_linear_gelu(input : Tensor, ln_gamma : Tensor, ln_beta : Tensor, weight : Tensor, bias : Tensor | Nil, output : Tensor, eps : Float32 = 1e-5_f32) : Nil

Fused LayerNorm + Linear + GELU

Source
get_pipeline(name : String) : Metal::ComputePipeline

Get or create pipeline for a kernel

Source
layernorm_forward(input : Tensor, gamma : Tensor, beta : Tensor, output : Tensor, eps : Float32 = 1e-5_f32) : Nil

LayerNorm forward input: [batch, features] gamma, beta: [features] output: [batch, features]

Source
linear_forward(input : Tensor, weight : Tensor, bias : Tensor | Nil, output : Tensor) : Nil

Linear forward: output = input @ weight^T + bias input: [batch, in_features] weight: [out_features, in_features] bias: [out_features] or nil output: [batch, out_features]

Source
linear_gelu_forward(input : Tensor, weight : Tensor, bias : Tensor | Nil, output : Tensor) : Nil

Linear forward with fused GELU activation

Source
nchw_to_nhwc(input : Tensor, output : Tensor) : Nil

Convert NCHW -> NHWC

Source
relu_forward(input : Tensor, output : Tensor)

ReLU forward (standalone)

Source
reshape_for_heads(input : Tensor, output : Tensor, batch : Int32, seq_len : Int32, num_heads : Int32, head_dim : Int32) : Nil

Reshape [batch, seq, embed] -> [batch * heads, seq, head_dim]

Source
reshape_from_heads(input : Tensor, output : Tensor, batch : Int32, seq_len : Int32, num_heads : Int32, head_dim : Int32) : Nil

Reshape [batch * heads, seq, head_dim] -> [batch, seq, embed]

Source
resize_normalize(input : Tensor, output : Tensor, out_h : Int32, out_w : Int32)

Fused resize + normalize for image preprocessing input: [h, w, 3] HWC in [0,1] output: [3, out_h, out_w] CHW normalized

Source
rmsnorm_forward(input : Tensor, gamma : Tensor, output : Tensor, eps : Float32 = 1e-5_f32) : Nil

RMSNorm forward

Source
rope_apply(input : Tensor, freqs : Tensor, output : Tensor, batch : Int32, seq_len : Int32, embed_dim : Int32, num_heads : Int32, head_dim : Int32, width : Int32) : Nil

Apply RoPE to [batch, seq, embed]

Source
softmax(input : Tensor, output : Tensor) : Nil

Softmax over last dimension input/output: [rows, cols]

Source
upsample_bilinear_2x(input : Tensor, output : Tensor)

Bilinear 2x upsampling (GPU, higher quality)

Source
upsample_nearest_2x(input : Tensor, output : Tensor)

Nearest neighbor 2x upsampling (GPU) input: [batch, h, w, channels] NHWC output: [batch, h2, w2, channels]

Source