Migrate conv2D and related ops to logical sampling (#52)

* migrate addScaledMat and conv2d to logical sampling and improve shader compiler

* fix conv2d zero paddig and make the project build

* migrate rest of conv shaders to logical

* replace zero pad with if

* migrate pool ops

* address review comments

Author: dsmilkov

.vscode/settings.json

@@ -9,6 +9,7 @@
   },
   "tslint.enable": true,
   "tslint.run": "onType",
+  "tslint.configFile": "tslint.json",
   "editor.tabSize": 2,
   "editor.insertSpaces": true,
   "files.insertFinalNewline": true,

demos/benchmarks/conv_gpu_benchmark.ts

@@ -14,76 +14,53 @@ limitations under the License.
 ==============================================================================*/
 
 import * as conv_util from '../../src/math/conv_util';
-import * as conv_gpu from '../../src/math/webgl/conv_gpu';
+import {Array1D, Array3D, Array4D, initializeGPU} from '../../src/math/ndarray';
+import {Conv2DProgram} from '../../src/math/webgl/conv_gpu';
 import {GPGPUContext} from '../../src/math/webgl/gpgpu_context';
-import * as test_util from '../../src/test_util';
+import * as gpgpu_math from '../../src/math/webgl/gpgpu_math';
+import {TextureManager} from '../../src/math/webgl/texture_manager';
 
 import {BenchmarkTest} from './benchmark';
 
 const OP_RUNS = 40;
 
 export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
-  const inputShapeRCD: [number, number, number] = [size, size, 1];
+  const gpgpu = new GPGPUContext();
+  const texManager = new TextureManager(gpgpu);
+  initializeGPU(gpgpu, texManager);
+
+  const inputDepth = 1;
+  const inputShape: [number, number, number] = [size, size, inputDepth];
   const outputDepth = 1;
   const fieldSize = 11;
   const stride = 1;
-  const zeroPad = conv_util.computeDefaultPad(inputShapeRCD, fieldSize, stride);
-  const outputShapeRCD: [number, number, number] =
-      conv_util.computeOutputShape3D(
-          inputShapeRCD, fieldSize, outputDepth, stride, zeroPad);
-
-  const inputTexShapeRC = conv_util.computeTexShapeFrom3D(inputShapeRCD);
-  const outputTexShapeRC = conv_util.computeTexShapeFrom3D(outputShapeRCD);
-  const weightsTexShapeRC = conv_util.computeWeightsTexShape(
-      inputShapeRCD[2], outputDepth, fieldSize);
-  const biasesTexShapeRC = conv_util.computeBiasesTexShape(outputDepth);
+  const zeroPad = conv_util.computeDefaultPad(inputShape, fieldSize, stride);
 
   const hasBias = true;
-  const gpgpu = new GPGPUContext();
-  const program = gpgpu.createProgram(conv_gpu.getFragmentShaderSource(
-      inputShapeRCD, outputDepth, fieldSize, stride, zeroPad, hasBias));
-
-  const inputTexture =
-      gpgpu.createMatrixTexture(inputTexShapeRC[0], inputTexShapeRC[1]);
-  const weightsTexture =
-      gpgpu.createMatrixTexture(weightsTexShapeRC[0], weightsTexShapeRC[1]);
-  const biasesTexture =
-      gpgpu.createMatrixTexture(biasesTexShapeRC[0], biasesTexShapeRC[1]);
-  const outputTexture =
-      gpgpu.createMatrixTexture(outputTexShapeRC[0], outputTexShapeRC[1]);
-
-  const inputData = test_util.randomArrayInRange(
-      inputTexShapeRC[0] * inputTexShapeRC[1], -1, 1);
-  const weightsData = test_util.randomArrayInRange(
-      weightsTexShapeRC[0] * weightsTexShapeRC[1], -1, 1);
-  const biasesData = test_util.randomArrayInRange(
-      biasesTexShapeRC[0] * biasesTexShapeRC[1], -1, 1);
-
-  gpgpu.uploadMatrixToTexture(
-      inputTexture, inputTexShapeRC[0], inputTexShapeRC[1], inputData);
-  gpgpu.uploadMatrixToTexture(
-      weightsTexture, weightsTexShapeRC[0], weightsTexShapeRC[1], weightsData);
-  gpgpu.uploadMatrixToTexture(
-      biasesTexture, biasesTexShapeRC[0], biasesTexShapeRC[1], biasesData);
+  const program = new Conv2DProgram(
+      inputShape, fieldSize, outputDepth, stride, zeroPad, hasBias);
+  const outputShape = program.outputShape as [number, number, number];
+  const out = Array3D.zeros(outputShape);
+  const x = Array3D.randUniform(inputShape, -1, 1);
+  const wShape = conv_util.computeWeightsShape4D(1, outputDepth, fieldSize);
+  const W = Array4D.randUniform(wShape, -1, 1);
+  const b = Array1D.randUniform([outputDepth], -1, 1);
+  const inputs = [x, W, b];
+  const binary = gpgpu_math.compileProgram(gpgpu, program, inputs, out);
 
   const start = performance.now();
   for (let i = 0; i < OP_RUNS; i++) {
-    conv_gpu.convolve(
-        gpgpu, program, inputTexture, weightsTexture, biasesTexture,
-        outputTexture, outputTexShapeRC);
+    gpgpu_math.runProgram(binary, inputs, out);
   }
