export function runProgram()

in tfjs-backend-webgl/src/gpgpu_math.ts [202:361]


export function runProgram<T extends Tensor, K extends Tensor>(
    gpgpu: GPGPUContext, binary: GPGPUBinary, inputs: TensorData[],
    output: TensorData, customUniformValues?: number[][]): void {
  if (!binary.program.enableShapeUniforms) {
    validateBinaryAndProgram(binary.inShapeInfos, inputs);
    validateBinaryAndProgram([binary.outShapeInfo], [output]);
  }

  const outTex = output.texData.texture;
  const outTexShape = output.texData.texShape;
  if (output.texData.isPacked) {
    gpgpu.setOutputPackedMatrixTexture(
        outTex.texture, outTexShape[0], outTexShape[1]);
  } else {
    gpgpu.setOutputMatrixTexture(
        outTex.texture, outTexShape[0], outTexShape[1]);
  }
  gpgpu.setProgram(binary.webGLProgram);

  // Set special uniforms (NAN, INFINITY)
  if (env().getNumber('WEBGL_VERSION') === 1) {
    if (binary.infLoc !== null) {
      gpgpu.gl.uniform1f(binary.infLoc, Infinity);
    }
  }
  if (binary.nanLoc !== null) {
    gpgpu.gl.uniform1f(binary.nanLoc, NaN);
  }

  // Set user-defined inputs
  inputs.forEach((input, i) => {
    const varName = binary.program.variableNames[i];
    const varLoc = binary.uniformLocations[varName];
    const varOffsetLoc = binary.uniformLocations[`offset${varName}`];
    const varShapeLoc = binary.inShapesLocations[`${varName}Shape`];
    const varTexShapeLoc = binary.inTexShapesLocations[`${varName}TexShape`];

    if (varShapeLoc) {
      const {uniformShape} = shader_compiler.getUniformInfoFromShape(
          binary.program.packedInputs, input.shape, input.texData.texShape);
      switch (uniformShape.length) {
        case 1:
          gpgpu.gl.uniform1iv(varShapeLoc, new Int32Array(uniformShape));
          break;
        case 2:
          gpgpu.gl.uniform2iv(varShapeLoc, new Int32Array(uniformShape));
          break;
        case 3:
          gpgpu.gl.uniform3iv(varShapeLoc, new Int32Array(uniformShape));
          break;
        case 4:
          gpgpu.gl.uniform4iv(varShapeLoc, new Int32Array(uniformShape));
          break;
        default:
          break;
      }
    }
    if (varTexShapeLoc) {
      gpgpu.gl.uniform2i(
          varTexShapeLoc, input.texData.texShape[0], input.texData.texShape[1]);
    }

    if (varLoc == null) {
      // The compiler inferred that this variable is not used in this shader.
      return;
    }

    if (input.isUniform) {
      // Upload the values of the tensor as uniform.
      if (util.sizeFromShape(input.shape) < 2) {
        gpgpu.gl.uniform1f(varLoc, input.uniformValues[0]);
      } else {
        let vals = input.uniformValues;
        if (!(vals instanceof Float32Array)) {
          vals = new Float32Array(vals);
        }
        gpgpu.gl.uniform1fv(varLoc, vals);
      }
      return;
    }

    // If the input was sliced, upload the flat offset index.
    if (input.texData.slice != null && varOffsetLoc != null) {
      gpgpu.gl.uniform1i(varOffsetLoc, input.texData.slice.flatOffset);
    }

    gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i);
  });

  const outShapeLoc = binary.outShapeLocation;
  if (outShapeLoc) {
    switch (output.shape.length) {
      case 1:
        gpgpu.gl.uniform1iv(outShapeLoc, new Int32Array(output.shape));
        break;
      case 2:
        gpgpu.gl.uniform2iv(outShapeLoc, new Int32Array(output.shape));
        break;
      case 3:
        gpgpu.gl.uniform3iv(outShapeLoc, new Int32Array(output.shape));
        break;
      case 4:
        gpgpu.gl.uniform4iv(outShapeLoc, new Int32Array(output.shape));
        break;
      default:
        break;
    }
  }
  if (binary.outShapeStridesLocation) {
    const strides = util.computeStrides(output.shape);
    switch (output.shape.length) {
      case 2:
        gpgpu.gl.uniform1iv(
            binary.outShapeStridesLocation, new Int32Array(strides));
        break;
      case 3:
        gpgpu.gl.uniform2iv(
            binary.outShapeStridesLocation, new Int32Array(strides));
        break;
      case 4:
        gpgpu.gl.uniform3iv(
            binary.outShapeStridesLocation, new Int32Array(strides));
        break;
      default:
        break;
    }
  }
  if (binary.outTexShapeLocation) {
    gpgpu.gl.uniform2i(
        binary.outTexShapeLocation, output.texData.texShape[0],
        output.texData.texShape[1]);
  }

  if (binary.program.customUniforms && customUniformValues) {
    binary.program.customUniforms.forEach((d, i) => {
      const customLoc = binary.customUniformLocations[i];
      const customValue = customUniformValues[i];
      if (d.type === 'float') {
        gpgpu.gl.uniform1fv(customLoc, customValue);
      } else if (d.type === 'vec2') {
        gpgpu.gl.uniform2fv(customLoc, customValue);
      } else if (d.type === 'vec3') {
        gpgpu.gl.uniform3fv(customLoc, customValue);
      } else if (d.type === 'vec4') {
        gpgpu.gl.uniform4fv(customLoc, customValue);
      } else if (d.type === 'int') {
        gpgpu.gl.uniform1iv(customLoc, customValue);
      } else if (d.type === 'ivec2') {
        gpgpu.gl.uniform2iv(customLoc, customValue);
      } else if (d.type === 'ivec3') {
        gpgpu.gl.uniform3iv(customLoc, customValue);
      } else if (d.type === 'ivec4') {
        gpgpu.gl.uniform4iv(customLoc, customValue);
      } else {
        throw Error(`uniform type ${d.type} is not supported yet.`);
      }
    });
  }
  gpgpu.executeProgram();
}