Method uploadToGPU

tfjs-backend-webgpu/src/backend_webgpu.ts:733–780

(dataId: DataId)

  }

  uploadToGPU(dataId: DataId): void {
    const tensorData = this.tensorMap.get(dataId);
    // Already on the GPU.
    if (tensorData.resource != null) {
      return;
    }

    const size = webgpu_util.GPUBytesPerElement(tensorData.dtype) *
        util.sizeFromShape(tensorData.shape);
    let buffer;
    const usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC |
        GPUBufferUsage.COPY_DST;
    if (tensorData.values) {
      buffer = this.bufferManager.acquireBuffer(size, usage, true);
      if (buffer.mapState === 'unmapped') {
        const stagingBuffer = this.bufferManager.acquireBuffer(
            size, GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC, true,
            false);
        const arrayBuffer = stagingBuffer.getMappedRange();
        if (tensorData.dtype === 'int32' || tensorData.dtype === 'bool') {
          new Int32Array(arrayBuffer).set(tensorData.values as TypedArray);
        } else {
          new Float32Array(arrayBuffer).set(tensorData.values as Float32Array);
        }
        stagingBuffer.unmap();
        this.ensureCommandEncoderReady();
        this.endComputePassEncoder();
        this.commandEncoder.copyBufferToBuffer(
            stagingBuffer, 0, buffer, 0, size);

        this.stagingPendingDisposal.push(stagingBuffer);
      } else {
        const arrayBuffer = buffer.getMappedRange();
        if (tensorData.dtype === 'int32' || tensorData.dtype === 'bool') {
          new Int32Array(arrayBuffer).set(tensorData.values as TypedArray);
        } else {
          new Float32Array(arrayBuffer).set(tensorData.values as Float32Array);
        }
        buffer.unmap();
      }

      // Once uploaded, don't store the values on cpu.
      tensorData.values = null;
    } else {
      buffer = this.bufferManager.acquireBuffer(size, usage);
    }
    tensorData.resource = buffer;
  }

  private makeUniforms(programUniform: ProgramUniform): GPUBindingResource {
    let currentOffset = 0;