-
-  gpgpu.downloadMatrixFromTexture(
-      outputTexture, outputTexShapeRC[0], outputTexShapeRC[1]);
-  const end = performance.now();
-
-  const avgTime = (end - start) / OP_RUNS;
-
-  gpgpu.deleteMatrixTexture(inputTexture);
-  gpgpu.deleteMatrixTexture(weightsTexture);
-  gpgpu.deleteMatrixTexture(biasesTexture);
-  gpgpu.deleteMatrixTexture(outputTexture);
-  gpgpu.deleteProgram(program);
+  out.getValues();
+  const avgTime = (performance.now() - start) / OP_RUNS;
+
+  x.dispose();
+  W.dispose();
+  b.dispose();
+  out.dispose();
+  texManager.dispose();
+  gpgpu.deleteProgram(binary.webGLProgram);
   gpgpu.dispose();
 
   return avgTime;

demos/benchmarks/conv_transpose_gpu_benchmark.ts

@@ -14,71 +14,48 @@ limitations under the License.
 ==============================================================================*/
 
 import * as conv_util from '../../src/math/conv_util';
-import * as conv_backprop_gpu from '../../src/math/webgl/conv_backprop_gpu';
+import {Array3D, Array4D, initializeGPU} from '../../src/math/ndarray';
+import {Conv2DTransposeProgram} from '../../src/math/webgl/conv_backprop_gpu';
 import {GPGPUContext} from '../../src/math/webgl/gpgpu_context';
-import * as test_util from '../../src/test_util';
-
+import * as gpgpu_math from '../../src/math/webgl/gpgpu_math';
+import {TextureManager} from '../../src/math/webgl/texture_manager';
 import {BenchmarkTest} from './benchmark';
 
-const OP_RUNS = 100;
+const OP_RUNS = 40;
 
 export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
-  const xShapeRCD: [number, number, number] = [size, size, 1];
+  const origInputDepth = 1;
   const origOutputDepth = 2;
+  const xShape: [number, number, number] = [size, size, 1];
   const fieldSize = 11;
   const origStride = 1;
   const origPad = 1;
 
   const gpgpu = new GPGPUContext();
+  const texManager = new TextureManager(gpgpu);
+  initializeGPU(gpgpu, texManager);
   gpgpu.enableAutomaticDebugValidation(true);
-  const origInputDepth = xShapeRCD[2];
-  const src = conv_backprop_gpu.getFragmentShaderConvTransposeSource(
-      xShapeRCD, fieldSize, origInputDepth, origStride, origPad, false);
-  const program = gpgpu.createProgram(src);
-
-  // Upload x.
-  const xTexShapeRC = conv_util.computeTexShapeFrom3D(xShapeRCD);
-  const xTex = gpgpu.createMatrixTexture(xTexShapeRC[0], xTexShapeRC[1]);
-  const xData =
-      test_util.randomArrayInRange(xTexShapeRC[0] * xTexShapeRC[1], -1, 1);
-  gpgpu.uploadMatrixToTexture(xTex, xTexShapeRC[0], xTexShapeRC[1], xData);
 
