Lazily pack and unpack textures in WebGL kernels

[annxingyuan]

Changes

• Add packedInput boolean property to GPGPUProgram to identify WebGL kernels that require packed inputs (batchMatMul, conv2DWithIm2Col).
• Such kernels no longer unpack their outputs, and they only pack their inputs if they are not already packed.
• In compileAndRun, check whether the GPGPUProgram requires packed inputs. If it does not, and an input texture is packed, unpack it before passing it to the program.

PERF

---

For repository owners only:

Please remember to apply all applicable tags to your pull request.
Tags: FEATURE, BREAKING, BUG, PERF, DEV, DOC, SECURITY

For more info see: /~https://github.com/tensorflow/tfjs/blob/master/DEVELOPMENT.md

---

This change is 

[nsthorat]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @nsthorat, and @dsmilkov)

---

src/kernels/backend_webgl.ts, line 1663 at r1 (raw file):

      }

      if (texData.usage === TextureUsage.PACK &&

maybe we should take PACKED / UNPACKED out of usage into its own field, since it's now kind of orthogonal to the usage.

---

src/kernels/backend_webgl.ts, line 1667 at r1 (raw file):

        const unpackProgram = new UnpackProgram(input.shape);
        input = this.compileAndRun(unpackProgram, [input]);
        texData = this.texData.get(input.dataId);

I think this is a leak -- you need to dispose the old data. Can you make sure there's a unit test that catches this?

---

src/kernels/backend_webgl.ts, line 1668 at r1 (raw file):

        input = this.compileAndRun(unpackProgram, [input]);
        texData = this.texData.get(input.dataId);
      }

shouldn't we also be doing the reverse here? If the shader wants packed inputs and the texture usage is not PACK then we pack? Then we don't have to do that all over the place in each kernel

---

src/kernels/webgl/gpgpu_math.ts, line 31 at r1 (raw file):

  outputShape: number[];
  userCode: string;
  packedInputs?: boolean;

probably should just something as simple as isPacked which is about whether the operation is packed, not that the inputs are packed

[annxingyuan]

Reviewable status: 0 of 1 approvals obtained (waiting on @dsmilkov)

---

src/kernels/backend_webgl.ts, line 1663 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

maybe we should take PACKED / UNPACKED out of usage into its own field, since it's now kind of orthogonal to the usage.

Done

---

src/kernels/backend_webgl.ts, line 1667 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

I think this is a leak -- you need to dispose the old data. Can you make sure there's a unit test that catches this?

Done

---

src/kernels/backend_webgl.ts, line 1668 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

shouldn't we also be doing the reverse here? If the shader wants packed inputs and the texture usage is not PACK then we pack? Then we don't have to do that all over the place in each kernel

Done

---

src/kernels/webgl/gpgpu_math.ts, line 31 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

probably should just something as simple as isPacked which is about whether the operation is packed, not that the inputs are packed

Done

[annxingyuan]

@nsthorat

If I dispose the input tensors to a kernel immediately after packing / unpacking them, gradient operations stop working because the inputs to engine.runKernel are disposed by the time they get passed to the tape node.

Maybe it would make sense to have the gradient scope dispose of the tensors instead? Do you have advice on how to do that in a straightforward way?

[nsthorat]

Reviewed 1 of 5 files at r1, 3 of 9 files at r2.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan and @dsmilkov)

---

src/kernels/webgl/tex_util.ts, line 49 at r2 (raw file):

  values: DataTypeMap[DataType];
  usage: TextureUsage;
  packed: boolean;

this should be isPacked too

---

src/kernels/webgl/texture_manager.ts, line 169 at r2 (raw file):

    shapeRowsCol: [number, number], physicalTexType: PhysicalTextureType,
    packed: boolean): string {
  return `${shapeRowsCol[0]}_${shapeRowsCol[1]}_${physicalTexType}${

you cam just make this _${packed} (note the underscore)

---

src/kernels/webgl/texture_manager.ts, line 149 at r3 (raw file):

    logicalTexType: TextureUsage, packed: boolean): PhysicalTextureType {
  if (packed) {
    return ENV.get('WEBGL_RENDER_FLOAT32_ENABLED') ?

do we ever have a packed bit being true for uploading / downloading? I assume that will be later?

---

src/kernels/webgl/webgl_util_test.ts, line 89 at r2 (raw file):

        [2, util.nearestLargerEven(tf.ENV.get('WEBGL_MAX_TEXTURE_SIZE') + 1)];
    const texShape =
        webgl_util.getTextureShapeFromLogicalShape(logicalShape, true);

pull true to a constant for readability

const isPacked = true;

[annxingyuan]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan and @dsmilkov)

---

src/kernels/webgl/tex_util.ts, line 49 at r2 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

this should be isPacked too

Done

---

src/kernels/webgl/texture_manager.ts, line 169 at r2 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

you cam just make this _${packed} (note the underscore)

Done

---

src/kernels/webgl/texture_manager.ts, line 149 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

do we ever have a packed bit being true for uploading / downloading? I assume that will be later?

Yes, I was planning to add it for this issue: tensorflow/tfjs#804

---

src/kernels/webgl/webgl_util_test.ts, line 89 at r2 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

pull true to a constant for readability

const isPacked = true;

Done

[dsmilkov]

Reviewed 2 of 9 files at r2, 1 of 2 files at r3, 6 of 6 files at r4.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan and @dsmilkov)

---

src/kernels/backend_webgl.ts, line 1652 at r4 (raw file):

      // TODO(/~https://github.com/tensorflow/tfjs/issues/821): Make it possible
      // for packed shaders to sample from uniforms.
      if (texData.texture == null && !(!texData.isPacked && program.isPacked) &&

this condition looks a bit suspicious since it's equivalent to texData.isPacked || !program.isPacked. Is that what we want?

---

src/kernels/backend_webgl.ts, line 1663 at r4 (raw file):

      }

      if (texData.isPacked !== !!program.isPacked) {

why the double !!

---

src/kernels/backend_webgl.ts, line 1668 at r4 (raw file):

        if (texData.isPacked) {
          preProcessProgram = new UnpackProgram(input.shape);
          processedInput = this.compileAndRun(preProcessProgram, [input]);

just a reminder: once you sync at master, use the private unpack() pack() methods added in previous PR

---

src/kernels/webgl/gpgpu_math.ts, line 31 at r4 (raw file):

  outputShape: number[];
  userCode: string;
  isPacked?: boolean;

let's rename this to "usesPackedTextures" be more explicit

[annxingyuan]

Note - this PR now depends on #1327 because the batchNorm tests need to be able to print 4D packed outputs