Decode textures on the GPU.

[annxingyuan]

Changes

• Add DecodeMatrixProgram.
• Remove tex_util methods for decoding textures on the CPU.
• Add method for getting texture shape for densely packed RGBA texture.

To see the logs from the Cloud Build CI, please join either
our discussion
or announcement mailing list.

---

This change is 

[nsthorat]

Reviewed 11 of 11 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, @nsthorat, and @pyu10055)

---

src/backends/webgl/backend_webgl.ts, line 366 at r1 (raw file):

  }

  readSync(dataId: DataId): DataValues {

does this need to be updated?

---

src/backends/webgl/backend_webgl.ts, line 507 at r1 (raw file):

    const vals =
        this.gpgpu
            .downloadByteEncodedFloatMatrixFromOutputTexture(

does this just work right now?

---

src/backends/webgl/backend_webgl.ts, line 2308 at r1 (raw file):

  }

  private decode(dataId: DataId) {

return type here

---

src/backends/webgl/backend_webgl.ts, line 2317 at r1 (raw file):

    const tmpTarget = this.makeTensorHandle(shape, 'float32') as TensorHandle &
        {size: number};
    tmpTarget.size = sizeFromShape(shape);

wait why do you have to do this?

---

src/backends/webgl/backend_webgl.ts, line 2448 at r1 (raw file):

    if (!ENV.getBool('WEBGL_LAZILY_UNPACK') &&
        this.texData.get(output.dataId).isPacked && !program.isPackShader &&
        !program.isDecodeShader) {

I was just thinking -- maybe this is better as an argument to compileAndRun instead of a field on the program? it feels like the wrong place for the field, wdyt?

---

src/backends/webgl/decode_matrix_gpu.ts, line 28 at r1 (raw file):

  outputShape: [number, number, number];

  constructor(outputShape: [number, number, number], texShape: [

just curious are you using the autoformatter here?

[annxingyuan]

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

---

src/backends/webgl/backend_webgl.ts, line 366 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

does this need to be updated?

readSync ultimately calls getValuesFromTexture, which has been updated.

---

src/backends/webgl/backend_webgl.ts, line 507 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

does this just work right now?

Yes, because the output is already as dense as possible, since there is no packed version of EncodeFloatProgram and each output value requires all four texel channels to represent.

---

src/backends/webgl/backend_webgl.ts, line 2308 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

return type here

Done

---

src/backends/webgl/backend_webgl.ts, line 2317 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

wait why do you have to do this?

Sorry must have been for earlier debugging - removed.

---

src/backends/webgl/decode_matrix_gpu.ts, line 28 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

just curious are you using the autoformatter here?

Yes.

[nsthorat]

Reviewed 1 of 1 files at r2.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, @nsthorat, and @pyu10055)

---

src/backends/webgl/backend_webgl.ts, line 2317 at r1 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

Sorry must have been for earlier debugging - removed.

thx, update the & {size: number} type above too

---

src/backends/webgl/backend_webgl.ts, line 2448 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

I was just thinking -- maybe this is better as an argument to compileAndRun instead of a field on the program? it feels like the wrong place for the field, wdyt?

thoughts on this?

[dsmilkov]

Really nice cleanup! Since the motivation for this is performance, do you have some numbers/benchmarks you can add to the PR description?

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

---

src/backends/webgl/backend_webgl.ts, line 2448 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

thoughts on this?

Maybe alternatively check if instanceof DecodeShaderProgram

---

src/backends/webgl/backend_webgl.ts, line 2327 at r2 (raw file):

    this.texData.get(tmpTarget.dataId).texShape =
        denseTexShape.map(
            d => d * 2) as [number, number];  // To undo the effect of isPacked

should you check for isPacked in order to do this?

---

src/backends/webgl/decode_matrix_gpu.ts, line 49 at r2 (raw file):

        vec4 result = vec4(0.);

        for(int i=0; i<4; i++) {

add space between for and (

---

src/backends/webgl/gpgpu_math.ts, line 34 at r2 (raw file):

  usesPackedTextures?: boolean;
  isDecodeShader?: boolean;
  isPackShader?: boolean;  // This property is used to single out the packing

if we remove isDecoderShader (per Nikhil's comment), we should consider removing isPackedShader also.

[annxingyuan]

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

---

src/backends/webgl/backend_webgl.ts, line 2317 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

thx, update the & {size: number} type above too

How should I update it? I still need the & {size: number} annotation otherwise the output no longer conforms to the return type signature of compileAndRun

Or did you mean update it in some other way?

---

src/backends/webgl/backend_webgl.ts, line 2448 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Maybe alternatively check if instanceof DecodeShaderProgram

Oops - yes Nikhil that makes sense. Daniel would that create issues with the minified build?

---

src/backends/webgl/backend_webgl.ts, line 2327 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

should you check for isPacked in order to do this?

No because I force the texture to be packed (see line 2318): we want each texel in the output to have 4 values so the data can be straightforwardly copied from the texture during readback.

---

src/backends/webgl/decode_matrix_gpu.ts, line 49 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add space between for and (

Done

---

src/backends/webgl/gpgpu_math.ts, line 34 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

if we remove isDecoderShader (per Nikhil's comment), we should consider removing isPackedShader also.

Done