matMulBiasReluGPU function

GPUTensor<Matrix> matMulBiasReluGPU(GPUTensor<Matrix> x, GPUTensor<Matrix> w, GPUTensor<Vector> b, CommandBuffer tape, List<GPUTensor> intermediates, { GPUTensor<Matrix>? reluOutTensor, GPUTensor<Matrix>? preReluOutTensor, })

///////////////////////////////// Fused Kernels (1000+) /// ///////////////////////////////// Executes a fused matrix multiplication, bias addition, and ReLU activation sequence. Multiplies matrix x by weight matrix w, broadcasts and adds the bias vector b, and applies a ReLU activation in a single optimized pass. Returns the activated output matrix and pushes the pre-activation tensor to the intermediates list to manage VRAM lifecycles. The operations are appended to tape for execution.

Implementation

/// Executes a fused matrix multiplication, bias addition, and ReLU activation sequence.
/// Multiplies matrix [x] by weight matrix [w], broadcasts and adds the bias vector [b],
/// and applies a ReLU activation in a single optimized pass. Returns the activated output matrix
/// and pushes the pre-activation tensor to the [intermediates] list to manage VRAM lifecycles.
/// The operations are appended to [tape] for execution.
GPUTensor<Matrix> matMulBiasReluGPU(GPUTensor<Matrix> x, GPUTensor<Matrix> w,
    GPUTensor<Vector> b, CommandBuffer tape, List<GPUTensor> intermediates,
    {GPUTensor<Matrix>? reluOutTensor, GPUTensor<Matrix>? preReluOutTensor}) {
  int M = x.shape[0];
  int K = x.shape[1];
  int N = w.shape[1];

  // Instantly reserve VRAM without building 33MB of Dart Lists!
  List<int> outShape = <int>[M, N];
  GPUTensor<Matrix> reluOut =
      reluOutTensor ?? GPUTensor<Matrix>.empty(outShape);
  GPUTensor<Matrix> preReluOut =
      preReluOutTensor ?? GPUTensor<Matrix>.empty(outShape);

  reluOut.subMap["preReluOut"] = preReluOut;
  // Add to trash so it doesn't leak VRAM!
  intermediates.add(preReluOut);

  tape.putInt(OP_MATMUL_BIAS_RELU_FORWARD);
  tape.putString(x.id);
  tape.putString(w.id);
  tape.putString(b.id);
  tape.putString(reluOut.id);
  tape.putString(preReluOut.id);

  int cost = (2 * M * K * N) + (M * N) + (M * N);

  reluOut.creator = GPUNode(
    <GPUTensor>[x, w, b],
    (CommandBuffer bTape) {
      bTape.putInt(OP_RELU_BACKWARD);
      bTape.putString(preReluOut.id);
      bTape.putString(reluOut.gradId);
      bTape.putString(preReluOut.gradId);

      bTape.putInt(OP_MATMUL);
      bTape.putString(x.id);
      bTape.putString(preReluOut.gradId);
      bTape.putString(w.gradId);
      bTape.putBool(true);
      bTape.putBool(false);
      bTape.putFloat(1.0);
      bTape.putFloat(1.0);
      bTape.putBool(true); // Tensor Cores ON

      bTape.putInt(OP_MATMUL);
      bTape.putString(preReluOut.gradId);
      bTape.putString(w.id);
      bTape.putString(x.gradId);
      bTape.putBool(false);
      bTape.putBool(true);
      bTape.putFloat(1.0);
      bTape.putFloat(1.0);
      bTape.putBool(true); // Tensor Cores ON

      bTape.putInt(OP_SUM_REDUCE_COLUMNS);
      bTape.putString(preReluOut.gradId);
      bTape.putString(b.gradId);
    },
    opName: 'matMulBiasReluGPU',
    cost: cost,
  );

  return reluOut;
}