webgpu

package
v0.9.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: May 27, 2026 License: Apache-2.0 Imports: 16 Imported by: 0

Documentation

Rendered for windows/amd64

Overview

Package webgpu implements the WebGPU backend for GPU-accelerated tensor operations. Uses gogpu/wgpu (github.com/gogpu/wgpu) for pure Go, zero-CGO WebGPU bindings.

Package webgpu implements the WebGPU backend for GPU-accelerated tensor operations.

Package webgpu implements the WebGPU backend for GPU-accelerated tensor operations.

Package webgpu implements the WebGPU backend for GPU-accelerated tensor operations.

Package webgpu implements the WebGPU backend for GPU-accelerated tensor operations.

Package webgpu provides embedded WGSL compute shaders for tensor operations.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func IsAvailable 

func IsAvailable() bool

IsAvailable checks if WebGPU with compute shader support is available. Returns false on software renderers that don't support compute pipelines, and also returns false if the underlying driver panics (e.g., missing GPU on CI runners). The panic recovery prevents process crashes on headless systems such as GitHub Actions Windows runners that have no Vulkan driver.

func ListAdapters 

func ListAdapters() ([]*wgpu.AdapterInfo, error)

ListAdapters returns information about all available GPU adapters.

Types

type Backend 

type Backend struct {

	// Lazy mode: when true, operations return lazy tensors that keep data on GPU
	// until Data() is explicitly called. This is the key optimization for
	// Phase 3 Integration - eliminates readBuffer() bottleneck.
	// Default: true for optimal performance.
	LazyMode bool
	// contains filtered or unexported fields
}

Backend implements tensor operations on GPU using WebGPU.

func New 

func New() (*Backend, error)

New creates a new WebGPU backend.

Backend selection via GOGPU_GRAPHICS_API environment variable (gogpu ecosystem standard): 

GOGPU_GRAPHICS_API=auto      (default) wgpu selects best available
GOGPU_GRAPHICS_API=vulkan    force Vulkan
GOGPU_GRAPHICS_API=dx12      force DirectX 12
GOGPU_GRAPHICS_API=metal     force Metal
GOGPU_GRAPHICS_API=gl        force OpenGL/ES
GOGPU_GRAPHICS_API=software  software compute — no GPU required (CI, testing)

func (*Backend) Abs added in v0.8.0 

func (b *Backend) Abs(x *tensor.RawTensor) *tensor.RawTensor

Abs computes element-wise absolute value on GPU.

Only supports float32 dtype for now. Will raise a panic if called with unsupported dtype.

func (*Backend) AbsGPU added in v0.8.0 

func (b *Backend) AbsGPU(t *GPUTensor) *GPUTensor

AbsGPU applies absolute value activation on GPU. Data stays on GPU - no CPU transfer occurs.

func (*Backend) AdapterInfo 

func (b *Backend) AdapterInfo() *wgpu.AdapterInfo

AdapterInfo returns information about the GPU adapter.

func (*Backend) Add 

func (b *Backend) Add(a, other *tensor.RawTensor) *tensor.RawTensor

Add performs element-wise addition on GPU. Supports float32 and int32 dtypes. In LazyMode (default), returns a lazy tensor that keeps data on GPU.

func (*Backend) AddBackwardGPU added in v0.6.0 

func (b *Backend) AddBackwardGPU(_, _, grad *GPUTensor) (*GPUTensor, *GPUTensor)

AddBackwardGPU computes gradients for element-wise addition. d(a+b)/da = 1, d(a+b)/db = 1.

func (*Backend) AddGPU added in v0.6.0 

func (b *Backend) AddGPU(a, c *GPUTensor) *GPUTensor

AddGPU performs element-wise addition on GPU tensors. Data stays on GPU - no CPU transfer occurs.

func (*Backend) AddScalar added in v0.3.0 

func (b *Backend) AddScalar(x *tensor.RawTensor, scalar any) *tensor.RawTensor

AddScalar adds a scalar to tensor elements on GPU.

func (*Backend) And added in v0.3.0 

func (b *Backend) And(a, other *tensor.RawTensor) *tensor.RawTensor

And performs element-wise logical AND on GPU. Supports mixed dtypes by casting to float32 (for boolean tensors from different sources).

func (*Backend) Argmax added in v0.3.0 

func (b *Backend) Argmax(x *tensor.RawTensor, dim int) *tensor.RawTensor

Argmax returns indices of maximum values along dimension on GPU.

func (*Backend) BatchMatMul added in v0.4.0 

func (b *Backend) BatchMatMul(a, other *tensor.RawTensor) *tensor.RawTensor

BatchMatMul performs batched matrix multiplication on GPU. Supports 3D tensors [batch, M, K] @ [batch, K, N] -> [batch, M, N] and 4D tensors [batch, heads, M, K] @ [batch, heads, K, N].

func (*Backend) Cast added in v0.3.0 

func (b *Backend) Cast(x *tensor.RawTensor, dtype tensor.DataType) *tensor.RawTensor

Cast converts tensor to different data type. Supports float32 and int32 as target types.

func (*Backend) Cat added in v0.3.0 

func (b *Backend) Cat(tensors []*tensor.RawTensor, dim int) *tensor.RawTensor

Cat concatenates tensors along the specified dimension.

In LazyMode with float32 inputs, dispatches GPU compute passes (catShader) to copy each input into the correct region of a pre-allocated output buffer — zero CPU allocation. Falls back to CPU concatenation otherwise.

func (*Backend) Chunk added in v0.3.0 

func (b *Backend) Chunk(x *tensor.RawTensor, n, dim int) []*tensor.RawTensor

Chunk splits tensor into n equal parts along the specified dimension.

