From 428e55b98050cd5840def33758e8471f045fd641 Mon Sep 17 00:00:00 2001
From: Yang Gu <yang.gu@intel.com>
Date: Mon, 29 May 2023 15:38:24 +0800
Subject: [PATCH] webgpu: Fix timestamp query (#7723)

If a pass that needs timestamp query follows a pass without timestamp
query, querySet may not be created as expected. This PR fixes this
issue.
---
 tfjs-backend-webgpu/src/backend_webgpu.ts | 44 ++++++++++++-----------
 1 file changed, 23 insertions(+), 21 deletions(-)

diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts
index f4f89130d37..b21c9be1030 100644
--- a/tfjs-backend-webgpu/src/backend_webgpu.ts
+++ b/tfjs-backend-webgpu/src/backend_webgpu.ts
@@ -950,28 +950,30 @@ export class WebGPUBackend extends KernelBackend {
     const shouldTimeProgram = this.activeTimers != null;
     this.ensureCommandEncoderReady();
 
-    if (!this.computePassEncoder) {
-      const computePassDescriptor: GPUComputePassDescriptor = {};
-      if (shouldTimeProgram && this.supportTimestampQuery) {
-        if (this.querySet == null) {
-          this.querySet = this.device.createQuerySet({
-            type: 'timestamp',
-            count: this.querySetCount,
-          });
-        }
-        computePassDescriptor.timestampWrites = [
-          {
-            querySet: this.querySet,
-            queryIndex: 0,
-            location: 'beginning',
-          },
-          {
-            querySet: this.querySet,
-            queryIndex: 1,
-            location: 'end',
-          }
-        ];
+    const computePassDescriptor: GPUComputePassDescriptor = {};
+    if (shouldTimeProgram && this.supportTimestampQuery) {
+      this.endComputePassEncoder();
+      if (this.querySet == null) {
+        this.querySet = this.device.createQuerySet({
+          type: 'timestamp',
+          count: this.querySetCount,
+        });
       }
+      computePassDescriptor.timestampWrites = [
+        {
+          querySet: this.querySet,
+          queryIndex: 0,
+          location: 'beginning',
+        },
+        {
+          querySet: this.querySet,
+          queryIndex: 1,
+          location: 'end',
+        }
+      ];
+      this.computePassEncoder =
+          this.commandEncoder.beginComputePass(computePassDescriptor);
+    } else if (!this.computePassEncoder) {
       this.computePassEncoder =
           this.commandEncoder.beginComputePass(computePassDescriptor);
     }