From 85c91de62ec677f4e5ddcfcf00aac00c8b6cf23a Mon Sep 17 00:00:00 2001
From: Yang Gu <yang.gu@intel.com>
Date: Thu, 25 Jan 2024 06:49:37 +0800
Subject: [PATCH] [js/webgpu] Fix issue of timestamp query (#19258)

When we enable webgpu profiling mode between session.create and
session.run, current implementation has a problem to create querySet
(and also queryResolveBuffer) if we share the commandEncoder with inputs
upload. This PR fixes this by moving the querySet creation to the place
we set queryType.
---
 web/lib/wasm/jsep/backend-webgpu.ts | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/web/lib/wasm/jsep/backend-webgpu.ts b/web/lib/wasm/jsep/backend-webgpu.ts
index afef7042a4280..8ca025d66550c 100644
--- a/web/lib/wasm/jsep/backend-webgpu.ts
+++ b/web/lib/wasm/jsep/backend-webgpu.ts
@@ -222,16 +222,6 @@ export class WebGpuBackend {
   getCommandEncoder(): GPUCommandEncoder {
     if (!this.commandEncoder) {
       this.commandEncoder = this.device.createCommandEncoder();
-
-      if (this.queryType !== 'none' && typeof this.querySet === 'undefined') {
-        this.querySet = this.device.createQuerySet({
-          type: 'timestamp',
-          count: this.maxDispatchNumber * 2,
-        });
-        this.queryResolveBuffer = this.device.createBuffer(
-            // eslint-disable-next-line no-bitwise
-            {size: this.maxDispatchNumber * 2 * 8, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.QUERY_RESOLVE});
-      }
     }
     return this.commandEncoder;
   }
@@ -654,6 +644,16 @@ export class WebGpuBackend {
       } else if (this.device.features.has('timestamp-query')) {
         this.queryType = 'at-passes';
       }
+
+      if (this.queryType !== 'none' && typeof this.querySet === 'undefined') {
+        this.querySet = this.device.createQuerySet({
+          type: 'timestamp',
+          count: this.maxDispatchNumber * 2,
+        });
+        this.queryResolveBuffer = this.device.createBuffer(
+            // eslint-disable-next-line no-bitwise
+            {size: this.maxDispatchNumber * 2 * 8, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.QUERY_RESOLVE});
+      }
     }
   }
   onRunStart(): void {