Optimize addN op by the multiple upload of textures

src/kernels/webgl/addn_gpu.ts

@@ -0,0 +1,51 @@
+/**
+ * @license
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GPGPUProgram} from './gpgpu_math';
+
+export class AddNProgram implements GPGPUProgram {
+  variableNames: string[];
+  outputShape: number[] = [];
+  userCode: string;
+
+  constructor(outputShape: number[], shapes: number[][]) {
+    this.outputShape = outputShape;
+    this.variableNames = shapes.map((_, i) => `T${i}`);
+
+    const snippets: string[] = [];
+    // Get target elements from every input tensor.
+    this.variableNames.forEach(variable => {
+      snippets.push(
+        `float v${variable} = get${variable}AtOutCoords();`
+      );
+    });
+
+    // Calculate the sum of all elements.
+    const operation = this.variableNames.map(variable => {
+      return `v${variable}`;
+    }).join(' + ');
+
+    this.userCode = `
+      void main() {
+        ${snippets.join('\n        ')}
+
+        float result = ${operation};
+        setOutput(result);
+      }
+    `;
+  }
+}

src/kernels/webgl/addn_packed_gpu.ts

@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GPGPUProgram} from './gpgpu_math';
+
+export class AddNPackedProgram implements GPGPUProgram {
+  variableNames: string[];
+  outputShape: number[] = [];
+  userCode: string;
+  usesPackedTextures = true;
+
+  constructor(outputShape: number[], shapes: number[][]) {
+    this.outputShape = outputShape;
+    this.variableNames = shapes.map((_, i) => `T${i}`);
+
+    const snippets: string[] = [];
+    // Get target elements from every input tensor.
+    this.variableNames.forEach(variable => {
+      snippets.push(
+        `vec4 v${variable} = get${variable}AtOutCoords();`
+      );
+    });
+
+    // Calculate the sum of all elements.
+    const operation = this.variableNames.map(variable => {
+      return `v${variable}`;
+    }).join(' + ');
+
+    this.userCode = `
+      void main() {
+        ${snippets.join('\n        ')}
+
+        vec4 result = ${operation};
+        setOutput(result);
+      }
+    `;
+  }
+}

src/kernels/webgl/backend_webgl.ts

@@ -119,7 +119,9 @@ import {UnaryOpProgram} from './unaryop_gpu';
 import * as unary_packed_op from './unaryop_packed_gpu';
 import {UnaryOpPackedProgram} from './unaryop_packed_gpu';
 import {UnpackProgram} from './unpack_gpu';
+import {AddNProgram} from './addn_gpu';
 import * as webgl_util from './webgl_util';
+import {AddNPackedProgram} from './addn_packed_gpu';
 
 type KernelInfo = {
   name: string; query: Promise<number>;
@@ -1468,11 +1470,31 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   addN<T extends Tensor>(tensors: T[]): T {
-    let res = tensors[0];
-    for (let i = 1; i < tensors.length; i++) {
-      res = this.add(res, tensors[i]) as T;
+    if (tensors.length === 1) {
+      return tensors[0];
     }
-    return res;
+
+    // Limit the number of uploaded textures for optimization.
+    if (tensors.length > ENV.get('WEBGL_MAX_TEXTURES_IN_SHADER')) {
+      const midIndex = Math.floor(tensors.length / 2);
+      const leftSide = this.addN(tensors.slice(0, midIndex));
+      const rightSide = this.addN(tensors.slice(midIndex));
+      return this.addN([leftSide, rightSide]);
+    }
+
+    const dtype = tensors
+        .map(t => t.dtype)
+        .reduce((d1, d2) => upcastType(d1, d2));
+    const shapes = tensors.map(t => t.shape);
+    // We can make sure shapes are identical in op level.
+    const usePackedOp = ENV.getBool('WEBGL_PACK_BINARY_OPERATIONS');
+    const program = usePackedOp ?
+        new AddNPackedProgram(tensors[0].shape, shapes) :
+        new AddNProgram(tensors[0].shape, shapes);
+    const output = usePackedOp ?
+        this.makePackedTensor(program.outputShape, dtype) as T :
+        this.makeOutputArray(program.outputShape, dtype) as T;
+    return this.compileAndRun<T>(program, tensors, output);
   }
 
   subtract(a: Tensor, b: Tensor): Tensor {