Docs / jax4j-engine-cuda / com.nablatensor.jax4j.backend.cuda

final class

CudaBackend

NVIDIA CUDA backend. All math runs in custom kernels that are compiled to PTX at runtime with NVRTC and launched through the CUDA Driver API. Tensors stay resident in device memory between operations.

Phase-0 note: a single driver context is created on first use; the notebook exercises the backend from one thread. Per-thread primary-context handling is tracked for a later phase.

View source ↗Full module reference on javadoc.io ↗

Methods

String name()

DeviceType deviceType()

boolean isAvailable()

int priority()

String deviceName()

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

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

DeviceBuffer full(float value, Shape shape, DType dtype, Device device)

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

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)

DeviceBuffer stackAxis0(DeviceBuffer[] inputs)

DeviceBuffer reduceSum(DeviceBuffer a)

DeviceBuffer reduceMax(DeviceBuffer a)

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

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

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

DeviceBuffer reluBackward(DeviceBuffer upstream, DeviceBuffer input)

DeviceBuffer sumAxis0(DeviceBuffer a)

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

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

DeviceBuffer argMaxAxis(DeviceBuffer a, int axis)

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

DeviceBuffer reshape(DeviceBuffer a, Shape target)

DeviceBuffer broadcastTo(DeviceBuffer a, Shape target)

Broadcasts entirely on-device: source strides (0 for dims being broadcast) and target dims are computed on the host (cheap, just a few ints) and passed as kernel scalars, capped at MAX_BROADCAST_RANK dims (padded with size-1/stride-0 entries) - every shape this project produces is rank <= 2. No download/upload round trip, unlike the previous host-side implementation.

void sync()

void release(DeviceBuffer buffer)

DeviceBuffer fused(Expr expr, DeviceBuffer[] inputs)

Runs a whole elementwise expression chain as one generated kernel launch, instead of one launch (and one intermediate device buffer) per op. The compiled kernel is cached by its generated source, so repeated fused calls for the same op chain (any shape) reuse it without recompiling.