[webgl]Add functions for parallel compilation

[lina128]

This is an experimental feature.
In this experimental feature, we use a flag (compileOnly) and two APIs (checkCompileCompletion/Async and getUniformLocations) to achieve parallel compilation. The usage is demonstrated in the unit test.

This feature is experimental, it assumes only one model is running by the engine and there's no cached shaders, basically an engine for only one model. It works like this:

1. Compilation round: Just compile shaders in parallel. Then check shaders link status in batch and then set uniform locations.
2. Warm up upload and download round: Upload textures to GPU in batch.
3. Actual inference as normal.

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

---

This change is 

[lina128]

cc @qjia7

[pyu10055]

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

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1122 at r1 (raw file):

  async checkCompileCompletionAsync(): Promise<boolean[]> {
    const ps = [];

rename to promises

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1121 at r2 (raw file):

    } else {
      for (const [, binary] of Object.entries(this.binaryCache)) {
        const p: Promise<boolean> = new Promise((resolve) => {

you can combine the array push with the creation of the promise.

---

tfjs-backend-webgl/src/backend_webgl_test.ts, line 759 at r2 (raw file):

  // because it's a global flag, the async test will affect other tests.
  it('does not have memory leak.', () => {
    console.log('hello world');

remove?

---

tfjs-backend-webgl/src/backend_webgl_test.ts, line 778 at r2 (raw file):

    expectArraysClose(data, [2, 2, 2]);
    tf.dispose([a0, b0, c0]);
    tf.removeBackend(customWebGLBackendName);

should we verify that after removing the backend, the getBinaryCache() call return 0 or the new backend starts with 0 ?

---

tfjs-backend-webgl/src/gpgpu_math.ts, line 49 at r1 (raw file):

}

export interface GPGPUBinary {

should GPGPUBinary extends from GPGPUBinaryLocations interface?

[pyu10055]

Reviewed 2 of 6 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @lina128)

---

tfjs-backend-webgl/src/backend_webgl.ts, line 835 at r3 (raw file):

      // When in precompile phase, we want to precompile all the shaders, no
      // matter whether texture is uploaded or not.
      if (texData.texture != null || engine().state.compileOnly) {

it is not clear what is this condition for? the original code does check null for texture, if when compileOnly is true will need to execute the block, is it changing the behave for compileOnly is false?

---

tfjs-backend-webgl/src/backend_webgl_test.ts, line 798 at r3 (raw file):

    // Pre-compile round.
    tf.engine().state.compileOnly = true;

Should this be a core flag that can be set like tf.env().set('ENGINE_COMPILE_ONLY', true)?

---

tfjs-backend-webgl/src/gpgpu_math.ts, line 107 at r3 (raw file):

  if (!engine().state.compileOnly) {
    const {

you can use deconstrcution here?

const locations = getUniformLocations(gpgpu, program, webGLProgram);
return {program, fragmentShader, source, webGLProgram, ...locations};

[lina128]

Reviewable status: 0 of 1 approvals obtained (waiting on @lina128 and @pyu10055)

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1122 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

rename to promises

Done.

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1121 at r2 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

you can combine the array push with the creation of the promise.

Done.

---

tfjs-backend-webgl/src/backend_webgl.ts, line 835 at r3 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

it is not clear what is this condition for? the original code does check null for texture, if when compileOnly is true will need to execute the block, is it changing the behave for compileOnly is false?

This seems unnecessary, removed.

---

tfjs-backend-webgl/src/backend_webgl_test.ts, line 759 at r2 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

remove?

Done.

---

tfjs-backend-webgl/src/backend_webgl_test.ts, line 778 at r2 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

should we verify that after removing the backend, the getBinaryCache() call return 0 or the new backend starts with 0 ?

After removing the backend, we won't be able to call getBinaryCache() because the instance no longer exists. For the new backend, it's a new instance, it will start with 0.

---

tfjs-backend-webgl/src/gpgpu_math.ts, line 49 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

should GPGPUBinary extends from GPGPUBinaryLocations interface?

Ah, good idea, done.

---

tfjs-backend-webgl/src/gpgpu_math.ts, line 107 at r3 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

you can use deconstrcution here?

const locations = getUniformLocations(gpgpu, program, webGLProgram);
return {program, fragmentShader, source, webGLProgram, ...locations};

Done.

[lina128]

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

---

tfjs-backend-webgl/src/backend_webgl_test.ts, line 798 at r3 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

Should this be a core flag that can be set like tf.env().set('ENGINE_COMPILE_ONLY', true)?

Done.

[pyu10055]

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

---

tfjs-core/src/engine.ts, line 159 at r4 (raw file):

  };

  compileOnly = false;

is this variable still needed?

[lina128]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @pyu10055)

---

tfjs-core/src/engine.ts, line 159 at r4 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

is this variable still needed?

Good catch, removed.

[vladmandic]

Has this PR stalled? I was looking forward to this improvement a lot.