In LazyMode with float32 inputs, dispatches GPU compute passes (chunkShader) to copy each slice into a freshly-allocated output buffer — zero CPU allocation. Falls back to CPU splitting otherwise.

func (*Backend) Clamp added in v0.8.1 

func (b *Backend) Clamp(x *tensor.RawTensor, minBound, maxBound any) *tensor.RawTensor

Clamp restricts tensor values element-wise to [minBound, maxBound].

Supports float32 and int32 dtypes. Will raise a panic if called with unsupported dtype.

func (*Backend) ClampGPU added in v0.8.1 

func (b *Backend) ClampGPU(t *GPUTensor, minValue, maxValue any) *GPUTensor

ClampGPU applies clamp activation on GPU: clamp(x, minValue, maxValue). Data stays on GPU - no CPU transfer occurs.

func (*Backend) ClearInputBufferCache added in v0.9.1 

func (b *Backend) ClearInputBufferCache()

ClearInputBufferCache is the exported version of clearInputBufferCache. Call this between training steps if weight tensors are replaced by new *RawTensor objects (e.g. after in-place optimizer updates that produce new tensors). Forgetting to call this after weight replacement causes stale GPU data to be used in forward passes.

func (*Backend) Conv2D 

func (b *Backend) Conv2D(input, kernel *tensor.RawTensor, stride, padding int) *tensor.RawTensor

Conv2D performs 2D convolution on GPU. Input shape: [batch, in_channels, height, width]. Kernel shape: [out_channels, in_channels, kH, kW].

func (*Backend) Conv2DInputBackward added in v0.7.1 

func (b *Backend) Conv2DInputBackward(input, kernel, grad *tensor.RawTensor, stride, padding int) *tensor.RawTensor

Conv2DInputBackward computes gradient with respect to input for Conv2D. Not yet implemented for WebGPU backend.

func (*Backend) Conv2DKernelBackward added in v0.7.1 

func (b *Backend) Conv2DKernelBackward(input, kernel, grad *tensor.RawTensor, stride, padding int) *tensor.RawTensor

Conv2DKernelBackward computes gradient with respect to kernel for Conv2D. Not yet implemented for WebGPU backend.

func (*Backend) Cos added in v0.3.0 

func (b *Backend) Cos(x *tensor.RawTensor) *tensor.RawTensor

Cos computes element-wise cosine on GPU.

func (*Backend) DeferReleaseGPUBuffer added in v0.9.1 

func (b *Backend) DeferReleaseGPUBuffer(bufferPtr unsafe.Pointer)

DeferReleaseGPUBuffer implements tensor.LazyBackend interface. Returns the buffer to pool for reuse. If no encoder is active (no pending command buffers reference the buffer), returns immediately. Otherwise queues for return after the next flushCommands/Submit.

func (*Backend) Device 

func (b *Backend) Device() tensor.Device

Device returns the compute device.

func (*Backend) Div 

func (b *Backend) Div(a, other *tensor.RawTensor) *tensor.RawTensor

Div performs element-wise division on GPU. Supports float32 and int32 dtypes. In LazyMode (default), returns a lazy tensor that keeps data on GPU.

func (*Backend) DivBackwardGPU added in v0.6.0 

func (b *Backend) DivBackwardGPU(a, c, grad *GPUTensor) (*GPUTensor, *GPUTensor)

DivBackwardGPU computes gradients for element-wise division. d(a/b)/da = 1/b, d(a/b)/db = -a/b^2.

func (*Backend) DivGPU added in v0.6.0 

func (b *Backend) DivGPU(a, c *GPUTensor) *GPUTensor

DivGPU performs element-wise division on GPU tensors. Data stays on GPU - no CPU transfer occurs.

func (*Backend) DivScalar added in v0.3.0 

func (b *Backend) DivScalar(x *tensor.RawTensor, scalar any) *tensor.RawTensor

DivScalar divides tensor elements by a scalar on GPU.

func (*Backend) Embedding added in v0.5.1 

func (b *Backend) Embedding(weight, indices *tensor.RawTensor) *tensor.RawTensor

Embedding performs embedding lookup on GPU. weight: [num_embeddings, embedding_dim], indices: int32 tensor. Returns: [...indices_shape, embedding_dim].

func (*Backend) Equal added in v0.3.0 

func (b *Backend) Equal(a, other *tensor.RawTensor) *tensor.RawTensor

Equal performs element-wise equality comparison on GPU. Always returns float32 tensor (0.0 for false, 1.0 for true).

func (*Backend) Erf added in v0.7.15 

func (b *Backend) Erf(x *tensor.RawTensor) *tensor.RawTensor

Erf computes element-wise error function on GPU.

func (*Backend) ErfGPU added in v0.7.15 

func (b *Backend) ErfGPU(t *GPUTensor) *GPUTensor

ErfGPU applies erf activation on GPU. Data stays on GPU - no CPU transfer occurs.

func (*Backend) Exp added in v0.3.0 

func (b *Backend) Exp(x *tensor.RawTensor) *tensor.RawTensor

Exp computes element-wise exponential on GPU.

func (*Backend) Expand added in v0.3.0 

func (b *Backend) Expand(x *tensor.RawTensor, newShape tensor.Shape) *tensor.RawTensor

Expand broadcasts tensor to new shape. GPU-accelerated for up to 6D tensors.

func (*Backend) FlashAttentionGPU added in v0.7.0 

func (b *Backend) FlashAttentionGPU(
	q, k, v *tensor.RawTensor,
	scale float32,
	causal bool,
	blockSize int,
) (*tensor.RawTensor, error)

FlashAttentionGPU executes Flash Attention 2 on GPU using WebGPU.

