[WebGPU] Make dataToGPU upload to GPU if data is on CPU (#8483)

tensorflow · Dec 19, 2024 · 2644bd0 · 2644bd0
1 parent cb6206c
commit 2644bd0
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts
@@ -594,16 +594,19 @@ export class WebGPUBackend extends KernelBackend {
    * @param dataId The source tensor.
    */
   override readToGPU(dataId: DataId): GPUData {
-    const srcTensorData = this.tensorMap.get(dataId);
-    const {values, dtype, shape, resource} = srcTensorData;
+    let srcTensorData = this.tensorMap.get(dataId);
+    const {values, dtype, shape} = srcTensorData;
+    let resource = srcTensorData.resource;
 
     if (dtype === 'complex64') {
       throw new Error('Does not support reading buffer for complex64 dtype.');
     }
 
     if (resource == null) {
       if (values != null) {
-        throw new Error('Data is not on GPU but on CPU.');
+        this.uploadToGPU(dataId);
+        srcTensorData = this.tensorMap.get(dataId);
+        resource = srcTensorData.resource;
       } else {
         throw new Error('There is no data on GPU or CPU.');
       }

diff --git a/tfjs-backend-webgpu/src/backend_webgpu_test.ts b/tfjs-backend-webgpu/src/backend_webgpu_test.ts
@@ -200,6 +200,18 @@ describeWebGPU('backend webgpu', () => {
     await c3.data();
     tf.env().set('WEBGPU_DEFERRED_SUBMIT_BATCH_SIZE', savedFlag);
   });
+
+  it('dataToGPU uploads to GPU if the tensor is on CPU', async () => {
+    const webGPUBackend = (tf.backend() as WebGPUBackend);
+    const data = [1,2,3,4,5];
+    const tensor = tf.tensor1d(data);
+    const res = tensor.dataToGPU();
+    expect(res.buffer).toBeDefined();
+    const resData = await webGPUBackend.getBufferData(res.buffer);
+    const values = tf.util.convertBackendValuesAndArrayBuffer(
+      resData, res.tensorRef.dtype);
+    expectArraysEqual(values, data);
+  });
 });
 
 describeWebGPU('backendWebGPU', () => {