From 69a434e96487a185ac450b20a21882b263b668f0 Mon Sep 17 00:00:00 2001 From: Xu Xing Date: Mon, 14 Nov 2022 12:08:02 +0800 Subject: [PATCH 1/2] Refactor type conversion for read back Bug: https://github.com/tensorflow/tfjs/issues/6965 --- tfjs-backend-cpu/src/backend_cpu.ts | 7 ++++--- tfjs-backend-webgpu/src/backend_webgpu.ts | 9 ++++----- .../src/backend_webgpu_test.ts | 13 ++++++------- tfjs-backend-webgpu/src/webgpu_util.ts | 12 ------------ tfjs-core/src/util_base.ts | 19 ++++++++++++++++++- 5 files changed, 32 insertions(+), 28 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 2fe76cdf23..854cb366c6 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -49,7 +49,8 @@ export class MathBackendCPU extends KernelBackend { this.data = new DataStorage(this, engine()); } - override write(values: backend_util.BackendValues, shape: number[], + override write( + values: backend_util.BackendValues, shape: number[], dtype: DataType): DataId { if (this.firstUse) { this.firstUse = false; @@ -138,8 +139,8 @@ export class MathBackendCPU extends KernelBackend { this.readSync(complexTensorInfos.imag.dataId) as Float32Array; return backend_util.mergeRealAndImagArrays(realValues, imagValues); } - - return this.data.get(dataId).values; + return util.convertBackendValuesAndArrayBuffer( + this.data.get(dataId).values, dtype); } bufferSync(t: TensorInfo): diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts index 4d64539727..1227f74890 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu.ts @@ -357,7 +357,7 @@ export class WebGPUBackend extends KernelBackend { } public async getBufferData(buffer: GPUBuffer, size: number): - Promise { + Promise { const staging = this.bufferManager.acquireBuffer( size, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ); this.ensureCommandEncoderReady(); @@ -383,7 +383,7 @@ export class WebGPUBackend extends KernelBackend { this.dummyContext.getCurrentTexture(); } - return values as backend_util.BackendValues; + return values; } private convertAndCacheOnCPU(dataId: DataId, data: backend_util.TypedArray): @@ -439,10 +439,9 @@ export class WebGPUBackend extends KernelBackend { } else { const bufferInfo = tensorData.resourceInfo as BufferInfo; const data = await this.getBufferData(bufferInfo.buffer, bufferInfo.size); - vals = webgpu_util.ArrayBufferToTypedArray( - data as ArrayBuffer, tensorData.dtype); + vals = util.convertBackendValuesAndArrayBuffer(data, tensorData.dtype); } - this.convertAndCacheOnCPU(dataId, vals); + this.convertAndCacheOnCPU(dataId, vals as backend_util.TypedArray); return vals; } diff --git a/tfjs-backend-webgpu/src/backend_webgpu_test.ts b/tfjs-backend-webgpu/src/backend_webgpu_test.ts index f5792bdbab..816165c3a8 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu_test.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu_test.ts @@ -22,7 +22,6 @@ const {expectArraysEqual, expectArraysClose} = test_util; import {WebGPUBackend, WebGPUMemoryInfo} from './backend_webgpu'; import {describeWebGPU} from './test_util'; -import * as webgpu_util from './webgpu_util'; describeWebGPU('backend webgpu cpu forwarding turned on', () => { let cpuForwardFlagSaved: boolean; @@ -274,8 +273,8 @@ describeWebGPU('keeping data on gpu ', () => { `Expected: float32`); } const resData = await webGPUBackend.getBufferData(res.buffer, res.bufSize); - const values = webgpu_util.ArrayBufferToTypedArray( - resData as ArrayBuffer, res.tensorRef.dtype); + const values = tf.util.convertBackendValuesAndArrayBuffer( + resData, res.tensorRef.dtype); expectArraysEqual(values, data); }); @@ -294,8 +293,8 @@ describeWebGPU('keeping data on gpu ', () => { `Expected: float32`); } const resData = await webGPUBackend.getBufferData(res.buffer, res.bufSize); - const values = webgpu_util.ArrayBufferToTypedArray( - resData as ArrayBuffer, res.tensorRef.dtype); + const values = tf.util.convertBackendValuesAndArrayBuffer( + resData, res.tensorRef.dtype); expectArraysEqual(values, data); }); @@ -340,8 +339,8 @@ describeWebGPU('keeping data on gpu ', () => { const res = result as unknown as GPUData; const resData = await webGPUBackend.getBufferData(res.buffer, res.bufSize); - const values = webgpu_util.ArrayBufferToTypedArray( - resData as ArrayBuffer, res.tensorRef.dtype); + const values = tf.util.convertBackendValuesAndArrayBuffer( + resData, res.tensorRef.dtype); expectArraysEqual(values, data); }); diff --git a/tfjs-backend-webgpu/src/webgpu_util.ts b/tfjs-backend-webgpu/src/webgpu_util.ts index 3c6867c89b..93b008d476 100644 --- a/tfjs-backend-webgpu/src/webgpu_util.ts +++ b/tfjs-backend-webgpu/src/webgpu_util.ts @@ -154,18 +154,6 @@ export function GPUBytesPerElement(dtype: DataType): number { } } -export function ArrayBufferToTypedArray(data: ArrayBuffer, dtype: DataType) { - if (dtype === 'float32') { - return new Float32Array(data); - } else if (dtype === 'int32') { - return new Int32Array(data); - } else if (dtype === 'bool' || dtype === 'string') { - return Uint8Array.from(new Int32Array(data)); - } else { - throw new Error(`Unknown dtype ${dtype}`); - } -} - export function isWebGPUSupported(): boolean { return ((typeof window !== 'undefined') || //@ts-ignore diff --git a/tfjs-core/src/util_base.ts b/tfjs-core/src/util_base.ts index 84fbfdaeb6..6d4ec37779 100644 --- a/tfjs-core/src/util_base.ts +++ b/tfjs-core/src/util_base.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -import {DataType, DataTypeMap, FlatVector, NumericDataType, RecursiveArray, TensorLike, TypedArray, WebGLData, WebGPUData} from './types'; +import {BackendValues, DataType, DataTypeMap, FlatVector, NumericDataType, RecursiveArray, TensorLike, TypedArray, WebGLData, WebGPUData} from './types'; /** * Shuffles the array in-place using Fisher-Yates algorithm. @@ -661,6 +661,23 @@ export function toNestedArray( return createNestedArray(0, shape, a, isComplex); } +export function convertBackendValuesAndArrayBuffer( + data: BackendValues|ArrayBuffer, dtype: DataType) { + // If is type Uint8Array[], return it directly. + if (Array.isArray(data)) { + return data; + } + if (dtype === 'float32') { + return data instanceof Float32Array ? data : new Float32Array(data); + } else if (dtype === 'int32') { + return data instanceof Int32Array ? data : new Int32Array(data); + } else if (dtype === 'bool' || dtype === 'string') { + return Uint8Array.from(new Int32Array(data)); + } else { + throw new Error(`Unknown dtype ${dtype}`); + } +} + export function makeOnesTypedArray( size: number, dtype: D): DataTypeMap[D] { const array = makeZerosTypedArray(size, dtype); From a74bd16a16ba8b74eb0b883a8d94a7145f100e22 Mon Sep 17 00:00:00 2001 From: Xu Xing Date: Thu, 24 Nov 2022 14:09:22 +0800 Subject: [PATCH 2/2] Cleanup --- tfjs-backend-webgpu/src/backend_webgpu.ts | 36 ++++++++++------------- 1 file changed, 15 insertions(+), 21 deletions(-) diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts index 1227f74890..ab78ef7777 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu.ts @@ -17,7 +17,7 @@ import './flags_webgpu'; -import {backend_util, buffer, DataStorage, DataType, engine, env, GPUData, KernelBackend, Rank, RecursiveArray, ShapeMap, Tensor, TensorBuffer, TensorInfo, TimingInfo, TypedArray, util, WebGPUData} from '@tensorflow/tfjs-core'; +import {backend_util, BackendValues, buffer, DataStorage, DataType, engine, env, GPUData, KernelBackend, Rank, RecursiveArray, ShapeMap, Tensor, TensorBuffer, TensorInfo, TimingInfo, TypedArray, util, WebGPUData} from '@tensorflow/tfjs-core'; import {AdapterInfo} from './adapter_info'; import {BufferManager} from './buffer_manager'; @@ -46,7 +46,7 @@ export type TextureInfo = { }; type TensorData = { - values: backend_util.BackendValues, + values: BackendValues, dtype: DataType, shape: number[], refCount: number, @@ -290,9 +290,8 @@ export class WebGPUBackend extends KernelBackend { } } - override write( - values: backend_util.BackendValues, shape: number[], - dtype: DataType): DataId { + override write(values: BackendValues, shape: number[], dtype: DataType): + DataId { if (dtype === 'complex64' && values != null) { throw new Error( `Cannot write to a complex64 dtype. ` + @@ -304,8 +303,8 @@ export class WebGPUBackend extends KernelBackend { } override move( - dataId: DataId, values: backend_util.BackendValues, shape: number[], - dtype: DataType, refCount: number): void { + dataId: DataId, values: BackendValues, shape: number[], dtype: DataType, + refCount: number): void { if (dtype === 'complex64') { throw new Error( `Cannot write to a complex64 dtype. ` + @@ -386,8 +385,8 @@ export class WebGPUBackend extends KernelBackend { return values; } - private convertAndCacheOnCPU(dataId: DataId, data: backend_util.TypedArray): - backend_util.TypedArray { + private convertAndCacheOnCPU(dataId: DataId, data: BackendValues): + BackendValues { const tensorData = this.tensorMap.get(dataId); this.releaseResource(dataId); tensorData.values = data; @@ -396,7 +395,7 @@ export class WebGPUBackend extends KernelBackend { // TODO: Remove once this is fixed: // https://github.com/tensorflow/tfjs/issues/1595 - override readSync(dataId: object): backend_util.BackendValues { + override readSync(dataId: object): BackendValues { const tensorData = this.tensorMap.get(dataId); const {values} = tensorData; @@ -408,7 +407,7 @@ export class WebGPUBackend extends KernelBackend { return values; } - override async read(dataId: object): Promise { + override async read(dataId: object): Promise { if (!this.tensorMap.has(dataId)) { throw new Error(`Tensor ${dataId} was not registered!`); } @@ -417,15 +416,11 @@ export class WebGPUBackend extends KernelBackend { const {values} = tensorData; if (values != null) { - // TODO(xing.xu@intel.com): Merge backend_util.BackendValues and - // backend_util.TypedArray. - return this.convertAndCacheOnCPU( - dataId, values as backend_util.TypedArray) as - backend_util.BackendValues; + return this.convertAndCacheOnCPU(dataId, values); } // Download the values from the GPU. - let vals: backend_util.BackendValues; + let vals: BackendValues; if (tensorData.dtype === 'complex64') { const ps = await Promise.all([ this.read(tensorData.complexTensorInfos.real.dataId), @@ -441,7 +436,7 @@ export class WebGPUBackend extends KernelBackend { const data = await this.getBufferData(bufferInfo.buffer, bufferInfo.size); vals = util.convertBackendValuesAndArrayBuffer(data, tensorData.dtype); } - this.convertAndCacheOnCPU(dataId, vals as backend_util.TypedArray); + this.convertAndCacheOnCPU(dataId, vals); return vals; } @@ -604,13 +599,12 @@ export class WebGPUBackend extends KernelBackend { makeTensorInfo( shape: number[], dtype: DataType, - values?: backend_util.BackendValues|string[]): TensorInfo { + values?: BackendValues|string[]): TensorInfo { if (dtype === 'string' && values != null && values.length > 0 && util.isString(values[0])) { values = (values as unknown as string[]).map(d => util.encodeString(d)); } - const dataId = - this.write(values as backend_util.BackendValues, shape, dtype); + const dataId = this.write(values as BackendValues, shape, dtype); return {dataId, shape, dtype}; }