This implementation uses tiled computation with online softmax to achieve O(N) memory complexity instead of O(N²) for standard attention.

Parameters:

  • q: Query tensor [batch, seqLen, numHeads, headDim]
  • k: Key tensor [batch, kvLen, numHeads, headDim]
  • v: Value tensor [batch, kvLen, numHeads, headDim]
  • scale: Attention scale factor (typically 1/sqrt(headDim))
  • causal: Whether to apply causal masking
  • blockSize: Tile size for blocked computation (64 or 128)

Returns:

  • *tensor.RawTensor: Output tensor [batch, seqLen, numHeads, headDim]

func (*Backend) FlushGPU added in v0.9.1 

func (b *Backend) FlushGPU()

FlushGPU submits all pending GPU commands and returns deferred release buffers to pool. Unlike ReclaimMemory, does NOT release live tensors — safe to call between training steps when carry state is alive.

func (*Backend) FromRawTensor added in v0.6.0 

func (b *Backend) FromRawTensor(t *tensor.RawTensor) *GPUTensor

FromRawTensor uploads a CPU tensor to GPU memory. This creates a new GPUTensor with data copied from the RawTensor.

func (*Backend) Gather added in v0.3.0 

func (b *Backend) Gather(input *tensor.RawTensor, dim int, indices *tensor.RawTensor) *tensor.RawTensor

Gather selects elements along dim using index tensor on GPU.

func (*Backend) Greater added in v0.3.0 

func (b *Backend) Greater(a, other *tensor.RawTensor) *tensor.RawTensor

Greater performs element-wise greater-than comparison on GPU. Always returns float32 tensor (0.0 for false, 1.0 for true).

func (*Backend) GreaterEqual added in v0.3.0 

func (b *Backend) GreaterEqual(a, other *tensor.RawTensor) *tensor.RawTensor

GreaterEqual performs element-wise greater-or-equal comparison on GPU. Always returns float32 tensor (0.0 for false, 1.0 for true).

func (*Backend) Log added in v0.3.0 

func (b *Backend) Log(x *tensor.RawTensor) *tensor.RawTensor

Log computes natural logarithm element-wise on GPU.

func (*Backend) Lower added in v0.3.0 

func (b *Backend) Lower(a, other *tensor.RawTensor) *tensor.RawTensor

Lower performs element-wise less-than comparison on GPU. Always returns float32 tensor (0.0 for false, 1.0 for true).

func (*Backend) LowerEqual added in v0.3.0 

func (b *Backend) LowerEqual(a, other *tensor.RawTensor) *tensor.RawTensor

LowerEqual performs element-wise less-or-equal comparison on GPU. Always returns float32 tensor (0.0 for false, 1.0 for true).

func (*Backend) MatMul 

func (b *Backend) MatMul(a, other *tensor.RawTensor) *tensor.RawTensor

MatMul performs matrix multiplication on GPU.

func (*Backend) MatMulBackwardGPU added in v0.6.0 

func (b *Backend) MatMulBackwardGPU(a, c, grad *GPUTensor) (*GPUTensor, *GPUTensor)

MatMulBackwardGPU computes gradients for matrix multiplication. d(A@B)/dA = grad@B^T, d(A@B)/dB = A^T@grad.

func (*Backend) MatMulGPU added in v0.6.0 

func (b *Backend) MatMulGPU(a, c *GPUTensor) *GPUTensor

MatMulGPU performs matrix multiplication on GPU tensors. Data stays on GPU - no CPU transfer occurs.

func (*Backend) MaxPool2D 

func (b *Backend) MaxPool2D(input *tensor.RawTensor, kernelSize, stride int) *tensor.RawTensor

MaxPool2D performs 2D max pooling on GPU. Input shape: [batch, channels, height, width].

func (*Backend) MaxPool2DBackward added in v0.7.1 

func (b *Backend) MaxPool2DBackward(input, grad *tensor.RawTensor, maxIndices []int, kernelSize, stride int) *tensor.RawTensor

MaxPool2DBackward computes gradient with respect to input for MaxPool2D. Not yet implemented for WebGPU backend.

func (*Backend) MeanDim added in v0.3.0 

func (b *Backend) MeanDim(x *tensor.RawTensor, dim int, keepDim bool) *tensor.RawTensor

MeanDim computes mean along a dimension.

In LazyMode, composes the lazy SumDim GPU path with MulScalar(1/dimSize), keeping all intermediate results on GPU — zero CPU allocation. Falls back to CPU arithmetic when LazyMode is disabled.

func (*Backend) MemoryStats 

func (b *Backend) MemoryStats() MemoryStats

MemoryStats returns current GPU memory usage statistics.

func (*Backend) Mul 

func (b *Backend) Mul(a, other *tensor.RawTensor) *tensor.RawTensor

Mul performs element-wise multiplication on GPU. Supports float32 and int32 dtypes. In LazyMode (default), returns a lazy tensor that keeps data on GPU.

func (*Backend) MulBackwardGPU added in v0.6.0 

func (b *Backend) MulBackwardGPU(a, c, grad *GPUTensor) (*GPUTensor, *GPUTensor)

MulBackwardGPU computes gradients for element-wise multiplication. d(a*b)/da = b, d(a*b)/db = a.

func (*Backend) MulGPU added in v0.6.0 

func (b *Backend) MulGPU(a, c *GPUTensor) *GPUTensor

MulGPU performs element-wise multiplication on GPU tensors. Data stays on GPU - no CPU transfer occurs.

func (*Backend) MulScalar added in v0.3.0 

func (b *Backend) MulScalar(x *tensor.RawTensor, scalar any) *tensor.RawTensor

MulScalar multiplies tensor elements by a scalar on GPU.

