Docs / nablatensor-tensor / com.nablatensor.tensor.spi

interface

ComputeBackend

A compute backend. Implementations are discovered with ServiceLoader and must be registered in META-INF/services/com.nablatensor.tensor.spi.ComputeBackend.

Every numerical operation is executed by a custom kernel owned by the backend (runtime-compiled CUDA/ROCm, or a generated/vectorized CPU kernel).

View source ↗Full module reference on javadoc.io ↗

Methods

String name()

DeviceType deviceType()

boolean isAvailable()

Whether this backend can run on the current machine (drivers/GPU present).

int priority()

Higher wins when Backend.AUTO selects a default (CUDA > Vulkan > ROCm > CPU).

DeviceBuffer upload(float[] data, Shape shape, DType dtype, Device device)

DeviceBuffer randomUniform(long seed, long counter, Shape shape, Device device)

Deterministic counter-based uniform values in [0, 1).

DeviceBuffer randomNormal(long seed, long counter, Shape shape, Device device)

Deterministic counter-based standard-normal values generated with Box-Muller.

float[] download(DeviceBuffer buffer)

DeviceBuffer binary(Op op, DeviceBuffer a, DeviceBuffer b)

DeviceBuffer scalar(Op op, DeviceBuffer a, double value)

DeviceBuffer unary(Op op, DeviceBuffer a)

DeviceBuffer transpose(DeviceBuffer a)

DeviceBuffer matmul(DeviceBuffer a, DeviceBuffer b)

DeviceBuffer batchedMatmul(DeviceBuffer a, DeviceBuffer b)

DeviceBuffer sliceAxis0(DeviceBuffer input, int index)

Copies one slice from axis zero, removing that axis from the result shape.

DeviceBuffer stackAxis0(DeviceBuffer[] inputs)

Stacks equally-shaped buffers along a new leading axis.

DeviceBuffer reduceSum(DeviceBuffer a)

Sums every element, returning a single-element buffer of shape (1).

DeviceBuffer reduceMax(DeviceBuffer a)

Maximum over every element, returning a single-element buffer of shape (1).

DeviceBuffer sumAxis0(DeviceBuffer a)

Sums a rank-2 buffer over axis 0, returning a rank-1 buffer of shape (cols).

DeviceBuffer sumAxis(DeviceBuffer a, int axis, boolean keepDims)

DeviceBuffer maxAxis(DeviceBuffer a, int axis, boolean keepDims)

DeviceBuffer argMaxAxis(DeviceBuffer a, int axis)

Returns axis indices in an F32 buffer. Indices above 2^24 cannot be represented exactly until integer tensor dtypes are supported.

DeviceBuffer maxAxisBackward(DeviceBuffer upstream, DeviceBuffer input, int axis)

Routes one upstream value per reduced slice to the first maximum input position.

DeviceBuffer reshape(DeviceBuffer a, Shape target)

Copies a buffer while changing only its shape. Backends may override this to keep the copy device-local.

DeviceBuffer conv2d(DeviceBuffer x, DeviceBuffer w, ConvSpec spec)

2-D convolution, no bias. x is (batch, inC*inH*inW), w is (outC, inC*k*k), result is (batch, outC*outH*outW).

DeviceBuffer conv2dGradInput(DeviceBuffer upstream, DeviceBuffer w, ConvSpec spec)

Gradient of conv2d with respect to its input.

DeviceBuffer conv2dGradWeight(DeviceBuffer x, DeviceBuffer upstream, ConvSpec spec)

Gradient of conv2d with respect to its weights.

DeviceBuffer reluBackward(DeviceBuffer upstream, DeviceBuffer input)

Routes upstream through a ReLU: input > 0 ? upstream : 0.

DeviceBuffer broadcastTo(DeviceBuffer a, Shape target)

Broadcasts a to target following broadcast's rules.

DeviceBuffer fused(Expr expr, DeviceBuffer[] inputs)

Executes a fused elementwise expression tree (expr, reading from inputs) as a single unit. The default walks the tree with the primitive unary/binary/scalar ops (no real fusion, one call per node); CUDA/CPU backends override this to run the whole chain as a single kernel launch / pass.

void sync()

Block until all queued kernels on this backend have completed.

void release(DeviceBuffer buffer)

Deterministically frees a buffer now, instead of waiting on GC/Cleaner timing.