Packed reverse

[jgartman]

Adds a packed kernel for reverse:

shape	Packed (ms)	Unpacked (ms)
[100, 100]	5.814	5.542
[200, 200]	6.002	5.640
[500, 500]	7.109	16.257
[1000, 1000]	18.289	44.871

Benchmark code:

export class ReverseGPUBenchmark implements BenchmarkTest {
  async run(size: number): Promise<number> {
    tf.setBackend('webgl');
    tf.ENV.set('WEBGL_PACK', false);
    const t: tf.Tensor<tf.Rank.R2> = tf.randomNormal([size, size]);

    const benchmark = () => tf.reverse(t, [0]) as tf.Tensor;

    const time = await util.warmupAndBenchmarkGPU(benchmark);

    t.dispose();

    return time;
  }
}

Compiled shader code for reverse of tensor of shape = [2, 3, 4] along axis = 1:

     void main() {
          ivec3 rc = getOutputCoords();
          vec4 result = vec4(0.);
          result.r = getChannel(getX(rc.x,3 - rc.y - 1,rc.z), vec2(3 - rc.y - 1,rc.z));
          if(rc.z + 1 < 4){
            result.g = getChannel(getX(rc.x,3 - rc.y - 1,(rc.z + 1)), vec2(3 - rc.y - 1,(rc.z + 1)));
          }
          if(rc.y + 1 < 3) {
            result.b = getChannel(getX(rc.x,3 - (rc.y + 1) - 1,rc.z), vec2(3 - (rc.y + 1) - 1,rc.z));
            if(rc.z + 1 < 4) {
              result.a = getChannel(getX(rc.x,3 - (rc.y + 1) - 1,(rc.z + 1)), vec2(3 - (rc.y + 1) - 1,(rc.z + 1)));
            }
          }
          setOutput(result);
        }
      

---

This change is 

[annxingyuan]

Hey @jgartman - thank you so much for this! I left some mostly stylistic suggestions - otherwise this is good to go.

Could you also add a test for reverse on an array with odd number of rows and columns (final two dimensions)?

Thanks - we'll merge this soon!

[annxingyuan]

Indent these lines.

[jgartman]

Done

[annxingyuan]

LLC.

[jgartman]

Done

[annxingyuan]

I realize that we have this pattern elsewhere in the packed shaders, but could we instead test for adding 1 to the coordinate rather than mutating the coordinate? That way we don't have to remember to clean up after ourselves and decrement the coordinate. I think it'll be a little bit easier to read as well.

[jgartman]

Thanks @annxingyuan, I refactored the code a bit and updated the compiled shader code at the top of the PR as well. Let me know what you think.

[annxingyuan]

What do you think about restructuring this shader to match the unpacked shader so that all the logic for rank = 1 is nested under one if-statement? I wonder whether it might be easier to read that way - avoids having to test for rank = 1 in so many places. If we did it this way we should also be able to inline upperRow and lowerRow in the body of void main below.

[jgartman]

Done

[annxingyuan]

Thanks @jgartman !