Package io.github.kirstenali.deepj.tensor

Class ComputeGraph

java.lang.Object

io.github.kirstenali.deepj.tensor.ComputeGraph

public final class ComputeGraph
extends Object

Collects GPU operations lazily and flushes them as a single command buffer.

Operations are recorded as a flat int[] command stream. Buffer IDs reference persistent GPU-side buffers managed by a GpuRuntime. Data stays GPU-resident between ops -- only uploaded at graph entry and downloaded on materialization.

This class is backend-agnostic: Metal, CUDA, Vulkan, etc. are all supported by supplying the appropriate GpuRuntime implementation.

Field Summary

Fields

Modifier and Type	Field	Description
static final int	OP_ADAMW_UPDATE
static final int	OP_ADD
static final int	OP_DIVIDE
static final int	OP_EXP
static final int	OP_GELU
static final int	OP_GELU_BACKWARD
static final int	OP_LAYERNORM_BACKWARD
static final int	OP_LOG
static final int	OP_MATMUL
static final int	OP_MULTIPLY
static final int	OP_MULTIPLY_SCALAR
static final int	OP_NEG
static final int	OP_RELU
static final int	OP_RELU_BACKWARD
static final int	OP_SIGMOID
static final int	OP_SOFTMAX_BACKWARD
static final int	OP_SOFTMAX_ROWS
static final int	OP_SQRT
static final int	OP_SUBTRACT
static final int	OP_TANH

Constructor Summary

Constructors

Constructor	Description
ComputeGraph(GpuRuntime runtime)	Create a ComputeGraph backed by the given GPU runtime.

Method Summary

Modifier and Type	Method	Description
Tensor	createOutputTensor(GpuBuffer buf)	Create a Tensor backed by a GpuBuffer.
GpuBuffer	ensureGpuBuffer(Tensor t)	Ensure a tensor has a GpuBuffer.
void	flush()	Flush all recorded ops to the GPU as a single command buffer.
boolean	isEmpty()
void	materialize(Tensor t)	Materialize a tensor: flush pending ops if needed, then download GPU data to CPU.
GpuBuffer	newOutputBuffer(int rows, int cols)	Allocate a new output buffer (result of a GPU op).
void	recordAdamWUpdate(GpuBuffer w, GpuBuffer g, GpuBuffer mt, GpuBuffer vt, float lr, float beta1, float beta2, float eps, float weightDecay, float bc1, float bc2, int n)	Record in-place AdamW update: [OP_ADAMW_UPDATE, wId, gId, mtId, vtId, lrBits, beta1Bits, beta2Bits, epsBits, weightDecayBits, bc1Bits, bc2Bits, n]
void	recordBinary(int opCode, GpuBuffer a, GpuBuffer b, GpuBuffer out)	Record a binary element-wise op: [opCode, aId, bId, outId, n]
void	recordLayerNormBackward(GpuBuffer dXHat, GpuBuffer xHat, GpuBuffer std, GpuBuffer out, int rows, int cols)	Record layer norm backward: [OP_LAYERNORM_BACKWARD, dXHatId, xHatId, stdId, outId, rows, cols]
void	recordMatmul(GpuBuffer a, GpuBuffer b, GpuBuffer out, int m, int n, int k)	Record matmul: [OP_MATMUL, aId, bId, outId, m, n, k]
void	recordMultiplyScalar(GpuBuffer in, GpuBuffer out, float scalar)	Record scalar multiply: [OP_MULTIPLY_SCALAR, inId, outId, scalarBits, n]
void	recordSoftmaxBackward(GpuBuffer gradOutput, GpuBuffer softmaxOut, GpuBuffer out, int rows, int cols)	Record softmax backward: [OP_SOFTMAX_BACKWARD, gradId, softmaxId, outId, rows, cols]
void	recordSoftmaxRows(GpuBuffer in, GpuBuffer out, int rows, int cols)	Record softmax rows: [OP_SOFTMAX_ROWS, inId, outId, rows, cols]
void	recordUnary(int opCode, GpuBuffer in, GpuBuffer out)	Record a unary op: [opCode, inId, outId, n]
void	releaseAll()	Release all GPU buffers and reset the graph completely.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Field Details

OP_ADD

public static final int OP_ADD

See Also:

• Constant Field Values

OP_SUBTRACT

public static final int OP_SUBTRACT

See Also:

• Constant Field Values

OP_MULTIPLY

public static final int OP_MULTIPLY

See Also:

• Constant Field Values

OP_DIVIDE

public static final int OP_DIVIDE

See Also:

• Constant Field Values

OP_MATMUL

public static final int OP_MATMUL

See Also:

• Constant Field Values

OP_MULTIPLY_SCALAR

public static final int OP_MULTIPLY_SCALAR

See Also:

• Constant Field Values

OP_SQRT

public static final int OP_SQRT

See Also:

