fix upload

microsoft · Sep 13, 2022 · a782667 · a782667
1 parent b6e7fba
commit a782667
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/js/web/lib/onnxjs/backends/webgpu/gpu-data-manager.ts b/js/web/lib/onnxjs/backends/webgpu/gpu-data-manager.ts
@@ -13,7 +13,7 @@ export interface GpuDataManager {
   /**
    * upload data to GPU. if the ID already exists in cache, returns the cached value without uploading anything.
    */
-  upload(id: GpuDataId, data: Tensor.NumberType, gpuDataType: GpuDataType): Promise<GpuData>;
+  upload(data: Tensor.NumberType, gpuDataType: GpuDataType): Promise<GpuData>;
   /**
    * create new data on GPU.
    */
@@ -54,7 +54,7 @@ class GpuDataManagerImpl implements GpuDataManager {
     this.downloadCache = new Map();
   }
 
-  async upload(id: GpuDataId, data: Tensor.NumberType, gpuDataType: GpuDataType): Promise<GpuData> {
+  async upload(data: Tensor.NumberType, gpuDataType: GpuDataType): Promise<GpuData> {
     if (gpuDataType !== GpuDataType.default) {
       throw new Error('we only support default GPU data type now');
     }
@@ -72,8 +72,8 @@ class GpuDataManagerImpl implements GpuDataManager {
     new Uint8Array(arrayBuffer).set(new Uint8Array(srcArrayBuffer, srcOffset, srcLength));
     gpuBuffer.unmap();
 
-    const gpuData = {id, type: GpuDataType.default, buffer: gpuBuffer};
-    this.storageCache.set(id, {gpuData, size: srcLength});
+    const gpuData = {id: Guid.create(), type: GpuDataType.default, buffer: gpuBuffer};
+    this.storageCache.set(gpuData.id, {gpuData, size: srcLength});
     return gpuData;
   }
 

diff --git a/js/web/lib/onnxjs/backends/webgpu/inference-handler.ts b/js/web/lib/onnxjs/backends/webgpu/inference-handler.ts
@@ -26,7 +26,7 @@ export class WebGpuInferenceHandler implements InferenceHandler {
     this.dataManager = createTensorDataManager(session.backend.device);
   }
 
-  private uploadGpuData(tensor: Tensor, textureType: GpuDataType): GpuData {
+  private async uploadGpuData(tensor: Tensor, textureType: GpuDataType): Promise<GpuData> {
     if (this.session.isInitializer(tensor.dataId)) {
       return this.session.dataManager.uploadTensorToGpu(tensor, textureType);
     }
@@ -46,7 +46,7 @@ export class WebGpuInferenceHandler implements InferenceHandler {
     // create info for input
     const inputDatas: GpuData[] = [];
     for (let i = 0; i < program.inputTypes.length; ++i) {
-      inputDatas[i] = this.uploadGpuData(inputs[i], program.inputTypes[i]);
+      inputDatas[i] = await this.uploadGpuData(inputs[i], program.inputTypes[i]);
     }
 
     const key = getProgramInfoUniqueKey(program, inputDatas);

diff --git a/js/web/lib/onnxjs/backends/webgpu/tensor-data-manager.ts b/js/web/lib/onnxjs/backends/webgpu/tensor-data-manager.ts
@@ -12,7 +12,7 @@ export interface TensorDataManager {
   /**
    * upload a CPU tensor to GPU.
    */
-  uploadTensorToGpu(tensor: Tensor, gpuDataType: GpuDataType): GpuData;
+  uploadTensorToGpu(tensor: Tensor, gpuDataType: GpuDataType): Promise<GpuData>;
 
   /**
    * create a new GPU tensor.
@@ -55,7 +55,7 @@ class TensorDataManagerImpl implements TensorDataManager {
     tensorIds.add(tensorId);
   }
 
-  uploadTensorToGpu(tensor: Tensor, gpuDataType: GpuDataType): GpuData {
+  async uploadTensorToGpu(tensor: Tensor, gpuDataType: GpuDataType): Promise<GpuData> {
     const gpuDataId = this.map.get(tensor.dataId);
     if (gpuDataId) {
       const gpuData = this.gpuDataManager.get(gpuDataId);
@@ -65,7 +65,7 @@ class TensorDataManagerImpl implements TensorDataManager {
       return gpuData;
     }
 
-    const gpuData = this.gpuDataManager.create(tensor.type, tensor.dims, gpuDataType);
+    const gpuData = await this.gpuDataManager.upload(tensor.numberData, gpuDataType);
     this.registerIdMapping(tensor.dataId, gpuData.id);
     return gpuData;
   }