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();
}