• Constant Field Values

OP_NEG

public static final int OP_NEG

See Also:

• Constant Field Values

OP_EXP

public static final int OP_EXP

See Also:

• Constant Field Values

OP_LOG

public static final int OP_LOG

See Also:

• Constant Field Values

OP_TANH

public static final int OP_TANH

See Also:

• Constant Field Values

OP_SIGMOID

public static final int OP_SIGMOID

See Also:

• Constant Field Values

OP_RELU

public static final int OP_RELU

See Also:

• Constant Field Values

OP_RELU_BACKWARD

public static final int OP_RELU_BACKWARD

See Also:

• Constant Field Values

OP_GELU

public static final int OP_GELU

See Also:

• Constant Field Values

OP_GELU_BACKWARD

public static final int OP_GELU_BACKWARD

See Also:

• Constant Field Values

OP_SOFTMAX_ROWS

public static final int OP_SOFTMAX_ROWS

See Also:

• Constant Field Values

OP_SOFTMAX_BACKWARD

public static final int OP_SOFTMAX_BACKWARD

See Also:

• Constant Field Values

OP_LAYERNORM_BACKWARD

public static final int OP_LAYERNORM_BACKWARD

See Also:

• Constant Field Values

OP_ADAMW_UPDATE

public static final int OP_ADAMW_UPDATE

See Also:

• Constant Field Values

Constructor Details

ComputeGraph

public ComputeGraph(GpuRuntime runtime)

Create a ComputeGraph backed by the given GPU runtime.

Parameters:

runtime - the native driver abstraction (Metal, CUDA, etc.)

Method Details

ensureGpuBuffer

public GpuBuffer ensureGpuBuffer(Tensor t)

Ensure a tensor has a GpuBuffer. If it already has one (from a previous op), reuse it. Otherwise allocate a new buffer and schedule upload of its CPU data.

newOutputBuffer

public GpuBuffer newOutputBuffer(int rows,
 int cols)

Allocate a new output buffer (result of a GPU op). Not yet allocated on native side.

createOutputTensor

public Tensor createOutputTensor(GpuBuffer buf)

Create a Tensor backed by a GpuBuffer. The data[][] is allocated but stale.

recordBinary

public void recordBinary(int opCode,
 GpuBuffer a,
 GpuBuffer b,
 GpuBuffer out)

Record a binary element-wise op: [opCode, aId, bId, outId, n]

recordMatmul

public void recordMatmul(GpuBuffer a,
 GpuBuffer b,
 GpuBuffer out,
 int m,
 int n,
 int k)

Record matmul: [OP_MATMUL, aId, bId, outId, m, n, k]

recordUnary

public void recordUnary(int opCode,
 GpuBuffer in,
 GpuBuffer out)

Record a unary op: [opCode, inId, outId, n]

recordMultiplyScalar

public void recordMultiplyScalar(GpuBuffer in,
 GpuBuffer out,
 float scalar)

Record scalar multiply: [OP_MULTIPLY_SCALAR, inId, outId, scalarBits, n]

recordSoftmaxRows

public void recordSoftmaxRows(GpuBuffer in,
 GpuBuffer out,
 int rows,
 int cols)

Record softmax rows: [OP_SOFTMAX_ROWS, inId, outId, rows, cols]

recordSoftmaxBackward

public void recordSoftmaxBackward(GpuBuffer gradOutput,
 GpuBuffer softmaxOut,
 GpuBuffer out,
 int rows,
 int cols)

Record softmax backward: [OP_SOFTMAX_BACKWARD, gradId, softmaxId, outId, rows, cols]

recordLayerNormBackward

public void recordLayerNormBackward(GpuBuffer dXHat,
 GpuBuffer xHat,
 GpuBuffer std,
 GpuBuffer out,
 int rows,
 int cols)

Record layer norm backward: [OP_LAYERNORM_BACKWARD, dXHatId, xHatId, stdId, outId, rows, cols]

recordAdamWUpdate

public void recordAdamWUpdate(GpuBuffer w,
 GpuBuffer g,
 GpuBuffer mt,
 GpuBuffer vt,
 float lr,
 float beta1,
 float beta2,
 float eps,
 float weightDecay,
 float bc1,
 float bc2,
 int n)

Record in-place AdamW update: [OP_ADAMW_UPDATE, wId, gId, mtId, vtId, lrBits, beta1Bits, beta2Bits, epsBits, weightDecayBits, bc1Bits, bc2Bits, n]

isEmpty

public boolean isEmpty()

flush

public void flush()

Flush all recorded ops to the GPU as a single command buffer. After flush, GPU buffers hold computed results; CPU data is stale.

materialize

public void materialize(Tensor t)

Materialize a tensor: flush pending ops if needed, then download GPU data to CPU.

releaseAll

public void releaseAll()

Release all GPU buffers and reset the graph completely.