Packed reverse (tensorflow#1617)

PERF

Author: jgartman

src/kernels/backend_webgl.ts

@@ -96,6 +96,7 @@ import {ResizeBilinearPackedProgram} from './webgl/resize_bilinear_packed_gpu';
 import {ResizeNearestNeigborBackpropProgram} from './webgl/resize_nearest_neighbor_backprop_gpu';
 import {ResizeNearestNeighborProgram} from './webgl/resize_nearest_neighbor_gpu';
 import {ReverseProgram} from './webgl/reverse_gpu';
+import {ReversePackedProgram} from './webgl/reverse_packed_gpu';
 import {ScatterProgram} from './webgl/scatter_gpu';
 import {SegmentOpProgram} from './webgl/segment_gpu';
 import {SelectProgram} from './webgl/select_gpu';
@@ -713,7 +714,9 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   reverse<T extends Tensor>(x: T, axis: number[]): T {
-    const program = new ReverseProgram(x.shape, axis);
+    const program = ENV.get('WEBGL_PACK_ARRAY_OPERATIONS') ?
+        new ReversePackedProgram(x.shape, axis) :
+        new ReverseProgram(x.shape, axis);
     return this.compileAndRun(program, [x]);
   }
 

src/kernels/webgl/reverse_packed_gpu.ts

@@ -0,0 +1,107 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC 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 {getChannels} from '../packing_util';
+import {GPGPUProgram} from './gpgpu_math';
+import {getCoordsDataType} from './shader_compiler';
+
+export class ReversePackedProgram implements GPGPUProgram {
+  variableNames = ['x'];
+  outputShape: number[];
+  userCode: string;
+  usesPackedTextures = true;
+
+  constructor(xShape: number[], axis: number[]) {
+    const rank = xShape.length;
+    if (rank > 4) {
+      throw new Error(
+          `WebGL backend: Reverse of rank-${rank} tensor is not yet supported`);
+    }
+    this.outputShape = xShape;
+    const channels = getChannels('rc', rank);
+    const nextColumn =
+        `${channels[rank - 1]} + 1 < ${this.outputShape[rank - 1]}`;
+    const nextRow = `${channels[rank - 2]} + 1 < ${this.outputShape[rank - 2]}`;
+    const type = getCoordsDataType(rank);
+    if (rank === 1) {
+      this.userCode = `
+        void main(){
+          int rc = getOutputCoords();
+          vec4 result = vec4(0.);
+          result.r = getChannel(getX(${xShape[0]} - rc - 1), rc);
+          if(${nextColumn}){
+              result.g = getChannel(getX(${xShape[0]} - (rc  + 1) - 1), rc + 1);
+          }
+          setOutput(result);
+        }
+      `;
+    } else {
+      this.userCode = `
+        void main() {
+          ${type} rc = getOutputCoords();
+          vec4 result = vec4(0.);
+          result.r = ${getR(channels.slice())};
+          if(${nextColumn}){
+            result.g = ${getG(channels.slice())};
+          }
+          if(${nextRow}) {
+            result.b = ${getB(channels.slice())};
+            if(${nextColumn}) {
+              result.a = ${getA(channels.slice())};
+            }
+          }
+          setOutput(result);
+        }
+    `;
+    }
+
+    function getR(channels: string[]): string {
+      return getChannel(channels);
+    }
+
+    function getG(channels: string[]): string {
+      channels[rank - 1] = '(' + channels[rank - 1] + ` + 1)`;
+      return getChannel(channels);
+    }
+
+    function getB(channels: string[]): string {
+      channels[rank - 2] = '(' + channels[rank - 2] + ` + 1)`;
+      return getChannel(channels);
+    }
+
+    function getA(channels: string[]): string {
+      channels[rank - 1] = '(' + channels[rank - 1] + ` + 1)`;
+      channels[rank - 2] = '(' + channels[rank - 2] + ` + 1)`;
+      return getChannel(channels);
+    }
+
+    function getChannel(channels: string[]): string {
+      const inCoordsArray = xShape.map((_, i) => getInCoord(i, channels));
+      const inCoords = inCoordsArray.join(',');
+      const innerDims = inCoordsArray.slice(-2).join(',');
+      return `getChannel(getX(${inCoords}), vec2(${innerDims}))`;
+    }
+
+    function getInCoord(i: number, channels1: string[]): string {
+      if (axis.indexOf(i) !== -1 && xShape[i] !== 1) {
+        return `${xShape[i]} - ${channels1[i]} - 1`;
+      } else {
+        return `${channels1[i]}`;
+      }
+    }
+  }
+}

src/ops/reverse_test.ts

@@ -71,6 +71,17 @@ describeWithFlags('reverse2d', ALL_ENVS, () => {
     expectArraysClose(result, [3, 2, 1, 6, 5, 4]);
   });
 
+  it('reverse a 2D array odd rows and columns at axis [0, 1]', () => {
+    const axis = [0, 1];
+    const a = tf.tensor2d(
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], [3, 5]);
+    const result = tf.reverse2d(a, axis);
+
+    expect(result.shape).toEqual(a.shape);
+    expectArraysClose(
+        result, [15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]);
+  });
+
   it('throws error with invalid input', () => {
     // tslint:disable-next-line:no-any
     const x: any = tf.tensor1d([1, 20, 300, 4]);