Package io.github.kirstenali.deepj.tensor

Class ComputeGraph

java.lang.Object
io.github.kirstenali.deepj.tensor.ComputeGraph
public final class ComputeGraph extends Object

  • Field Details

  • Constructor Details

    • ComputeGraph

      public ComputeGraph(GpuRuntime runtime)

  • Method Details

    • isEmpty

      public boolean isEmpty()

    • flush

      public void flush()

    • materialize

      public void materialize(Tensor t)

    • releaseAll

      public void releaseAll()

    • releaseTemporary

      public void releaseTemporary()

    • recordHeadCacheWrite

      public void recordHeadCacheWrite(GpuBuffer update, GpuBuffer cache, int heads, int offset)

    • recordHeadCacheWrite

      public void recordHeadCacheWrite(GpuBuffer update, GpuBuffer cache, int heads, int length, int capacity, int width, int offset)

    • recordHeadPrefixMatmul

      public void recordHeadPrefixMatmul(GpuBuffer left, GpuBuffer cache, GpuBuffer output, int heads, int length, boolean transposeCache)

    • recordHeadPrefixMatmul

      public void recordHeadPrefixMatmul(GpuBuffer left, GpuBuffer cache, GpuBuffer output, int heads, int length, boolean transposeCache, int rows, int columns, int cacheWidth)

    • recordCausalAttention

      public void recordCausalAttention(GpuBuffer query, GpuBuffer key, GpuBuffer value, GpuBuffer output, GpuBuffer normalization, int heads, float scale)

    • recordCausalAttention

      public void recordCausalAttention(GpuBuffer query, GpuBuffer key, GpuBuffer value, GpuBuffer output, GpuBuffer normalization, int queryHeads, int keyValueHeads, float scale)

    • recordCausalAttentionBackward

      public void recordCausalAttentionBackward(GpuBuffer query, GpuBuffer key, GpuBuffer value, GpuBuffer output, GpuBuffer gradient, GpuBuffer normalization, GpuBuffer delta, GpuBuffer dq, GpuBuffer dk, GpuBuffer dv, int heads, float scale)

    • recordCausalAttentionBackward

      public void recordCausalAttentionBackward(GpuBuffer query, GpuBuffer key, GpuBuffer value, GpuBuffer output, GpuBuffer gradient, GpuBuffer normalization, GpuBuffer delta, GpuBuffer dq, GpuBuffer dk, GpuBuffer dv, int queryHeads, int keyValueHeads, float scale)

    • recordCrossEntropy

      public void recordCrossEntropy(GpuBuffer logits, GpuBuffer targets, GpuBuffer losses, GpuBuffer gradient, int rows, int cols)

    • recordRmsNorm

      public void recordRmsNorm(GpuBuffer input, GpuBuffer gamma, GpuBuffer output, GpuBuffer normalized, GpuBuffer rms, int rows, int cols, float epsilon)

    • recordRmsNormBackward

      public void recordRmsNormBackward(GpuBuffer gradient, GpuBuffer normalized, GpuBuffer rms, GpuBuffer gamma, GpuBuffer output, int rows, int cols)

    • recordSwiGlu

      public void recordSwiGlu(GpuBuffer gate, GpuBuffer up, GpuBuffer fused)

    • recordSwiGluBackward

      public void recordSwiGluBackward(GpuBuffer gradient, GpuBuffer gate, GpuBuffer up, GpuBuffer gateGradient, GpuBuffer upGradient)

    • recordSumScalar

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

    • recordTranspose

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

    • recordRowBroadcast

      public void recordRowBroadcast(int opCode, GpuBuffer a, GpuBuffer rowVec, GpuBuffer out, int rows, int cols)

    • recordColBroadcast

      public void recordColBroadcast(int opCode, GpuBuffer a, GpuBuffer colVec, GpuBuffer out, int rows, int cols)

    • recordReduction

      public void recordReduction(int opCode, GpuBuffer in, GpuBuffer out, int rows, int cols)

    • recordSoftmaxRows

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

    • recordSoftmaxBackward

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

    • recordScaledSoftmaxBackward

      public void recordScaledSoftmaxBackward(GpuBuffer gradient, GpuBuffer probabilities, GpuBuffer output, float scale)

    • recordLayerNormBackward

      public void recordLayerNormBackward(GpuBuffer dXHat, GpuBuffer xHat, GpuBuffer std, GpuBuffer out, int rows, int 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)

    • recordStorageCopy

      public void recordStorageCopy(GpuBuffer input, GpuBuffer output, int source, int target, int count)

    • recordBinary

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

    • recordMatmul

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

    • recordBatchedMatmul

      public void recordBatchedMatmul(GpuBuffer left, GpuBuffer right, GpuBuffer out, int batches, int leftRows, int leftCols, int rightRows, int rightCols, boolean transposeLeft, boolean transposeRight)

    • recordHeadPermutation

      public void recordHeadPermutation(int opCode, GpuBuffer input, GpuBuffer output, int sequenceLength, int heads, int headDim, int modelWidth)

    • recordCausalMask

      public void recordCausalMask(GpuBuffer input, GpuBuffer output, int rows, int sequenceLength)

    • recordCausalSoftmax

      public void recordCausalSoftmax(GpuBuffer input, GpuBuffer output, int rows, int sequenceLength, float scale)

    • recordGatherRows

      public void recordGatherRows(GpuBuffer input, GpuBuffer indices, GpuBuffer output, int inputRows, int columns, int outputRows)

    • recordRotary

      public void recordRotary(GpuBuffer input, GpuBuffer cosine, GpuBuffer sine, GpuBuffer output, int rows, int sequenceLength, int headDim, boolean inverse)

    • recordUnary

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

    • recordScalarUnary

      public void recordScalarUnary(int opCode, GpuBuffer in, GpuBuffer out, float scalar)

    • recordPow

      public void recordPow(GpuBuffer in, GpuBuffer out, float exponent)

    • recordClamp

      public void recordClamp(GpuBuffer in, GpuBuffer out, float min, float max)

    • recordScatterAddRows

      public void recordScatterAddRows(GpuBuffer target, GpuBuffer indices, GpuBuffer grad, int targetRows, int targetCols, int nIndices)

    • recordScatterAddRowsAtomic

      public void recordScatterAddRowsAtomic(GpuBuffer target, GpuBuffer indices, GpuBuffer grad, int targetRows, int targetCols, int nIndices)

    • recordSumAbs

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

    • recordSumSquares

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

    • recordSumSquaresScalar

      public void recordSumSquaresScalar(GpuBuffer input, GpuBuffer total)

    • recordCrossEntropyLoss

      public void recordCrossEntropyLoss(GpuBuffer logits, GpuBuffer targets, GpuBuffer out, int rows, int cols)

    • recordCrossEntropyGradient

      public void recordCrossEntropyGradient(GpuBuffer logits, GpuBuffer targets, GpuBuffer out, int rows, int cols)

    • layoutCausal

      public Tensor layoutCausal(Tensor input, int length, float scale, int mode)

    • layoutRotary

      public Tensor layoutRotary(Tensor input, Tensor cosine, Tensor sine, int length, boolean inverse)

    • layoutCrossEntropy

      public CrossEntropyResult layoutCrossEntropy(Tensor logits, int[] targets, boolean[] mask, int mode)

    • layoutAdamW

      public void layoutAdamW(Tensor weight, Tensor gradient, Tensor first, Tensor second, float lr, float beta1, float beta2, float epsilon, float decay, float bc1, float bc2)

    • layoutSwiGluBackward

      public SwiGluBackwardResult layoutSwiGluBackward(Tensor gradient, Tensor gate, Tensor up)

    • layoutNormalize

      public NormalizationResult layoutNormalize(Tensor input, Tensor gamma, Tensor beta, float epsilon, boolean centered)

    • layoutNormalizationBackward

      public Tensor layoutNormalizationBackward(Tensor gradient, Tensor normalized, Tensor denominator, Tensor gamma, boolean centered)

    • layoutIndexSelect

      public Tensor layoutIndexSelect(Tensor input, int axis, int[] indices)

    • layoutIndexAdd

      public void layoutIndexAdd(Tensor target, int axis, int[] indices, Tensor updates)

    • layoutIndexGradient

      public Tensor layoutIndexGradient(Tensor gradient, int[] shape, int axis, int[] indices)

    • layoutSliceGradient

      public Tensor layoutSliceGradient(Tensor gradient, int[] dimensions, int axis, int start, int step)

    • layoutMatmul

      public Tensor layoutMatmul(Tensor left, Tensor right)

    • layoutMaximumGradient

      public Tensor layoutMaximumGradient(Tensor input, Tensor maximum, Tensor gradient, int[] axes, boolean keep)

    • layoutReduction

      public Tensor layoutReduction(Tensor input, int[] axes, boolean keep, int operation, int correction)

    • layoutTotal

      public Tensor layoutTotal(Tensor input, int operation)

    • layoutSoftmax

      public Tensor layoutSoftmax(Tensor input, int axis)

    • recordMultiplyScalar

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

    • layoutBinary

      public Tensor layoutBinary(Tensor left, Tensor right, int operation)

    • layoutUnary

      public Tensor layoutUnary(Tensor input, int operation, float first, float second)

    • assignToDevice

      public void assignToDevice(Tensor target, Tensor source)

    • layoutBinaryInto

      public void layoutBinaryInto(Tensor target, Tensor operand, int operation)

    • layoutUnaryInto

      public void layoutUnaryInto(Tensor target, int operation, float scalar)

    • transferLayout

      public Tensor transferLayout(Tensor input, DType dtype)

    • zeroToDevice

      public Tensor zeroToDevice(DType dtype, int... shape)

    • lazyZeros

      public Tensor lazyZeros(DType dtype, int... shape)
      Zeros that stay unallocated until first use; a device zero-fill is queued only when a kernel consumes them.

    • castToDevice

      public Tensor castToDevice(Tensor input, DType dtype)

    • ensureGpuBuffer

      public GpuBuffer ensureGpuBuffer(Tensor t)

    • ensureStorageBuffer

      public GpuBuffer ensureStorageBuffer(Tensor t)
      Physical shared storage only; logical ND kernels must supply their own layouts.

    • newOutputBuffer

      public GpuBuffer newOutputBuffer(int rows, int cols)

    • newOutputBuffer

      public GpuBuffer newOutputBuffer(int rows, int cols, DType dtype)

    • newZeroBuffer

      public GpuBuffer newZeroBuffer(int rows, int cols)
      Queues explicit zero initialization before any operations consume this buffer.

    • newZeroBuffer

      public GpuBuffer newZeroBuffer(int rows, int cols, DType dtype)

    • createOutputTensor

      public Tensor createOutputTensor(GpuBuffer buf)

    • createOutputTensor

      public Tensor createOutputTensor(GpuBuffer buf, int... shape)

    • createOutputLike

      public Tensor createOutputLike(GpuBuffer buf, Tensor source)
      Shares only immutable canonical geometry, never storage, gradients or caller-owned arrays.

    • copyToDevice

      public Tensor copyToDevice(Tensor source)

    • markUpdated

      public void markUpdated(Tensor tensor)

    • bindTensorToBuffer

      public void bindTensorToBuffer(Tensor t, GpuBuffer buf)