func (*Backend) Name 

func (b *Backend) Name() string

Name returns the backend name.

func (*Backend) NewBatch added in v0.6.0 

func (b *Backend) NewBatch() *CommandBatch

NewBatch creates a new command batch for accumulating operations. The batch will use a single CommandEncoder for all operations.

func (*Backend) Not added in v0.3.0 

func (b *Backend) Not(x *tensor.RawTensor) *tensor.RawTensor

Not performs element-wise logical NOT on GPU.

func (*Backend) NotEqual added in v0.3.0 

func (b *Backend) NotEqual(a, other *tensor.RawTensor) *tensor.RawTensor

NotEqual performs element-wise inequality comparison on GPU. Always returns float32 tensor (0.0 for false, 1.0 for true).

func (*Backend) OnesGPU added in v0.6.0 

func (b *Backend) OnesGPU(shape tensor.Shape, dtype tensor.DataType) *GPUTensor

OnesGPU creates a GPU tensor filled with ones. Data is initialized to ones on CPU then uploaded to GPU.

func (*Backend) Or added in v0.3.0 

func (b *Backend) Or(a, other *tensor.RawTensor) *tensor.RawTensor

Or performs element-wise logical OR on GPU. Supports mixed dtypes by casting to float32 (for boolean tensors from different sources).

func (*Backend) RandGPU added in v0.6.0 

func (b *Backend) RandGPU(shape tensor.Shape, dtype tensor.DataType) *GPUTensor

RandGPU creates a random GPU tensor with uniform distribution [0, 1). Data is generated on CPU using math/rand then uploaded to GPU.

func (*Backend) ReLU 

func (b *Backend) ReLU(x *tensor.RawTensor) *tensor.RawTensor

ReLU applies ReLU activation: max(0, x).

func (*Backend) ReLUBackwardGPU added in v0.6.0 

func (b *Backend) ReLUBackwardGPU(input, grad *GPUTensor) *GPUTensor

ReLUBackwardGPU computes gradients for ReLU activation. d(ReLU(x))/dx = 1 if x > 0, else 0. grad_input = grad * (input > 0).

func (*Backend) ReLUGPU added in v0.6.0 

func (b *Backend) ReLUGPU(t *GPUTensor) *GPUTensor

ReLUGPU applies ReLU activation on GPU: max(0, x). Data stays on GPU - no CPU transfer occurs.

func (*Backend) ReadGPUBuffer added in v0.6.0 

func (b *Backend) ReadGPUBuffer(bufferPtr unsafe.Pointer, size uint64) ([]byte, error)

ReadGPUBuffer implements tensor.LazyBackend interface. Reads data from a GPU Storage buffer (Storage | CopySrc) to CPU memory. bufferPtr must point to a *wgpu.Buffer created with BufferUsageStorage|CopySrc.

Deferred staging: a temporary MapRead buffer is created here on demand, the result buffer is copied into it, and the staging buffer is released immediately after readback. This means NO staging buffers exist during op execution — only one result buffer per lazy tensor.

Sequence:

  1. flushCommands() — submit all pending compute command buffers in one batch.
  2. Poll(PollWait) — block until the GPU completes all submitted commands.
  3. Create a transient MapRead staging buffer.
  4. CopyBufferToBuffer(resultBuf → staging), Submit, Poll.
  5. Map staging, copy bytes to CPU slice, Unmap, release staging.

func (*Backend) ReclaimMemory added in v0.9.1 

func (b *Backend) ReclaimMemory()

ReclaimMemory implements tensor.MemoryReclaimer. Releases all non-persistent live GPU tensors tracked by the backend, then flushes pending commands and blocks until the GPU completes them. Persistent tensors (optimizer moments, model weights) survive — only transient intermediates (forward pass, NoGrad blocks, masks) are released.

func (*Backend) RegisterLiveGPU added in v0.9.1 

func (b *Backend) RegisterLiveGPU(l *tensor.LazyGPUData)

RegisterLiveGPU registers a LazyGPUData in the backend's live tensor set.

func (*Backend) Release 

func (b *Backend) Release()

Release releases all WebGPU resources. Must be called when the backend is no longer needed.

func (*Backend) ReleaseGPUBuffer added in v0.6.0 

func (b *Backend) ReleaseGPUBuffer(bufferPtr unsafe.Pointer)

ReleaseGPUBuffer implements tensor.LazyBackend interface. Returns the buffer to the GPU pool for reuse (ADR-016 ExclusivePool).

func (*Backend) Reshape 

func (b *Backend) Reshape(t *tensor.RawTensor, newShape tensor.Shape) *tensor.RawTensor

Reshape returns a tensor with new shape backed by the same data.

In LazyMode, when the source tensor has unrealized GPU data, performs a GPU-to-GPU buffer copy (zero CPU allocation) to create the reshaped result. Falls back to CPU copy otherwise.

func (*Backend) Rsqrt added in v0.3.0 

func (b *Backend) Rsqrt(x *tensor.RawTensor) *tensor.RawTensor

Rsqrt computes element-wise reciprocal square root on GPU.

func (*Backend) ScatterAdd added in v0.8.2 

func (b *Backend) ScatterAdd(dest *tensor.RawTensor, dim int, indices, src *tensor.RawTensor) *tensor.RawTensor

ScatterAdd performs a general scatter-add matching Gather backward semantics.

For each element in src (same shape as indices), accumulates into result along dim at the position given by the corresponding index value. Follows Burn's float_scatter_add.

In LazyMode, dispatches a GPU compute shader that keeps the result on GPU using a per-destination-element approach (no f32 atomics required). Falls back to CPU for non-lazy mode.