-  // Upload weights.
-  const wTexShapeRC = conv_util.computeWeightsTexShape(
+  const hasBias = false;
+  const program = new Conv2DTransposeProgram(
+      xShape, fieldSize, origInputDepth, origStride, origPad, hasBias);
+  const outputShape = program.outputShape as [number, number, number];
+  const out = Array3D.zeros(outputShape);
+  const x = Array3D.randUniform(xShape, -1, 1);
+  const wShape = conv_util.computeWeightsShape4D(
       origInputDepth, origOutputDepth, fieldSize);
-  const wData =
-      test_util.randomArrayInRange(wTexShapeRC[0] * wTexShapeRC[1], -1, 1);
-  const wTex = gpgpu.createMatrixTexture(wTexShapeRC[0], wTexShapeRC[1]);
-  gpgpu.uploadMatrixToTexture(wTex, wTexShapeRC[0], wTexShapeRC[1], wData);
-
-  // Figure out the output shape by dilating the input.
-  const dilatedRC =
-      conv_util.computeDilatedRC([xShapeRCD[0], xShapeRCD[1]], origStride);
-  const pad = fieldSize - 1 - origPad;
-  const resultShapeRCD = conv_util.computeOutputShape3D(
-      [dilatedRC[0], dilatedRC[1], origOutputDepth], fieldSize, origInputDepth,
-      1, pad);
-
-  const resultTexRC = conv_util.computeTexShapeFrom3D(resultShapeRCD);
-  const resultTex = gpgpu.createMatrixTexture(resultTexRC[0], resultTexRC[1]);
-
+  const W = Array4D.randUniform(wShape, -1, 1);
+  const inputs = [x, W];
+  const binary = gpgpu_math.compileProgram(gpgpu, program, inputs, out);
   const start = performance.now();
   for (let i = 0; i < OP_RUNS; i++) {
-    conv_backprop_gpu.convTranspose(
-        gpgpu, program, xTex, wTex, null, resultTex, resultTexRC);
+    gpgpu_math.runProgram(binary, inputs, out);
   }
+  out.getValues();
+  const avgTime = (performance.now() - start) / OP_RUNS;
 
-  gpgpu.downloadMatrixFromTexture(resultTex, resultTexRC[0], resultTexRC[1]);
-
-  const end = performance.now();
-
-  const avgTime = (end - start) / OP_RUNS;
-
-  gpgpu.deleteMatrixTexture(resultTex);
-  gpgpu.deleteMatrixTexture(xTex);
-  gpgpu.deleteMatrixTexture(wTex);
-  gpgpu.deleteProgram(program);
+  texManager.dispose();
+  gpgpu.deleteProgram(binary.webGLProgram);
   gpgpu.dispose();
-
   return avgTime;
 };

demos/benchmarks/logsumexp_gpu_benchmark.ts

@@ -13,14 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+import {Array2D, initializeGPU, Scalar} from '../../src/math/ndarray';
 import {GPGPUContext} from '../../src/math/webgl/gpgpu_context';
-import {LogSumExpProgram} from '../../src/math/webgl/logsumexp_gpu';
 import * as gpgpu_math from '../../src/math/webgl/gpgpu_math';
-import {Scalar, Array2D, initializeGPU} from '../../src/math/ndarray';
+import {LogSumExpProgram} from '../../src/math/webgl/logsumexp_gpu';
 import {TextureManager} from '../../src/math/webgl/texture_manager';
+
 import {BenchmarkTest} from './benchmark';
 
-const OP_RUNS = 100;
+const OP_RUNS = 2;
 
 export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
   const gpgpu = new GPGPUContext();
@@ -35,7 +36,7 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
   for (let i = 0; i < OP_RUNS; i++) {
     gpgpu_math.runProgram(binary, [a], out);
   }
-
+  out.getValues();
   const avgTime = (performance.now() - start) / OP_RUNS;
   a.dispose();
   out.dispose();

demos/benchmarks/max_pool_backprop_gpu_benchmark.ts

@@ -14,69 +14,49 @@ limitations under the License.
 ==============================================================================*/
 
 import * as conv_util from '../../src/math/conv_util';
+import {Array3D, initializeGPU, NDArray} from '../../src/math/ndarray';
 import {GPGPUContext} from '../../src/math/webgl/gpgpu_context';
+import * as gpgpu_math from '../../src/math/webgl/gpgpu_math';
 // tslint:disable-next-line:max-line-length
-import * as max_pool_backprop_gpu from '../../src/math/webgl/max_pool_backprop_gpu';
-import * as test_util from '../../src/test_util';
-import * as util from '../../src/util';
-
+import {MaxPool2DBackpropProgram} from '../../src/math/webgl/max_pool_backprop_gpu';
+import {TextureManager} from '../../src/math/webgl/texture_manager';
 import {BenchmarkTest} from './benchmark';
 
