[webgpu] Support WEBGPU_PRINT_SHADER #7523

axinging · 2023-03-28T07:35:58Z

To see the printed shader, turn on 'Verbose' in console.
Below is a test page(example url: index.html?WEBGPU_PRINT_SHADER=all, index.html?WEBGPU_PRINT_SHADER=binary, index.html?WEBGPU_PRINT_SHADER=binary,depth):

async function testWebGPUPrintShader() {
  tf.env().set('WEBGPU_CPU_FORWARD', false);
  await tf.setBackend('webgpu');
  await tf.ready();
  const re = getURLState(location.search);
  tf.env().set('WEBGPU_PRINT_SHADER', re);
  console.log(tf.env().get('WEBGPU_PRINT_SHADER'));
  // depthwise, matches 'depth'.
  {
    const fSize = 2;
    const pad = 'valid';
    const stride = 1;
    const chMul = 1;
    const inDepth = 1;

    const x = tf.tensor4d(
        [
          0.230664, 0.987388, 0.0685208, 0.419224, 0.887861, 0.731641,
          0.0741907, 0.409265, 0.351377
        ],
        [1, 3, 3, inDepth]);
    const w = tf.tensor4d(
        [0.303873, 0.229223, 0.144333, 0.803373],
        [fSize, fSize, inDepth, chMul],
    );

    const result = tf.depthwiseConv2d(x, w, stride, pad);
  }

  // add(sub,mul), matches 'binary'(Full binary list: https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgpu/src/binary_op_util.ts).
  {
    const a = tf.tensor2d([1, 2], [1, 2]);
    const b = tf.tensor2d([1, 2], [1, 2]);
    const c = tf.add(a, b);
  }

  // maxPool, matches 'pool'.
  {
    const x = tf.tensor3d([1, 2, 3, 4, 5, 6, 7, 9, 8], [3, 3, 1]);

    const result = tf.maxPool(x, 2, 1, 0);
  }
}

function getURLState(url) {
  let params = new URLSearchParams(url);
  const keys = [...params.keys()];
  if (keys.length === 0) return '';
  let printShaderString = '';
  if (params.has('WEBGPU_PRINT_SHADER')) {
    printShaderString = params.get('WEBGPU_PRINT_SHADER');
  }
  return printShaderString;
}

Bug: #7516

This change is

gyagp · 2023-03-29T01:35:16Z

e2e/benchmarks/local-benchmark/index.js

@@ -56,6 +58,7 @@ const TUNABLE_FLAG_NAME_MAP = {
 if (tf.engine().backendNames().includes('webgpu')) {
  TUNABLE_FLAG_NAME_MAP['WEBGPU_DEFERRED_SUBMIT_BATCH_SIZE'] =
      'deferred submit batch size';
+  TUNABLE_FLAG_NAME_MAP['WEBGPU_PRINT_SHADER'] = 'Print shader(s)';


Just “Print shader” to be consistent.

gyagp · 2023-03-29T01:38:54Z

tfjs-backend-webgpu/src/backend_webgpu.ts

      pipeline = webgpu_program.compileProgram(
-          this.device, program, inputsData, output);
+          this.device, program, inputsData, output, printShader);


To add another parameter "printShader" here couldn't save the checks. We may just put the check in compileProgram().

gyagp · 2023-03-29T01:43:12Z

tfjs-backend-webgpu/src/webgpu_program.ts

@@ -65,6 +65,10 @@ export const compileProgram =
        layout: 'auto'
      });