Returns a new tensor with the same shape as dest. dest is not modified.

func (*Backend) SelectAdd added in v0.8.2 

func (b *Backend) SelectAdd(dest *tensor.RawTensor, dim int, indices, src *tensor.RawTensor) *tensor.RawTensor

SelectAdd performs a scatter-add along the specified dimension.

Used primarily in Embedding backward to accumulate gradient rows into the weight gradient tensor. In LazyMode, dispatches a GPU compute shader that keeps the result on GPU without requiring f32 atomics (per-destination-row approach). Falls back to CPU for non-lazy mode.

func (*Backend) SetLazyMode added in v0.6.0 

func (b *Backend) SetLazyMode(enabled bool)

SetLazyMode enables or disables lazy evaluation mode. When enabled (default), operations return lazy tensors that keep data on GPU until Data() is explicitly called. This dramatically improves performance by eliminating unnecessary GPU→CPU transfers. When disabled, operations immediately transfer results to CPU (slower).

func (*Backend) SiLU added in v0.8.1 

func (b *Backend) SiLU(x *tensor.RawTensor) *tensor.RawTensor

SiLU applies SiLU (Swish) activation: x * sigmoid(x).

func (*Backend) Sigmoid 

func (b *Backend) Sigmoid(x *tensor.RawTensor) *tensor.RawTensor

Sigmoid applies sigmoid activation: 1 / (1 + exp(-x)).

func (*Backend) SigmoidBackwardGPU added in v0.6.0 

func (b *Backend) SigmoidBackwardGPU(output, grad *GPUTensor) *GPUTensor

SigmoidBackwardGPU computes gradients for sigmoid activation. d(sigmoid(x))/dx = sigmoid(x) * (1 - sigmoid(x)).

func (*Backend) SigmoidGPU added in v0.6.0 

func (b *Backend) SigmoidGPU(t *GPUTensor) *GPUTensor

SigmoidGPU applies sigmoid activation on GPU: 1 / (1 + exp(-x)). Data stays on GPU - no CPU transfer occurs.

func (*Backend) Sign added in v0.8.0 

func (b *Backend) Sign(x *tensor.RawTensor) *tensor.RawTensor

Sign computes element-wise sign function on GPU.

Only supports float32 dtype for now. Will raise a panic if called with unsupported dtype.

func (*Backend) SignGPU added in v0.8.0 

func (b *Backend) SignGPU(t *GPUTensor) *GPUTensor

SignGPU applies sign activation on GPU. Data stays on GPU - no CPU transfer occurs.

func (*Backend) Sin added in v0.3.0 

func (b *Backend) Sin(x *tensor.RawTensor) *tensor.RawTensor

Sin computes element-wise sine on GPU.

func (*Backend) Softmax 

func (b *Backend) Softmax(x *tensor.RawTensor, dim int) *tensor.RawTensor

Softmax applies softmax along the specified dimension. Supports N-dimensional tensors with dim=-1 (last dimension).

func (*Backend) SoftmaxBackwardGPU added in v0.6.0 

func (b *Backend) SoftmaxBackwardGPU(output, grad *GPUTensor, dim int) *GPUTensor

SoftmaxBackwardGPU computes gradients for softmax activation. d_input[i] = s[i] * (grad[i] - sum(s * grad)) where s = softmax output.

func (*Backend) SoftmaxGPU added in v0.6.0 

func (b *Backend) SoftmaxGPU(t *GPUTensor, dim int) *GPUTensor

SoftmaxGPU applies softmax activation along the specified dimension. For now, only last dimension (dim=-1) is supported efficiently on GPU. Data stays on GPU - no CPU transfer occurs.

func (*Backend) Sqrt added in v0.3.0 

func (b *Backend) Sqrt(x *tensor.RawTensor) *tensor.RawTensor

Sqrt computes element-wise square root on GPU.

func (*Backend) Squeeze added in v0.3.0 

func (b *Backend) Squeeze(x *tensor.RawTensor, dim int) *tensor.RawTensor

Squeeze removes a dimension of size 1 at the specified position.

func (*Backend) Sub 

func (b *Backend) Sub(a, other *tensor.RawTensor) *tensor.RawTensor

Sub performs element-wise subtraction on GPU. Supports float32 and int32 dtypes. In LazyMode (default), returns a lazy tensor that keeps data on GPU.

func (*Backend) SubBackwardGPU added in v0.6.0 

func (b *Backend) SubBackwardGPU(_, _, grad *GPUTensor) (*GPUTensor, *GPUTensor)

SubBackwardGPU computes gradients for element-wise subtraction. d(a-b)/da = 1, d(a-b)/db = -1.

func (*Backend) SubGPU added in v0.6.0 

func (b *Backend) SubGPU(a, c *GPUTensor) *GPUTensor

SubGPU performs element-wise subtraction on GPU tensors. Data stays on GPU - no CPU transfer occurs.

func (*Backend) SubScalar added in v0.3.0 

func (b *Backend) SubScalar(x *tensor.RawTensor, scalar any) *tensor.RawTensor

SubScalar subtracts a scalar from tensor elements on GPU.

func (*Backend) Sum added in v0.3.0 

func (b *Backend) Sum(x *tensor.RawTensor) *tensor.RawTensor

Sum computes the sum of all elements on GPU.

func (*Backend) SumDim added in v0.3.0 

func (b *Backend) SumDim(x *tensor.RawTensor, dim int, keepDim bool) *tensor.RawTensor

SumDim sums along a dimension.

In LazyMode, dispatches a GPU compute shader (sumDimGeneralShader) that keeps data on GPU — zero CPU allocation for the result. Falls back to CPU reduction when LazyMode is disabled or the source tensor is not GPU-resident.

