[webgpu] Add program label for shader (#5905)

* [webgpu] Add program label for shader This is used for tracing. * Add label for pipeline Co-authored-by: Jiajia Qin <[email protected]> Co-authored-by: Ping Yu <[email protected]>
tensorflow · Dec 2, 2021 · a2b5ccf · a2b5ccf
1 parent 5b33e5c
commit a2b5ccf
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tfjs-backend-webgpu/src/kernels/webgpu_program.ts b/tfjs-backend-webgpu/src/kernels/webgpu_program.ts
@@ -73,9 +73,13 @@ export const compileProgram =
 
       const source = shader_preprocessor.makeShader(
           inputsData, outputData, program, isFromPixel);
-      const module = device.createShaderModule({code: source});
-      const pipeline = device.createComputePipeline(
-          {layout: pipelineLayout, compute: {module, entryPoint: 'main'}});
+      const module = device.createShaderModule(
+          {code: source, label: program.constructor.name});
+      const pipeline = device.createComputePipeline({
+        layout: pipelineLayout,
+        compute: {module, entryPoint: 'main'},
+        label: program.constructor.name
+      });
 
       return pipeline;
     };