LayersModel#predict() results in all zeros when using WebGPU backend in Deno

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Monterey 12.6
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
• TensorFlow.js installed from (npm or script link): https://cdn.skypack.dev/@tensorflow/tfjs
• TensorFlow.js version (use command below): 3.20.0
• Browser version: Deno (built from source at commit 2929ec9f)
• Tensorflow.js Converter Version: N/A

Describe the current behavior

I came from https://github.com/tensorflow/tfjs/issues/6746 and made https://github.com/denoland/deno/pull/15853 to have deno compatible with the WebGPU backend. Model predictions sometimes resulting in an all-zeros tensor, the problem does not exist with a CPU backend.

Describe the expected behavior

Output tensors should contain non-zero numbers.

Standalone code to reproduce the issue

import * as tf from "https://cdn.skypack.dev/@tensorflow/tfjs?dts";
import "https://cdn.skypack.dev/@tensorflow/tfjs-backend-webgpu";

await tf.ready();
const model = tf.sequential();
model.add(tf.layers.dense({ units: 1, inputShape: [1] });
const output = await model.predict(tf.tensor([1])).array();

console.log(output); // prints [[0]]

Other info / logs I am on an Apple M1 laptop.

Backstory and tracking issues:

1. Deno uses deno_webgpu from gfx-rs/wgpu, which in turn uses gfx-rs/naga for all shader related things.
2. gfx-rs/wgpu#2994
3. gfx-rs/naga#2072
4. gfx-rs/naga#2071
5. gfx-rs/naga#1829
6. Turns out Naga is lexical scoping instead of functional/modulo scoping and they are changing that in gfx-rs/naga#2065

[haoyunfeix]

Interpreters with function scoping should be compatible with lexical scoping, moving fn main upwards should increase compatibility without apparent down sides.

What do you think? @gyagp @haoyunfeix

@vicary Done with https://github.com/tensorflow/tfjs/pull/6918, for shader compiler issues 1 and 2 mentioned above we intend not to fix in TFJS but track them on naga project(https://github.com/gfx-rs/naga/issues/2071 and https://github.com/gfx-rs/naga/issues/2080).

BTW, I indeed try to fix 1 and 2 in TFJS(https://github.com/tensorflow/tfjs/compare/master...haoyunfeix:tfjs:test_6842?expand=1) and glad to see WebGPU on deno could get the same result as CPU backend. I posted all resources(updated webgpu build and test code) here in case you are interested in.

upload.zip

[haoyunfeix]

@haoyunfeix would you mind sharing your temporary fix while you are preparing the upstream issue?

Sure. My temporary fix as below, to move main implantation into _start entry point: Before:

        fn _start(@builtin(local_invocation_id) LocalId : vec3<u32>,
                  @builtin(global_invocation_id) GlobalId : vec3<u32>,
                  @builtin(num_workgroups) NumWorkgroups : vec3<u32>) {
          localId = LocalId;
          globalId = GlobalId;
          numWorkgroups = NumWorkgroups;
          main(getGlobalIndex());
        }

        fn main(index : i32)
       {
          // Fill in the shared memory buffer.
          let localIndex = i32(localId.x);
          if(localIndex < 2) {
            sharedBuf[localIndex] = f32(B[localIndex]);
          }
          workgroupBarrier();

          if(index < uniforms.size) {
            let coords = getCoordsFromIndex(index);
            let a = getAByOutputIndex(index);
          let b = sharedBuf[coords[1]];
            setOutputAtIndex(index, binaryOperation(a, b));
          }
        }

After:

        fn _start(@builtin(local_invocation_id) LocalId : vec3<u32>,
                  @builtin(global_invocation_id) GlobalId : vec3<u32>,
                  @builtin(num_workgroups) NumWorkgroups : vec3<u32>) {
          localId = LocalId;
          globalId = GlobalId;
          numWorkgroups = NumWorkgroups;
          var index = getGlobalIndex();
          let localIndex = i32(localId.x);
          if(localIndex < 2) {
            sharedBuf[localIndex] = f32(B[localIndex]);
          }
          workgroupBarrier();

          if(index < uniforms.size) {
            let coords = getCoordsFromIndex(index);
            let a = getAByOutputIndex(index);
          let b = sharedBuf[coords[1]];
            setOutputAtIndex(index, binaryOperation(a, b));
          }
        }

You could see that the other _start outside function like getCoordsFromIndex and getAByOutputIndex can compile and works well, so maybe we shouldn’t call a function named main? I’ll take a double check.