webgpu

package
v0.9.19 Latest Latest
Warning

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

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

Documentation

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 

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 

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 

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 

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 

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

AddScalar adds a scalar to tensor elements on GPU.

func (*Backend) And 

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 

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

Argmax returns indices of maximum values along dimension on GPU.

func (*Backend) BatchMatMul 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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

Cos computes element-wise cosine on GPU.

func (*Backend) DeferReleaseGPUBuffer 

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 

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 

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 

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

DivScalar divides tensor elements by a scalar on GPU.

func (*Backend) Embedding 

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 

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 

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

Erf computes element-wise error function on GPU.

func (*Backend) ErfGPU 

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

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

func (*Backend) Exp 

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

Exp computes element-wise exponential on GPU.

func (*Backend) Expand 

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 

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 

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 

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 

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 

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 

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 

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

Log computes natural logarithm element-wise on GPU.

func (*Backend) Lower 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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

Not performs element-wise logical NOT on GPU.

func (*Backend) NotEqual 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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

Rsqrt computes element-wise reciprocal square root on GPU.

func (*Backend) ScatterAdd 

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 

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 

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 

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 

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 

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 

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 

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

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

func (*Backend) Sin 

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 

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 

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 

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

Sqrt computes element-wise square root on GPU.

func (*Backend) Squeeze 

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 

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 

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 

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

SubScalar subtracts a scalar from tensor elements on GPU.

func (*Backend) Sum 

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

Sum computes the sum of all elements on GPU.

func (*Backend) SumDim 

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 

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 

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 

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 

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 

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

UnregisterLiveGPU removes a LazyGPUData from the live tensor set.

func (*Backend) Unsqueeze 

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 

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 

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 

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 

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 

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 

func (batch *CommandBatch) Count() int

Count returns the number of operations in the batch.

func (*CommandBatch) Submit 

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 

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 

func QueryDeviceLimits(device *wgpu.Device) DeviceLimits

QueryDeviceLimits reads limits from a wgpu device.

type ExclusivePool 

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 

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 

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 

func (p *ExclusivePool) Cleanup(explicit bool)

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

func (*ExclusivePool) Destroy 

func (p *ExclusivePool) Destroy()

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

func (*ExclusivePool) Release 

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

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

func (*ExclusivePool) Stats 

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

Stats returns pool statistics.

type GPUTape 

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 

func NewGPUTape(b *Backend) *GPUTape

NewGPUTape creates a new gradient tape for GPU operations.

func (*GPUTape) Backward 

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 

func (tape *GPUTape) Clear()

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

func (*GPUTape) Disable 

func (tape *GPUTape) Disable()

Disable disables operation recording.

func (*GPUTape) Enable 

func (tape *GPUTape) Enable()

Enable enables operation recording.

func (*GPUTape) IsEnabled 

func (tape *GPUTape) IsEnabled() bool

IsEnabled returns true if the tape is currently recording operations.

func (*GPUTape) NumOps 

func (tape *GPUTape) NumOps() int

NumOps returns the number of recorded operations.

func (*GPUTape) Record 

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 

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 

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 

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

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

func (*GPUTensor) ByteSize 

func (t *GPUTensor) ByteSize() uint64

ByteSize returns the total memory size in bytes.

func (*GPUTensor) DType 

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

DType returns the tensor's data type.

func (*GPUTensor) Eval 

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 

func (t *GPUTensor) Grad() *GPUTensor

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

func (*GPUTensor) Item 

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 

func (t *GPUTensor) NumElements() int

NumElements returns the total number of elements in the tensor.

func (*GPUTensor) Release 

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 

func (t *GPUTensor) RequiresGrad() bool

RequiresGrad returns whether this tensor requires gradient computation.

func (*GPUTensor) SetRequiresGrad 

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 

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

Shape returns the tensor's shape.

func (*GPUTensor) ToCPU 

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 

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 

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 

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 

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 

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

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

func (*TieredPool) BudgetBytes 

func (tp *TieredPool) BudgetBytes() uint64

BudgetBytes returns the configured GPU memory budget.

func (*TieredPool) Cleanup 

func (tp *TieredPool) Cleanup(explicit bool)

Cleanup runs cleanup on all pool tiers.

func (*TieredPool) Destroy 

func (tp *TieredPool) Destroy()

Destroy releases all pages in all tiers.

func (*TieredPool) DeviceLimitsInfo 

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

DeviceLimitsInfo returns the device limits used for pool configuration.

func (*TieredPool) Release 

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

Release returns a buffer to the appropriate pool.

func (*TieredPool) Stats 

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

Stats returns aggregate pool statistics.

func (*TieredPool) TotalReserved 

func (tp *TieredPool) TotalReserved() uint64

TotalReserved returns total GPU bytes reserved across all tiers.

type TrainingScope 

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 

func NewTrainingScope() *TrainingScope

NewTrainingScope creates a TrainingScope for tracking intermediate tensors.

func (*TrainingScope) Release 

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 

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.