From e043ac417f021776552febc5021bc2748b47d572 Mon Sep 17 00:00:00 2001
From: Jiajia Qin <jiajia.qin@intel.com>
Date: Wed, 4 Sep 2019 17:27:21 +0800
Subject: [PATCH] webgpu: Don't release the uniform buffer under delayed mode

In delayed mode, the command list hasn't been executed actually.
So, the uniform buffer can't be recycled.

This patch also fixes the issue that PostNet works incorrectly under
delayed mode.
---
 tfjs-backend-webgpu/src/backend_webgpu.ts      |  9 ++++++---
 tfjs-backend-webgpu/src/backend_webgpu_test.ts | 16 ++++++++--------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts
index b0a3680f46f..b3d16224b88 100644
--- a/tfjs-backend-webgpu/src/backend_webgpu.ts
+++ b/tfjs-backend-webgpu/src/backend_webgpu.ts
@@ -459,12 +459,15 @@ export class WebGPUBackend extends KernelBackend {
     });
     this.commandQueueOwnedIds.add(output.dataId);
 
+    const uniformInfo = {
+      byteSize: uniformData.byteLength,
+      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM,
+      buffer: uniforms.resource.buffer};
+    this.disposalQueue.push(uniformInfo);
+
     if (ENV.get('WEBGPU_IMMEDIATE_EXECUTION_ENABLED')) {
       this.submitQueue();
     }
-    this.releaseBuffer(
-        uniforms.resource.buffer, uniformData.byteLength,
-        GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM);
 
     if (shouldTimeProgram) {
       query = this.endTimer(query);
diff --git a/tfjs-backend-webgpu/src/backend_webgpu_test.ts b/tfjs-backend-webgpu/src/backend_webgpu_test.ts
index 04e8a7cbb8a..e392747d563 100644
--- a/tfjs-backend-webgpu/src/backend_webgpu_test.ts
+++ b/tfjs-backend-webgpu/src/backend_webgpu_test.ts
@@ -43,7 +43,7 @@ describeWebGPU('backend webgpu', () => {
 
     expect(endNumBytes - startNumBytes).toEqual(48);
     expect(endNumTensors - startNumTensors).toEqual(2);
-    expect(endNumBytesInGPU - startNumBytesInGPU).toEqual(48);
+    expect(endNumBytesInGPU - startNumBytesInGPU).toEqual(24);
 
     tf.test_util.expectArraysClose(
         dData, new Float32Array([9, 12, 15, 19, 26, 33]));
@@ -120,7 +120,7 @@ describeWebGPU('backend webgpu', () => {
     tf.ENV.set('WEBGPU_IMMEDIATE_EXECUTION_ENABLED', savedFlag);
   });
 
-  it('should recycle buffers in delayed mode', () => {
+  it('should not recycle buffers in delayed mode', () => {
     const savedFlag = tf.ENV.get('WEBGPU_IMMEDIATE_EXECUTION_ENABLED');
     tf.ENV.set('WEBGPU_IMMEDIATE_EXECUTION_ENABLED', false);
     const backend = tf.backend() as WebGPUBackend;
@@ -139,8 +139,8 @@ describeWebGPU('backend webgpu', () => {
     const freeBuffersAfterFirstMatMul = bufferManager.getNumFreeBuffers();
     const usedBuffersAfterFirstMatMul = bufferManager.getNumUsedBuffers();
     expect(freeBuffersAfterFirstMatMul - freeBuffersAfterFirstMul)
-        .toEqual(1);  // from released uniform
-    expect(usedBuffersAfterFirstMatMul - usedBuffersAfterFirstMul).toEqual(2);
+        .toEqual(0);
+    expect(usedBuffersAfterFirstMatMul - usedBuffersAfterFirstMul).toEqual(3);
 
     const a2 = tf.tensor2d([2, 4, 6, 8], [2, 2]);
     const b2 = tf.tensor2d([0.5, 0.5, 0.5, 0.5], [2, 2]);
@@ -149,15 +149,15 @@ describeWebGPU('backend webgpu', () => {
     const freeBuffersAfterSecondMul = bufferManager.getNumFreeBuffers();
     const usedBuffersAfterSecondMul = bufferManager.getNumUsedBuffers();
     expect(freeBuffersAfterSecondMul - freeBuffersAfterFirstMatMul)
-        .toEqual(0);  // released a uniform buffer and reused a buffer
-    expect(usedBuffersAfterSecondMul - usedBuffersAfterFirstMatMul).toEqual(3);
+        .toEqual(0);
+    expect(usedBuffersAfterSecondMul - usedBuffersAfterFirstMatMul).toEqual(4);
 
     const f2 = tf.tensor2d([1, 2, 3, 4, 5, 6], [2, 3]);
     tf.matMul(c2, f2);
     const freeBuffersAfterSecondMatMul = bufferManager.getNumFreeBuffers();
     const usedBuffersAfterSecondMatMul = bufferManager.getNumUsedBuffers();
     expect(freeBuffersAfterSecondMatMul - freeBuffersAfterSecondMul).toEqual(0);
-    expect(usedBuffersAfterSecondMatMul - usedBuffersAfterSecondMul).toEqual(2);
+    expect(usedBuffersAfterSecondMatMul - usedBuffersAfterSecondMul).toEqual(3);
     tf.ENV.set('WEBGPU_IMMEDIATE_EXECUTION_ENABLED', savedFlag);
   });
 
@@ -174,4 +174,4 @@ describeWebGPU('backend webgpu', () => {
     // Now that data has been downloaded to the CPU, dataSync should work.
     expect(() => c.dataSync()).not.toThrow();
   });
-});
\ No newline at end of file
+});