encapsulateGPUGraph<T> function

Tensor<T> encapsulateGPUGraph<T>(List<Tensor> cpuDynamicInputs, List<GPUTensor> gpuDynamicInputs, List<Tensor> cpuStaticParams, List<GPUTensor> gpuStaticParams, GPUTensor gpuOutput, T initialOutputValue, Uint8List forwardTape, Uint8List backwardTape,)

Implementation

Tensor<T> encapsulateGPUGraph<T>(
    List<Tensor> cpuDynamicInputs,
    List<GPUTensor> gpuDynamicInputs,
    List<Tensor> cpuStaticParams,
    List<GPUTensor> gpuStaticParams,
    GPUTensor gpuOutput,
    T initialOutputValue,
    Uint8List forwardTape,
    Uint8List backwardTape) {
  for (int i = 0; i < cpuDynamicInputs.length; i = i + 1) {
    GPUEngine.load(gpuDynamicInputs[i].id, cpuDynamicInputs[i].dataPtr,
        gpuDynamicInputs[i].shape);
  }

  GPUEngine.run(forwardTape);

  Tensor<T> out = Tensor<T>(initialOutputValue);
  GPUEngine.retrieve(gpuOutput.id, out.dataPtr);

  List<Tensor> allCpuNodes = [...cpuDynamicInputs, ...cpuStaticParams];
  List<GPUTensor> allGpuNodes = [...gpuDynamicInputs, ...gpuStaticParams];

  out.creator = Node(allCpuNodes, () {
    // Step 1: Zero the GPU gradients
    CommandBuffer zeroTape = CommandBuffer();
    for (int i = 0; i < allGpuNodes.length; i = i + 1) {
      zeroTape.putInt(OP_ZERO_GRAD);
      zeroTape.putString('${allGpuNodes[i].id}_grad');
    }

    // Fixed: Use .bytes() instead of .buffer
    GPUEngine.run(zeroTape.bytes());

    // Step 2: Push the "starting" gradient for the output back to GPU
    GPUEngine.load('${gpuOutput.id}_grad', out.gradPtr, gpuOutput.shape);

    // Step 3: Run the GPU backward tape
    GPUEngine.run(backwardTape);

    // Step 4: Pull gradients back to CPU and add them
    for (int i = 0; i < allCpuNodes.length; i = i + 1) {
      int numElements = 1;
      List<int> sList = allGpuNodes[i].shape;
      for (int s = 0; s < sList.length; s = s + 1) {
        numElements = numElements * sList[s];
      }

      Pointer<Float> tempGradPtr = calloc<Float>(numElements);
      GPUEngine.retrieve('${allGpuNodes[i].id}_grad', tempGradPtr);
      Float32List tempGradView = tempGradPtr.asTypedList(numElements);

      for (int k = 0; k < numElements; k = k + 1) {
        allCpuNodes[i].grad[k] = allCpuNodes[i].grad[k] + tempGradView[k];
      }

      calloc.free(tempGradPtr);
    }
  }, opName: 'gpu_graph_block');

  return out;
}