+      if (printShader) {
+        console.log(program.constructor.name);


Maybe we can put some decoration around the name so that it can be easily found.

gyagp · 2023-03-29T01:44:54Z

tfjs-backend-webgpu/src/flags_webgpu.ts

+/**
+ * Whether print shader.
+ */
+ENV.registerFlag('WEBGPU_PRINT_SHADER', () => false);


Do you want to wait for the string flag to be landed?

Yes, this change depends on #7509

axinging · 2023-03-30T01:47:28Z

@gyagp @qjia7 @xhcao PTAL

qjia7 · 2023-03-30T02:31:42Z

tfjs-backend-webgpu/src/webgpu_program.ts

+
+      const printShaderArray = printShaderString.split(',');
+      if (printShaderString === 'all' ||
+          printShaderArray.some(item => program.shaderKey.includes(item))) {


Do you need to convert program.shaderKey to lower case?

done, thanks

qjia7 · 2023-03-30T02:34:37Z

tfjs-backend-webgpu/src/webgpu_program.ts

+        throw new Error(
+            `WEBGPU_PRINT_SHADER doesn't support value ` +
+            `'${printShaderString}'. ` +
+            `It should be 'all' or program name`);


How can the user know the supported program name?

I changed this to shader key.

gyagp

Sorry for the delay!
The changes in WebGPU backend looks good. But I don't like the changes related to untunableFlagConfig.

gyagp · 2023-03-31T07:50:18Z

tfjs-backend-webgpu/src/webgpu_program.ts

@@ -54,7 +54,7 @@ export interface WebGPUProgram {

 export const compileProgram =
    (device: GPUDevice, program: WebGPUProgram, inputsData: InputInfo[],
-     output: TensorInfo): GPUComputePipeline => {
+     output: TensorInfo, key: string): GPUComputePipeline => {


Please rename this to shaderKey

gyagp · 2023-03-31T07:52:33Z

tfjs-backend-webgpu/src/backend_webgpu.ts

@@ -818,7 +818,7 @@ export class WebGPUBackend extends KernelBackend {
      pipeline = this.pipelineCache[key];
    } else {
      pipeline = webgpu_program.compileProgram(
-          this.device, program, inputsData, output);
+          this.device, program, inputsData, output, key);


Please rename this as shaderKey

gyagp · 2023-04-06T03:12:42Z

e2e/benchmarks/local-benchmark/index.html

@@ -113,6 +114,9 @@ <h2>TensorFlow.js Model Benchmark</h2>
      if (params.has('task')) {
        task = params.get('task');
      }
+      if (params.has('WEBGPU_PRINT_SHADER')) {
+        untunableFlagsConfig['WEBGPU_PRINT_SHADER'] = params.get('WEBGPU_PRINT_SHADER');


Can we just tf.env().setFlags() here and avoid all the changes related to untunableFlagsConfig?

No, this line https://github.com/tensorflow/tfjs/pull/7523/files#diff-622e3e249bd7bb5327a80ac8034614cdec1761a725466edce4dd57439d734fb8R756 will reset all env flags.

axinging

Thanks @gyagp , PTAL

axinging · 2023-04-07T01:49:27Z

e2e/benchmarks/local-benchmark/index.html

@@ -113,6 +114,9 @@ <h2>TensorFlow.js Model Benchmark</h2>
      if (params.has('task')) {
        task = params.get('task');
      }
+      if (params.has('WEBGPU_PRINT_SHADER')) {
+        untunableFlagsConfig['WEBGPU_PRINT_SHADER'] = params.get('WEBGPU_PRINT_SHADER');


No, this line https://github.com/tensorflow/tfjs/pull/7523/files#diff-622e3e249bd7bb5327a80ac8034614cdec1761a725466edce4dd57439d734fb8R756 will reset all env flags.

axinging · 2023-04-07T01:49:35Z

tfjs-backend-webgpu/src/backend_webgpu.ts

@@ -818,7 +818,7 @@ export class WebGPUBackend extends KernelBackend {
      pipeline = this.pipelineCache[key];
    } else {
      pipeline = webgpu_program.compileProgram(
-          this.device, program, inputsData, output);
+          this.device, program, inputsData, output, key);


axinging · 2023-04-07T01:49:42Z

tfjs-backend-webgpu/src/webgpu_program.ts

@@ -54,7 +54,7 @@ export interface WebGPUProgram {

 export const compileProgram =
    (device: GPUDevice, program: WebGPUProgram, inputsData: InputInfo[],
-     output: TensorInfo): GPUComputePipeline => {
+     output: TensorInfo, key: string): GPUComputePipeline => {


qjia7 · 2023-04-07T09:27:46Z

@axinging Can you remove the e2e related changes? The webgpu part LGTM. Then we can merge this PR.

Offline synced with Yang, for e2e changes, we would like to support any env flag as the url flag. It will be very convenient for us if want to test any env flag which is not in the current tunable flags list. In this way, we don't have to add new flag one by one to e2e , such as WEBGPU_PRINT_SHADER. For implementation, we can always treat unrecognized url flags as env flags and cache them like you did in this PR. You may need to carefully process the env flags if they are both in url flags and tunable flags list.

Bug: tensorflow#7516

gyagp

LGTM. Thanks for splitting this!

axinging changed the title ~~Support WEBGPU_PRINT_SHADER~~ [webgpu] Support WEBGPU_PRINT_SHADER Mar 28, 2023

gyagp reviewed Mar 29, 2023

View reviewed changes

axinging force-pushed the print_shader branch 2 times, most recently from cf99132 to a61a2c7 Compare March 30, 2023 01:34

axinging marked this pull request as ready for review March 30, 2023 01:44

axinging force-pushed the print_shader branch from a61a2c7 to 8016b8f Compare March 30, 2023 01:59

qjia7 reviewed Mar 30, 2023

View reviewed changes

axinging force-pushed the print_shader branch 4 times, most recently from 63e1ace to 4eb15f5 Compare March 31, 2023 07:45

gyagp reviewed Apr 6, 2023

View reviewed changes

axinging force-pushed the print_shader branch from 4eb15f5 to 86682c8 Compare April 6, 2023 06:28

axinging commented Apr 7, 2023

View reviewed changes

axinging added 2 commits April 10, 2023 09:57

Support WEBGPU_PRINT_SHADER

bc65bee

Bug: tensorflow#7516

Remove e2e change

534f4da

axinging force-pushed the print_shader branch from 8e5c822 to 534f4da Compare April 10, 2023 01:58

gyagp approved these changes Apr 10, 2023

View reviewed changes

qjia7 approved these changes Apr 10, 2023

View reviewed changes

gyagp merged commit 8db53b0 into tensorflow:master Apr 10, 2023

axinging mentioned this pull request Apr 10, 2023

Integrate TFJS kernels into TVM webatintel/tvm-web#4

Open

gaikwadrahul8 mentioned this pull request Apr 25, 2023

Support print shader on webgpu backend #7516

Closed

axinging deleted the print_shader branch May 9, 2023 05:19

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[webgpu] Support WEBGPU_PRINT_SHADER #7523

[webgpu] Support WEBGPU_PRINT_SHADER #7523

axinging commented Mar 28, 2023 •

edited

Loading

gyagp Mar 29, 2023

axinging Mar 29, 2023

gyagp Mar 29, 2023

axinging Mar 29, 2023

gyagp Mar 29, 2023

axinging Mar 29, 2023

gyagp Mar 29, 2023

axinging Mar 29, 2023

axinging commented Mar 30, 2023

qjia7 Mar 30, 2023

axinging Mar 31, 2023

qjia7 Mar 30, 2023

axinging Mar 31, 2023

gyagp left a comment

gyagp Mar 31, 2023

axinging Apr 7, 2023

gyagp Mar 31, 2023

axinging Apr 7, 2023

gyagp Apr 6, 2023

axinging Apr 7, 2023

axinging left a comment

axinging Apr 7, 2023

axinging Apr 7, 2023

axinging Apr 7, 2023

qjia7 commented Apr 7, 2023

gyagp left a comment

[webgpu] Support WEBGPU_PRINT_SHADER #7523

[webgpu] Support WEBGPU_PRINT_SHADER #7523

Conversation

axinging commented Mar 28, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

axinging commented Mar 30, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gyagp left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

axinging left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qjia7 commented Apr 7, 2023

gyagp left a comment

Choose a reason for hiding this comment

axinging commented Mar 28, 2023 •

edited

Loading