-const OP_RUNS = 100;
+const OP_RUNS = 40;
 
 export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
-  const dyShapeRCD: [number, number, number] = [size, size, 1];
-  const outputDepth = 1;
-  const fieldSize = 11;
-  const stride = 1;
-  const zeroPad = conv_util.computeDefaultPad(dyShapeRCD, fieldSize, stride);
-  const outputShapeRCD: [number, number, number] =
-      conv_util.computeOutputShape3D(
-          dyShapeRCD, fieldSize, outputDepth, stride, zeroPad);
-
-  const dyTexShapeRC = conv_util.computeTexShapeFrom3D(dyShapeRCD);
-  const outputTexShapeRC = conv_util.computeTexShapeFrom3D(outputShapeRCD);
-
   const gpgpu = new GPGPUContext();
-  const program = gpgpu.createProgram(
-      max_pool_backprop_gpu.getFragmentShaderMaxPoolBackprop(
-          dyShapeRCD, fieldSize, stride, zeroPad));
-
-  const dyTexture = gpgpu.createMatrixTexture(dyTexShapeRC[0], dyTexShapeRC[1]);
-  const maxPositionsTexture =
-      gpgpu.createMatrixTexture(dyTexShapeRC[0], dyTexShapeRC[1]);
-  const outputTexture =
-      gpgpu.createMatrixTexture(outputTexShapeRC[0], outputTexShapeRC[1]);
+  const texManager = new TextureManager(gpgpu);
+  initializeGPU(gpgpu, texManager);
 
-  const dyData =
-      test_util.randomArrayInRange(dyTexShapeRC[0] * dyTexShapeRC[1], -1, 1);
-  const maxPositionsData = new Float32Array(util.sizeFromShape(dyShapeRCD));
-  for (let i = 0; i < maxPositionsData.length; i++) {
-    maxPositionsData[i] = Math.floor(Math.random() * fieldSize * fieldSize);
+  const outputDepth = 1;
+  const dyShape: [number, number, number] = [size, size, outputDepth];
+  const fSize = 11;
+  const stride = 1;
+  const zeroPad = conv_util.computeDefaultPad(dyShape, fSize, stride);
+  const program = new MaxPool2DBackpropProgram(dyShape, fSize, stride, zeroPad);
+  const res = NDArray.zeros(program.outputShape);
+  const dy = Array3D.randUniform(dyShape, -1, 1);
+  const positionsData = new Float32Array(dy.size);
+  for (let i = 0; i < positionsData.length; i++) {
+    positionsData[i] = Math.floor(Math.random() * fSize * fSize);
   }
-
-  gpgpu.uploadMatrixToTexture(
-      dyTexture, dyTexShapeRC[0], dyTexShapeRC[1], dyData);
-  gpgpu.uploadMatrixToTexture(
-      maxPositionsTexture, dyTexShapeRC[0], dyTexShapeRC[1], maxPositionsData);
+  const positions = Array3D.new(dyShape, positionsData);
+  const binary =
+      gpgpu_math.compileProgram(gpgpu, program, [dy, positions], res);
 
   const start = performance.now();
   for (let i = 0; i < OP_RUNS; i++) {
-    max_pool_backprop_gpu.maxPoolBackprop(
-        gpgpu, program, dyTexture, maxPositionsTexture, outputTexture,
-        outputTexShapeRC);
+    gpgpu_math.runProgram(binary, [dy, positions], res);
   }
-
-  gpgpu.downloadMatrixFromTexture(
-      outputTexture, outputTexShapeRC[0], outputTexShapeRC[1]);
-  const end = performance.now();
-
-  const avgTime = (end - start) / OP_RUNS;
-
-  gpgpu.deleteMatrixTexture(dyTexture);
-  gpgpu.deleteMatrixTexture(maxPositionsTexture);
-  gpgpu.deleteMatrixTexture(outputTexture);
-  gpgpu.deleteProgram(program);
+  res.getValues();
+  const avgTime = (performance.now() - start) / OP_RUNS;
+
+  dy.dispose();
+  positions.dispose();
+  res.dispose();
+  texManager.dispose();
+  gpgpu.deleteProgram(binary.webGLProgram);
   gpgpu.dispose();
 
   return avgTime;