Method runWebGPUProgram

tfjs-backend-webgpu/src/backend_webgpu.ts:850–899

(
      program: webgpu_program.WebGPUProgram, inputs: TensorInfo[],
      outputDtype: DataType, programDefinedUniform?: ProgramUniform,
      output?: TensorInfo)

  }

  public runWebGPUProgram(
      program: webgpu_program.WebGPUProgram, inputs: TensorInfo[],
      outputDtype: DataType, programDefinedUniform?: ProgramUniform,
      output?: TensorInfo): TensorInfo {
    if (!output) {
      output = this.makeTensorInfo(program.outputShape, outputDtype);
    }
    if (util.sizeFromShape(output.shape) === 0) {
      // Short-circuit the computation since the result is empty (has 0 in its
      // shape).
      this.tensorMap.get(output.dataId).values =
          util.getTypedArrayFromDType(output.dtype as 'float32', 0);
      return output;
    }
    this.uploadToGPU(output.dataId);
    program.dispatch = reshapeDispatch(this.device, program);

    const inputsData = inputs.map((input: TensorInfo, i: number) => {
      if (input.dtype === 'complex64') {
        throw new Error(
            `GPGPUProgram does not support complex64 input. For complex64 ` +
            `dtypes, please separate the program into real and imaginary ` +
            `parts.`);
      }
      this.uploadToGPU(input.dataId);

      return {
        // Returning dtype from tensorMap because it reflects dtype
        // of underlying buffer, rather than abstract dtype.
        dtype: this.tensorMap.get(input.dataId).dtype,
        shape: input.shape,
        name: program.variableNames[i]
      };
    });

    program.shaderKey =
        webgpu_program.makeShaderKey(program, inputsData, output);

    const parallelCompilation = env().getBool('WEBGPU_ENGINE_COMPILE_ONLY');
    if (!(program.shaderKey in this.pipelineCache)) {
      this.pipelineCache[program.shaderKey] = webgpu_program.compileProgram(
          this.device, program, inputsData, output, parallelCompilation);
    }
    program.pipeline = this.pipelineCache[program.shaderKey];

    if (!parallelCompilation) {
      this.recordAndSubmit(program, output, inputs, programDefinedUniform);
    }
    return output;
  }

  private recordAndSubmit(
      program: webgpu_program.WebGPUProgram, output: TensorInfo,