[WebGL] Fix Vao usage for parallel compile feature #7587

Linchenn · 2023-04-14T00:14:43Z

Problem

If we create VAO for each webgl program, then we have to set VAO for each program and bind it before execution. Right now, setting VAO happens in createProgram and createProgram is involved in compilation stage. However, setting VAO has blocking-call getAttribLocation, so setting VAO should be called after compilation stage.

This should work similar as what we do for getUniformLocations. We moved getUniformLocations call out of createProgram/compilation-stage, because it has blocking-call.

Evaluation

before #6913 is merged (tfjs 4.0.0)

checkCompileCompletion takes the majority time for warm up.

after #6913 is merged (tfjs 4.1.0)

checkCompileCompletion takes the trivial time for warm up.

getAttribLocation is blocking in the compilation stage.

after this PR

checkCompileCompletion takes the majority time for warm up.

codes

I used the following codes to reproduce the problem

<!DOCTYPE html>
<html>
  <body>

    // Local build
    <script src="https://unpkg.com/@tensorflow/tfjs-core@latest/dist/tf-core.js"></script>
    <script src="https://unpkg.com/@tensorflow/tfjs-converter@latest/dist/tf-converter.js"></script>
    <script src="./dist/bin/tfjs-backend-webgl/dist/tf-backend-webgl.js"></script>

    // After https://github.com/tensorflow/tfjs/pull/6913/files is merged
    <!-- <script src="https://unpkg.com/@tensorflow/[email protected]/dist/tf-core.js"></script>
    <script src="https://unpkg.com/@tensorflow/[email protected]/dist/tf-converter.js"></script>
    <script src="https://unpkg.com/@tensorflow/[email protected]/dist/tf-backend-webgl.js"></script> -->

    // Before https://github.com/tensorflow/tfjs/pull/6913/files is merged
    <!-- <script src="https://unpkg.com/@tensorflow/[email protected]/dist/tf-core.js"></script>
    <script src="https://unpkg.com/@tensorflow/[email protected]/dist/tf-converter.js"></script>
    <script src="https://unpkg.com/@tensorflow/[email protected]/dist/tf-backend-webgl.js"></script> -->


    <script>
      const url = `https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v3_small_075_224/classification/5/default/1`;
      model = tf.loadGraphModel(url, {fromTFHub: true});

      function normalWarmUp() {
        const input = tf.randomNormal([1, 224, 224, 3]);
        return model.predict(input);
      }

      function parallelWarmUp() {
        const input = tf.randomNormal([1, 224, 224, 3]);
        tf.env().set('ENGINE_COMPILE_ONLY', true);
        model.predict(input);
        tf.env().set('ENGINE_COMPILE_ONLY', false);
        tf.backend().checkCompileCompletion();
        tf.backend().getUniformLocations();
        return model.predict(input);
      }

      function cleanUp() {
        tf.backend().binaryCache = {};
      }

      async function benchmark(fn) {
        model = await model;
        const start = performance.now();
        fn().dataSync();
        return performance.now() - start;
      }

    </script>
  </body>
</html>

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

This change is

Linchenn · 2023-04-14T00:19:30Z

The problem is caused by #6913, which is trying to bind VAO before each webgl program execution. I think the idea makes sense, even though TFJS's vertex shaders are same, because users may use the same canvas and write their own webgl program with different vertex shaders.

Linchenn · 2023-04-14T00:23:58Z

tfjs-backend-webgl/src/backend_webgl.ts