func (*Backend) SumDimGPU added in v0.6.0 

func (b *Backend) SumDimGPU(t *GPUTensor, dim int, keepDim bool) *GPUTensor

SumDimGPU computes sum along the last dimension. Input: [batch, dim], Output: [batch].

func (*Backend) Tanh 

func (b *Backend) Tanh(x *tensor.RawTensor) *tensor.RawTensor

Tanh applies tanh activation.

func (*Backend) TanhBackwardGPU added in v0.6.0 

func (b *Backend) TanhBackwardGPU(output, grad *GPUTensor) *GPUTensor

TanhBackwardGPU computes gradients for tanh activation. d(tanh(x))/dx = 1 - tanh(x)^2.

func (*Backend) TanhGPU added in v0.6.0 

func (b *Backend) TanhGPU(t *GPUTensor) *GPUTensor

TanhGPU applies tanh activation on GPU. Data stays on GPU - no CPU transfer occurs.

func (*Backend) Transpose 

func (b *Backend) Transpose(t *tensor.RawTensor, axes ...int) *tensor.RawTensor

Transpose transposes the tensor by permuting its dimensions. GPU-accelerated for 2D (optimized) and ND tensors (up to 6D).

func (*Backend) TransposeGPU added in v0.6.0 

func (b *Backend) TransposeGPU(t *GPUTensor, axes ...int) *GPUTensor

TransposeGPU transposes a 2D tensor on GPU. Data stays on GPU - no CPU transfer occurs.

func (*Backend) UnregisterLiveGPU added in v0.9.1 

func (b *Backend) UnregisterLiveGPU(l *tensor.LazyGPUData)

UnregisterLiveGPU removes a LazyGPUData from the live tensor set.

func (*Backend) Unsqueeze added in v0.3.0 

func (b *Backend) Unsqueeze(x *tensor.RawTensor, dim int) *tensor.RawTensor

Unsqueeze adds a dimension of size 1 at the specified position.

func (*Backend) UploadTensor added in v0.6.0 

func (b *Backend) UploadTensor(raw *tensor.RawTensor) *GPUTensor

UploadTensor uploads a CPU tensor to GPU memory. Returns a GPUTensor that can be used for lazy GPU operations.

func (*Backend) Where added in v0.3.0 

func (b *Backend) Where(condition, x, y *tensor.RawTensor) *tensor.RawTensor

Where performs conditional element selection on GPU. result[i] = condition[i] != 0 ? x[i] : y[i].

func (*Backend) ZerosGPU added in v0.6.0 

func (b *Backend) ZerosGPU(shape tensor.Shape, dtype tensor.DataType) *GPUTensor

ZerosGPU creates a zero-filled GPU tensor. Data is initialized to zeros on GPU.

type BufferPool 

type BufferPool struct {
	// contains filtered or unexported fields
}

BufferPool manages GPU buffer reuse to reduce allocation overhead. Buffers are categorized by size and usage flags.

func NewBufferPool 

func NewBufferPool(device *wgpu.Device) *BufferPool

NewBufferPool creates a new buffer pool for the given device.

func (*BufferPool) Acquire 

func (p *BufferPool) Acquire(size uint64, usage gputypes.BufferUsage) *wgpu.Buffer

Acquire gets a buffer from the pool or creates a new one. Returns a buffer that matches or exceeds the requested size and usage.

func (*BufferPool) Clear 

func (p *BufferPool) Clear()

Clear releases all pooled buffers. Should be called when the backend is released.

func (*BufferPool) Release 

func (p *BufferPool) Release(buffer *wgpu.Buffer, size uint64, usage gputypes.BufferUsage)

Release returns a buffer to the pool for reuse. If the pool is full, the buffer is immediately released.

func (*BufferPool) Stats 

func (p *BufferPool) Stats() (allocated, released, hits, misses uint64, pooledCount int)

Stats returns statistics about buffer pool usage.

type BufferSize 

type BufferSize int

BufferSize represents different buffer size categories for pooling. 

const (
	// SmallBuffer for tensors < 4KB.
	SmallBuffer BufferSize = iota
	// MediumBuffer for tensors 4KB-1MB.
	MediumBuffer
	// LargeBuffer for tensors > 1MB.
	LargeBuffer
)

type CommandBatch added in v0.6.0 

type CommandBatch struct {
	// contains filtered or unexported fields
}

CommandBatch accumulates GPU operations for single submission. Instead of submitting each operation separately (causing GPU overhead), we collect all operations in a batch and submit them together.

func (*CommandBatch) Add added in v0.6.0 

func (batch *CommandBatch) Add(name string, output *GPUTensor, execFunc func()) *CommandBatch

Add adds an operation to the batch. The operation function should encode the compute pass but NOT submit it. Returns the batch for method chaining.

func (*CommandBatch) Count added in v0.6.0 

func (batch *CommandBatch) Count() int

Count returns the number of operations in the batch.

func (*CommandBatch) Submit added in v0.6.0 

func (batch *CommandBatch) Submit()

Submit executes all batched operations in a single GPU submission. This dramatically reduces GPU overhead compared to submitting each operation separately.

Example performance difference: 

3 separate submissions: encode → submit → wait (×3) = ~1.5ms overhead
1 batched submission:   encode → encode → encode → submit → wait = ~0.5ms overhead

The batch is consumed after Submit() and cannot be reused.

type DeviceLimits added in v0.9.1 

type DeviceLimits struct {
	MaxBufferSize           uint64 // device.Limits().MaxBufferSize
	MaxStorageBufferBinding uint64 // device.Limits().MaxStorageBufferBindingSize
	MinAlignment            uint64 // device.Limits().MinUniformBufferOffsetAlignment
}

