Fix bug where gpgpu binary was holding onto stale inputs/output (#30)

dsmilkov · web-flow · commit 1c20d786d4b3 · 2017-08-12T09:30:40.000-04:00
* fix bug where gpgpu binary was holding onto stale inputs/output

* address comments
diff --git a/demos/benchmarks/logsumexp_gpu_benchmark.ts b/demos/benchmarks/logsumexp_gpu_benchmark.ts
@@ -33,7 +33,7 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
 
   const start = performance.now();
   for (let i = 0; i < OP_RUNS; i++) {
-    gpgpu_math.runProgram(binary);
+    gpgpu_math.runProgram(binary, [a], out);
   }
 
   const avgTime = (performance.now() - start) / OP_RUNS;
diff --git a/demos/benchmarks/mulmat_gpu_benchmark.ts b/demos/benchmarks/mulmat_gpu_benchmark.ts
@@ -47,7 +47,7 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
 
   const start = performance.now();
   for (let i = 0; i < OP_RUNS; i++) {
-    gpgpu_math.runProgram(binary);
+    gpgpu_math.runProgram(binary, [aArr, bArr], resArr);
   }
   gpgpu.downloadMatrixFromTexture(resultTexture, size, size);
   const avgTime = (performance.now() - start) / OP_RUNS;
diff --git a/src/math/math_gpu.ts b/src/math/math_gpu.ts
@@ -87,7 +87,7 @@ export class NDArrayMathGPU extends NDArrayMath {
   private gpgpu: GPGPUContext;
   private textureManager: TextureManager;
   private programCache: {[key: string]: WebGLProgram} = {};
-  private binaryCache: {[key: string]: GPGPUBinary<NDArray, NDArray>} = {};
+  private binaryCache: {[key: string]: GPGPUBinary} = {};
   private gpgpuCreatedLocally: boolean;
 
   constructor(gpgpu?: GPGPUContext, safeMode = true) {
@@ -984,14 +984,12 @@ export class NDArrayMathGPU extends NDArrayMath {
         newShapeRCD, {texture: resultTexture, textureShapeRC: resultTexShape});
   }
 
-  private getAndSaveBinary<K extends NDArray, T extends NDArray>(
-      key: string,
-      getBinary: () => GPGPUBinary<K,T>):
-      GPGPUBinary<K,T> {
+  private getAndSaveBinary(key: string, getBinary: () => GPGPUBinary):
+      GPGPUBinary {
     if (!(key in this.binaryCache)) {
       this.binaryCache[key] = getBinary();
     }
-    return this.binaryCache[key] as GPGPUBinary<K, T>;
+    return this.binaryCache[key];
   }
 
   private getAndSaveProgram(programKey: string, getShaderSource: () => string):
diff --git a/src/math/webgl/argmaxequals_gpu_test.ts b/src/math/webgl/argmaxequals_gpu_test.ts
@@ -29,7 +29,7 @@ function uploadArgMaxEqualsDownload(
   const out = Scalar.new(0);
   const program = new ArgMaxEqualsProgram(aArr.size, bArr.size);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [aArr, bArr], out);
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [aArr, bArr], out);
   const result = out.get();
   aArr.dispose();
   bArr.dispose();
diff --git a/src/math/webgl/argminmax_gpu_test.ts b/src/math/webgl/argminmax_gpu_test.ts
@@ -29,7 +29,7 @@ function uploadArgMinMaxDownload(
   const out = Scalar.new(0);
   const program = new ArgMinMaxProgram(arr.size, op);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [arr], out);
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [arr], out);
 
   const result = out.get();
   arr.dispose();
diff --git a/src/math/webgl/gpgpu_math.ts b/src/math/webgl/gpgpu_math.ts
@@ -17,6 +17,7 @@ import {NDArray} from '../ndarray';
 
 import {GPGPUContext} from './gpgpu_context';
 import * as shader_compiler from './shader_compiler';
+import {ShapeInfo} from './shader_compiler';
 import * as util from '../../util';
 
 export interface GPGPUProgram {
@@ -26,47 +27,47 @@ export interface GPGPUProgram {
   userCode: string;
 }
 
-export interface GPGPUBinary<T extends NDArray, K extends NDArray> {
+export interface GPGPUBinary {
   webGLProgram: WebGLProgram;
   program: GPGPUProgram;
   gpgpu: GPGPUContext;
   source: string;
-  inputs: T[];
-  output: K;
+  inShapeInfos: ShapeInfo[];
+  outShapeInfo: ShapeInfo;
 }
 
 export function compileProgram<T extends NDArray, K extends NDArray>(
     gpgpu: GPGPUContext, program: GPGPUProgram, inputs: T[],
-    output: K): GPGPUBinary<T,K> {
+    output: K): GPGPUBinary {
   const userCode = program.userCode;
-  const programInputs = program.variableNames.map((x, i) => {
-    const fullShape = {
-      shape: inputs[i].shape,
+  const inputInfos = program.variableNames.map((x, i) => {
+    const shapeInfo = {
+      logicalShape: inputs[i].shape,
       texShape: inputs[i].getTextureShapeRC()
     };
-    return {name: x, fullShape};
+    return {name: x, shapeInfo};
   });
-
-  const outFullShape = {
-    shape: output.shape,
+  const inShapeInfos = inputInfos.map(x => x.shapeInfo);
+  const outShapeInfo = {
+    logicalShape: output.shape,
     texShape: output.getTextureShapeRC()
   };
-  const source = shader_compiler.makeShader(programInputs, outFullShape,
+  const source = shader_compiler.makeShader(inputInfos, outShapeInfo,
       userCode);
   return {
     program,
     source,
     webGLProgram: gpgpu.createProgram(source),
     gpgpu,
-    inputs,
-    output
+    inShapeInfos,
+    outShapeInfo
   };
 }
 
-function validateBinaryAndProgram(aArrays: NDArray[], bArrays: NDArray[]) {
-  aArrays.forEach((a, i) => {
-    const shapeA = a.shape;
-    const texShapeA = a.getTextureShapeRC();
+function validateBinaryAndProgram(shapeInfos: ShapeInfo[], bArrays: NDArray[]) {
+  shapeInfos.forEach((s, i) => {
+    const shapeA = s.logicalShape;
+    const texShapeA = s.texShape;
     const shapeB = bArrays[i].shape;
     const texShapeB = bArrays[i].getTextureShapeRC();
 
@@ -82,17 +83,10 @@ function validateBinaryAndProgram(aArrays: NDArray[], bArrays: NDArray[]) {
 }
 
 export function runProgram<T extends NDArray, K extends NDArray>(
-    binary: GPGPUBinary<T,K>, inputs?: T[], output?: K): void {
-  if (inputs == null) {
-    inputs = binary.inputs;
-  } else {
-    validateBinaryAndProgram(binary.inputs, inputs);
-  }
-  if (output == null) {
-    output = binary.output;
-  } else {
-    validateBinaryAndProgram([binary.output], [output]);
-  }
+    binary: GPGPUBinary, inputs: T[], output: K): void {
+  validateBinaryAndProgram(binary.inShapeInfos, inputs);
+  validateBinaryAndProgram([binary.outShapeInfo], [output]);
+
   const outTex = output.getTexture();
   const outTexShape = output.getTextureShapeRC();
   const gpgpu = binary.gpgpu;
diff --git a/src/math/webgl/logsumexp_gpu_test.ts b/src/math/webgl/logsumexp_gpu_test.ts
@@ -66,7 +66,7 @@ export function uploadLogSumExpDownload(a: Float32Array, rows: number,
   const rScalar = Scalar.new(0);
   const program = new LogSumExpProgram(aArr.size);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [aArr], rScalar);
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [aArr], rScalar);
   const result = rScalar.get();
   textureManager.dispose();
   gpgpu.deleteProgram(binary.webGLProgram);
diff --git a/src/math/webgl/minmax_gpu_test.ts b/src/math/webgl/minmax_gpu_test.ts
@@ -28,7 +28,7 @@ function uploadMinMaxDownload(a: Float32Array, rows: number,
   const out = Scalar.new(0);
   const program = new MinMaxProgram(arr.size, op);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [arr], out);
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [arr], out);
   const result = out.get();
   arr.dispose();
   textureManager.dispose();
diff --git a/src/math/webgl/mulmat_gpu_test.ts b/src/math/webgl/mulmat_gpu_test.ts
@@ -279,8 +279,8 @@ describe('mulmat_gpu (multiple matrices)', () => {
     gpgpu.uploadMatrixToTexture(b, bShape[0], bShape[1], bData);
     gpgpu.uploadMatrixToTexture(c, cShape[0], cShape[1], cData);
 
-    gpgpu_math.runProgram(axbProgram);
-    gpgpu_math.runProgram(abxcProgram);
+    gpgpu_math.runProgram(axbProgram, [aArr, bArr], abArr);
+    gpgpu_math.runProgram(abxcProgram, [abArr, cArr], rArr);
     const result = gpgpu.downloadMatrixFromTexture(r, rShape[0], rShape[1]);
     const expected = test_util.cpuMultiplyMatrix(
         test_util.cpuMultiplyMatrix(aData, 4, 2, bData, 2, 12), 4, 12, cData,
@@ -362,7 +362,7 @@ export function uploadMultiplyMatrixDownload(
   gpgpu.uploadMatrixToTexture(aTexture, aNumRows, aNumCols, a);
   gpgpu.uploadMatrixToTexture(bTexture, bNumRows, bNumCols, b);
 
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [aArr, bArr], resArr);
 
   const result =
       gpgpu.downloadMatrixFromTexture(resultTexture, outNumRows, outNumCols);
diff --git a/src/math/webgl/reducesum_gpu_test.ts b/src/math/webgl/reducesum_gpu_test.ts
@@ -79,7 +79,7 @@ export function uploadReduceSumDownload(a: Float32Array, rows: number,
 
   const program = new ReduceSumProgram(arr.size);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [arr], out);
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [arr], out);
 
   const result = out.get();
   arr.dispose();
diff --git a/src/math/webgl/shader_compiler.ts b/src/math/webgl/shader_compiler.ts
@@ -15,64 +15,66 @@ limitations under the License.
 
 import * as util from '../../util';
 
-export type NDArrayShape = {
-  shape: number[],
+export type ShapeInfo = {
+  logicalShape: number[],
   texShape: [number, number];
 };
 
-export type Input = {
+export type InputInfo = {
   name: string,
-  fullShape: NDArrayShape
+  shapeInfo: ShapeInfo
 };
 
 export function makeShader(
-    inputs: Input[], output: NDArrayShape, userCode: string): string {
+    inputsInfo: InputInfo[], outputShape: ShapeInfo,
+    userCode: string): string {
   const inputPrefixSnippet =
-      inputs.map(x => `uniform sampler2D ${x.name};`).join('\n');
+      inputsInfo.map(x => `uniform sampler2D ${x.name};`).join('\n');
   const inputSamplingSnippet =
-      inputs.map(x => getInputSamplingSnippet(x, output)).join('\n');
-  const outTexShape = output.texShape;
+      inputsInfo.map(x => getInputSamplingSnippet(x, outputShape)).join('\n');
+  const outTexShape = outputShape.texShape;
   const outputSamplingSnippet =
-      getOutputSamplingSnippet(output.shape, outTexShape);
+      getOutputSamplingSnippet(outputShape.logicalShape, outTexShape);
   const source = [
     SHADER_PREFIX, inputPrefixSnippet, SAMPLE_1D_SNIPPET, SAMPLE_2D_SNIPPET,
     SAMPLE_3D_SNIPPET, inputSamplingSnippet, outputSamplingSnippet, userCode
   ].join('\n');
   return source;
 }
 
-function getInputSamplingSnippet(input: Input, output: NDArrayShape) {
-  const fullShape = input.fullShape;
-  const shape = fullShape.shape;
-  const texShape = fullShape.texShape;
-  const outTexShape = output.texShape;
+function getInputSamplingSnippet(inInfo: InputInfo, outShapeInfo: ShapeInfo) {
+  const shape = inInfo.shapeInfo.logicalShape;
+  const texShape = inInfo.shapeInfo.texShape;
+  const outTexShape = outShapeInfo.texShape;
 
   let res = '';
   switch (shape.length) {
     case 0:
-      res += getSamplerScalar(input.name);
+      res += getSamplerScalar(inInfo.name);
       break;
     case 1:
-      res += getSampler1D(input.name, texShape);
+      res += getSampler1D(inInfo.name, texShape);
       break;
     case 2:
-      res += getSampler2D(input.name, shape as [number, number], texShape);
+      res += getSampler2D(inInfo.name, shape as [number, number], texShape);
       break;
     case 3:
-      res += getSampler3D(input.name, shape as [number, number, number],
-          texShape);
+      res += getSampler3D(
+          inInfo.name, shape as [number, number, number], texShape);
       break;
     default:
-      throw new Error(`${fullShape.shape.length}-D input sampling is not yet ` +
-                      `supported`);
+      throw new Error(
+          `${shape.length}-D input sampling` +
+          ` is not yet supported`);
   }
   // If input and output have matching logical shapes, add
   // getTexNameAtOutCoord() method that samples the input texture using the
   // output coordinates.
-  if (util.arraysEqual(input.fullShape.shape, output.shape)) {
-    res += getSamplerAtOutputCoords(input.name, texShape, outTexShape);
+  if (util.arraysEqual(
+          inInfo.shapeInfo.logicalShape, outShapeInfo.logicalShape)) {
+    res += getSamplerAtOutputCoords(inInfo.name, texShape, outTexShape);
   }
-  res += getSamplerFlat(input.name, texShape);
+  res += getSamplerFlat(inInfo.name, texShape);
   return res;
 }
 
diff --git a/src/math/webgl/unaryop_gpu_test.ts b/src/math/webgl/unaryop_gpu_test.ts
@@ -26,7 +26,7 @@ export function uploadUnaryDownload(a: NDArray, op: UnaryOp): Float32Array {
   const out = Array2D.zerosLike(a);
   const program = new UnaryOpProgram(a.shape, op);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [a], out);
-  gpgpu_math.runProgram(binary);
+  gpgpu_math.runProgram(binary, [a], out);
   const result = out.getValues();
   textureManager.dispose();
   gpgpu.deleteProgram(binary.webGLProgram);

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {`
`33`	`33`
`34`	`34`	`const start = performance.now();`
`35`	`35`	`for (let i = 0; i < OP_RUNS; i++) {`
`36`		`- gpgpu_math.runProgram(binary);`
	`36`	`+ gpgpu_math.runProgram(binary, [a], out);`
`37`	`37`	`}`
`38`	`38`
`39`	`39`	`const avgTime = (performance.now() - start) / OP_RUNS;`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {`
`47`	`47`
`48`	`48`	`const start = performance.now();`
`49`	`49`	`for (let i = 0; i < OP_RUNS; i++) {`
`50`		`- gpgpu_math.runProgram(binary);`
	`50`	`+ gpgpu_math.runProgram(binary, [aArr, bArr], resArr);`
`51`	`51`	`}`
`52`	`52`	`gpgpu.downloadMatrixFromTexture(resultTexture, size, size);`
`53`	`53`	`const avgTime = (performance.now() - start) / OP_RUNS;`