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 Details

  • 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[] buffer is allocated but stale.

    • bindTensorToBuffer

      public void bindTensorToBuffer(Tensor t, GpuBuffer buf)
      Rebind an existing tensor to a GPU buffer and track ownership for lifecycle management.

    • 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]

    • recordScalarUnary

      public void recordScalarUnary(int opCode, GpuBuffer in, GpuBuffer out, float scalar)
      Record scalar add/divide: [opCode, inId, outId, scalarBits, n]

    • recordPow

      public void recordPow(GpuBuffer in, GpuBuffer out, float exponent)
      Record pow: [OP_POW, inId, outId, exponentBits, n]

    • recordClamp

      public void recordClamp(GpuBuffer in, GpuBuffer out, float min, float max)
      Record clamp: [OP_CLAMP, inId, outId, minBits, maxBits, n]

    • recordScatterAddRows

      public void recordScatterAddRows(GpuBuffer target, GpuBuffer indices, GpuBuffer grad, int targetRows, int targetCols, int nIndices)
      Record scatter-add-rows: [OP_SCATTER_ADD_ROWS, targetId, indicesId, gradId, targetRows, targetCols, nIdx]

    • recordScatterAddRowsAtomic

      public void recordScatterAddRowsAtomic(GpuBuffer target, GpuBuffer indices, GpuBuffer grad, int targetRows, int targetCols, int nIndices)
      Record scatter-add-rows (atomic): [OP_SCATTER_ADD_ROWS_ATOMIC, targetId, indicesId, gradId, targetRows, targetCols, nIdx]

    • recordSumAbs

      public void recordSumAbs(GpuBuffer in, GpuBuffer out, int rows, int cols)
      Record sum-abs row reduction: [OP_SUM_ABS, inId, outId, rows, cols]

    • recordCrossEntropyLoss

      public void recordCrossEntropyLoss(GpuBuffer logits, GpuBuffer targets, GpuBuffer out, int rows, int cols)
      Record cross-entropy row losses: [OP_CROSS_ENTROPY_LOSS, logitsId, targetsId, outId, rows, cols]

    • recordCrossEntropyGradient

      public void recordCrossEntropyGradient(GpuBuffer logits, GpuBuffer targets, GpuBuffer out, int rows, int cols)
      Record cross-entropy gradient: [OP_CROSS_ENTROPY_GRADIENT, logitsId, targetsId, outId, rows, cols]

    • recordSumScalar

      public void recordSumScalar(GpuBuffer in, GpuBuffer out, int rows, int cols)
      Record scalar sum reduction: [OP_SUM_SCALAR, inId, outId, rows, cols]

    • recordTranspose

      public void recordTranspose(GpuBuffer in, GpuBuffer out, int rows, int cols)
      Record transpose: [OP_TRANSPOSE, inId, outId, rows, cols]

    • recordRowBroadcast

      public void recordRowBroadcast(int opCode, GpuBuffer a, GpuBuffer rowVec, GpuBuffer out, int rows, int cols)
      Record row broadcast: [opCode, aId, rowVecId, outId, rows, cols]

    • recordColBroadcast

      public void recordColBroadcast(int opCode, GpuBuffer a, GpuBuffer colVec, GpuBuffer out, int rows, int cols)
      Record col broadcast: [opCode, aId, colVecId, outId, rows, cols]

    • recordReduction

      public void recordReduction(int opCode, GpuBuffer in, GpuBuffer out, int rows, int cols)
      Record row/col reduction: [opCode, inId, outId, rows, cols]

    • 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.