DeviceLimits holds GPU device properties used for pool configuration. Queried from wgpu adapter/device at backend initialization.

func QueryDeviceLimits added in v0.9.1 

func QueryDeviceLimits(device *wgpu.Device) DeviceLimits

QueryDeviceLimits reads limits from a wgpu device.

type ExclusivePool added in v0.9.1 

type ExclusivePool struct {
	// contains filtered or unexported fields
}

ExclusivePool manages GPU buffer reuse following the Burn/CubeCL pattern. Each page holds exactly one wgpu.Buffer (wgpu constraint: buffer is either read-only or read_write, not both). Pages are reused across training steps to eliminate allocation churn.

After warmup (1 training step), allocation count drops to near-zero.

Reference: cubecl-runtime/src/memory_management/memory_pool/exclusive_pool.rs.

func (*ExclusivePool) Accept added in v0.9.1 

func (p *ExclusivePool) Accept(size uint64) bool

Accept returns true if this pool handles allocations of the given size. Mirrors CubeCL ExclusiveMemoryPool::accept().

func (*ExclusivePool) Acquire added in v0.9.1 

func (p *ExclusivePool) Acquire(size uint64) (*wgpu.Buffer, error)

Acquire returns a GPU buffer of at least `size` bytes. Reuses a free page if available, otherwise allocates a new one.

func (*ExclusivePool) Cleanup added in v0.9.1 

func (p *ExclusivePool) Cleanup(explicit bool)

Cleanup deallocates pages that have been unused for gpuPoolFreeThresh consecutive cycles.

func (*ExclusivePool) Destroy added in v0.9.1 

func (p *ExclusivePool) Destroy()

Destroy releases all pages. Called from Backend.Release().

func (*ExclusivePool) Release added in v0.9.1 

func (p *ExclusivePool) Release(buffer *wgpu.Buffer)

Release returns a buffer to the pool for reuse. The buffer is NOT destroyed.

func (*ExclusivePool) Stats added in v0.9.1 

func (p *ExclusivePool) Stats() (total, inUse, free int, bytesReserved uint64)

Stats returns pool statistics.

type GPUTape added in v0.6.0 

type GPUTape struct {
	// contains filtered or unexported fields
}

GPUTape records GPU operations for backward pass. All operations and gradients stay on GPU for maximum performance.

func NewGPUTape added in v0.6.0 

func NewGPUTape(b *Backend) *GPUTape

NewGPUTape creates a new gradient tape for GPU operations.

func (*GPUTape) Backward added in v0.6.0 

func (tape *GPUTape) Backward(loss *GPUTensor) map[*GPUTensor]*GPUTensor

Backward computes gradients for all inputs by walking the tape in reverse. All operations stay on GPU - no CPU transfers occur.

Algorithm:

  1. Start with loss gradient (typically ones for scalar loss)
  2. Walk operations in reverse order
  3. For each operation, compute input gradients using chain rule
  4. Accumulate gradients when the same tensor is used multiple times

Returns a map from GPUTensor to its accumulated gradient (also GPUTensor).

func (*GPUTape) Clear added in v0.6.0 

func (tape *GPUTape) Clear()

Clear resets the tape, removing all recorded operations. Recording state is preserved.

func (*GPUTape) Disable added in v0.6.0 

func (tape *GPUTape) Disable()

Disable disables operation recording.

func (*GPUTape) Enable added in v0.6.0 

func (tape *GPUTape) Enable()

Enable enables operation recording.

func (*GPUTape) IsEnabled added in v0.6.0 

func (tape *GPUTape) IsEnabled() bool

IsEnabled returns true if the tape is currently recording operations.

func (*GPUTape) NumOps added in v0.6.0 

func (tape *GPUTape) NumOps() int

NumOps returns the number of recorded operations.

func (*GPUTape) Record added in v0.6.0 

func (tape *GPUTape) Record(name string, inputs []*GPUTensor, output *GPUTensor, backward func(*GPUTensor) []*GPUTensor)

Record records an operation for backward pass. The backward function should compute gradients for all inputs given the output gradient.

type GPUTensor added in v0.6.0 

type GPUTensor struct {
	// contains filtered or unexported fields
}

GPUTensor holds tensor data in GPU memory without transferring to CPU. This enables efficient GPU-to-GPU operations without the overhead of readBuffer() calls.

func (*GPUTensor) Backward added in v0.6.0 

func (t *GPUTensor) Backward()

Backward computes gradients for this tensor. This is a convenience method that creates a gradient of ones and calls tape.Backward().

func (*GPUTensor) Buffer added in v0.6.0 

func (t *GPUTensor) Buffer() *wgpu.Buffer

Buffer returns the underlying GPU buffer. This is exposed for internal backend operations.

func (*GPUTensor) ByteSize added in v0.6.0 

func (t *GPUTensor) ByteSize() uint64

ByteSize returns the total memory size in bytes.

func (*GPUTensor) DType added in v0.6.0 

func (t *GPUTensor) DType() tensor.DataType

DType returns the tensor's data type.

func (*GPUTensor) Eval added in v0.6.0 

func (t *GPUTensor) Eval() *GPUTensor

Eval forces computation of lazy tensor using batched submission. Collects all dependencies and submits them in a single GPU command buffer. This reduces GPU overhead compared to submitting each operation separately.

func (*GPUTensor) Grad added in v0.6.0 

func (t *GPUTensor) Grad() *GPUTensor

Grad returns the accumulated gradient for this tensor. Returns nil if no gradient has been computed.

func (*GPUTensor) Item added in v0.6.0 

func (t *GPUTensor) Item() float32

