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 +});