@@ -1277,6 +1277,7 @@ export class MathBackendWebGL extends KernelBackend {

  getUniformLocations() {
    for (const binary of Object.values(this.binaryCache)) {
+      this.gpgpu.buildVao(binary.webGLProgram);


It does not make sense to add this code within getUniformLocations and theoratically we should move it to a new function, like 'setVao'. However, since we have told some users to use this to utilize parallel compilation feature, making setVao and requiring users to call it are breaking change for them.

Nit: Can you add a TODO here explaining this? Thanks!

Done. Thanks!

pyu10055

thanks for fixing the issue!

pyu10055 · 2023-04-17T16:45:52Z

tfjs-backend-webgl/src/gpgpu_math.ts

@@ -253,6 +254,7 @@ export function runProgram<T extends Tensor, K extends Tensor>(
        outTex.texture, outTexShape[0], outTexShape[1]);
  }
  gpgpu.setProgram(binary.webGLProgram);


should this two line be the same as just call gpgpu.buildVao(binary.webGLProgram)?

IICU, buildVao is to initialize VAO, while here it is using initialized VAO. Specifically, VAO initialization has to bind all attributes, while using initialized VAO does not need to.

mattsoulanille

LGTM with a nit.

mattsoulanille · 2023-04-17T21:55:11Z

tfjs-backend-webgl/src/backend_webgl.ts

@@ -1277,6 +1277,7 @@ export class MathBackendWebGL extends KernelBackend {

  getUniformLocations() {
    for (const binary of Object.values(this.binaryCache)) {
+      this.gpgpu.buildVao(binary.webGLProgram);


Nit: Can you add a TODO here explaining this? Thanks!

mattsoulanille

LGTM. Thanks for adding the TODO.

Fix parallel compile feature

c96e4dc

Linchenn requested a review from pyu10055 April 14, 2023 00:16

Merge branch 'master' into fixVao

d5611b2

Linchenn commented Apr 14, 2023

View reviewed changes

pyu10055 reviewed Apr 17, 2023

View reviewed changes

Merge branch 'master' into fixVao

06021c0

pyu10055 approved these changes Apr 17, 2023

View reviewed changes

Linchenn requested a review from mattsoulanille April 17, 2023 21:49

mattsoulanille approved these changes Apr 17, 2023

View reviewed changes

Linchenn added 3 commits April 17, 2023 14:59

Update backend_webgl.ts

ba3b41b

Merge branch 'fixVao' of https://github.com/Linchenn/tfjs into fixVao

1562d36

typo

264f77c

mattsoulanille approved these changes Apr 17, 2023

View reviewed changes

Linchenn merged commit 82421ec into tensorflow:master Apr 17, 2023

Linchenn deleted the fixVao branch April 17, 2023 23:16

Linchenn mentioned this pull request Apr 17, 2023

Parallel shader compilation is broken #7577

Closed

wingman-jr-addon mentioned this pull request May 5, 2023

Performance - Check TF.js 4.5.0 wingman-jr-addon/wingman_jr#190

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WebGL] Fix Vao usage for parallel compile feature #7587

[WebGL] Fix Vao usage for parallel compile feature #7587

Linchenn commented Apr 14, 2023 •

edited

Loading

Linchenn commented Apr 14, 2023

Linchenn Apr 14, 2023

mattsoulanille Apr 17, 2023

Linchenn Apr 17, 2023

pyu10055 left a comment

pyu10055 Apr 17, 2023

Linchenn Apr 17, 2023

mattsoulanille left a comment

mattsoulanille Apr 17, 2023

mattsoulanille left a comment

[WebGL] Fix Vao usage for parallel compile feature #7587

[WebGL] Fix Vao usage for parallel compile feature #7587

Conversation

Linchenn commented Apr 14, 2023 • edited Loading

Problem

Evaluation

before #6913 is merged (tfjs 4.0.0)

after #6913 is merged (tfjs 4.1.0)

after this PR

codes

Linchenn commented Apr 14, 2023

Linchenn Apr 14, 2023

Choose a reason for hiding this comment

mattsoulanille Apr 17, 2023

Choose a reason for hiding this comment

Linchenn Apr 17, 2023

Choose a reason for hiding this comment

pyu10055 left a comment

Choose a reason for hiding this comment

pyu10055 Apr 17, 2023

Choose a reason for hiding this comment

Linchenn Apr 17, 2023

Choose a reason for hiding this comment

mattsoulanille left a comment

Choose a reason for hiding this comment

mattsoulanille Apr 17, 2023

Choose a reason for hiding this comment

mattsoulanille left a comment

Choose a reason for hiding this comment

Linchenn commented Apr 14, 2023 •

edited

Loading