Item returns the single scalar value from a tensor. This is useful for extracting loss values during training. Panics if tensor has more than one element.

func (*GPUTensor) NumElements added in v0.6.0 

func (t *GPUTensor) NumElements() int

NumElements returns the total number of elements in the tensor.

func (*GPUTensor) Release added in v0.6.0 

func (t *GPUTensor) Release()

Release releases the GPU buffer and frees memory. This should be called when the tensor is no longer needed.

func (*GPUTensor) RequiresGrad added in v0.6.0 

func (t *GPUTensor) RequiresGrad() bool

RequiresGrad returns whether this tensor requires gradient computation.

func (*GPUTensor) SetRequiresGrad added in v0.6.0 

func (t *GPUTensor) SetRequiresGrad(requires bool) *GPUTensor

SetRequiresGrad sets whether this tensor requires gradient computation. Returns the tensor for method chaining. Note: PyTorch uses requires_grad_ (underscore suffix for in-place). In Go, we use SetRequiresGrad for clarity.

func (*GPUTensor) Shape added in v0.6.0 

func (t *GPUTensor) Shape() tensor.Shape

Shape returns the tensor's shape.

func (*GPUTensor) ToCPU added in v0.6.0 

func (t *GPUTensor) ToCPU() *tensor.RawTensor

ToCPU transfers tensor data from GPU to CPU memory. This is an expensive operation and should be used sparingly. Returns a new RawTensor with data copied from GPU.

func (*GPUTensor) ZeroGrad added in v0.6.0 

func (t *GPUTensor) ZeroGrad()

ZeroGrad clears the accumulated gradient.

type MemoryStats 

type MemoryStats struct {
	// Total bytes allocated since backend creation
	TotalAllocatedBytes uint64
	// Peak memory usage in bytes
	PeakMemoryBytes uint64
	// Number of currently active buffers
	ActiveBuffers int64
	// Buffer pool statistics
	PoolAllocated uint64
	PoolReleased  uint64
	PoolHits      uint64
	PoolMisses    uint64
	PooledBuffers int
}

MemoryStats represents GPU memory usage statistics.

type TieredPool added in v0.9.1 

type TieredPool struct {
	// contains filtered or unexported fields
}

TieredPool routes allocations to size-appropriate ExclusivePool buckets. Smaller allocations go to smaller pools, reducing fragmentation. Each pool has its own dealloc period — smaller pools clean up more aggressively.

func NewTieredPool added in v0.9.1 

func NewTieredPool(device *wgpu.Device) *TieredPool

NewTieredPool creates a multi-tier GPU buffer pool from device limits. Budget is configurable via BORN_GPU_BUDGET_MB env var.

Reference: CubeCL generate_bucket_sizes() + MemoryManagement::from_configuration().

func (*TieredPool) Acquire added in v0.9.1 

func (tp *TieredPool) Acquire(size uint64) (*wgpu.Buffer, error)

Acquire routes the allocation to the first pool that accepts the size. If allocation fails (GPU OOM), calls the onOOM callback to flush pending GPU commands and free completed buffers, then retries once.

func (*TieredPool) BucketSizes added in v0.9.1 

func (tp *TieredPool) BucketSizes() []uint64

BucketSizes returns the max allocation size for each tier (for diagnostics).

func (*TieredPool) BudgetBytes added in v0.9.1 

func (tp *TieredPool) BudgetBytes() uint64

BudgetBytes returns the configured GPU memory budget.

func (*TieredPool) Cleanup added in v0.9.1 

func (tp *TieredPool) Cleanup(explicit bool)

Cleanup runs cleanup on all pool tiers.

func (*TieredPool) Destroy added in v0.9.1 

func (tp *TieredPool) Destroy()

Destroy releases all pages in all tiers.

func (*TieredPool) DeviceLimitsInfo added in v0.9.1 

func (tp *TieredPool) DeviceLimitsInfo() (maxPageSize, alignment uint64)

DeviceLimitsInfo returns the device limits used for pool configuration.

func (*TieredPool) Release added in v0.9.1 

func (tp *TieredPool) Release(buffer *wgpu.Buffer)

Release returns a buffer to the appropriate pool.

func (*TieredPool) Stats added in v0.9.1 

func (tp *TieredPool) Stats() (total, inUse, free int, bytesReserved uint64)

Stats returns aggregate pool statistics.

func (*TieredPool) TotalReserved added in v0.9.1 

func (tp *TieredPool) TotalReserved() uint64

TotalReserved returns total GPU bytes reserved across all tiers.

type TrainingScope added in v0.9.1 

type TrainingScope struct {
	// contains filtered or unexported fields
}

TrainingScope tracks intermediate GPU tensors produced during a training step for bulk release after the step completes. This is the primary mechanism for explicit GPU memory lifecycle management under ADR-015.

Usage: 

scope := webgpu.NewTrainingScope()
defer scope.Release()
// ... forward + backward pass, call scope.Track(t) on intermediates ...

func NewTrainingScope added in v0.9.1 

func NewTrainingScope() *TrainingScope

NewTrainingScope creates a TrainingScope for tracking intermediate tensors.

func (*TrainingScope) Release added in v0.9.1 

func (s *TrainingScope) Release()

Release calls ReleaseGPU on all tracked tensors and resets the scope for reuse. Idempotent — safe to call multiple times. The GPU buffers are enqueued for deferred destruction via wgpu's DestroyQueue and released after the next queue.Submit completes.

func (*TrainingScope) Track added in v0.9.1 

func (s *TrainingScope) Track(t *tensor.RawTensor)

Track registers a tensor for release when Release is called. Safe to call with nil (no-op). Does not take ownership — the caller is responsible for not using the tensor after calling scope.Release.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.