From 4d753b74a5bc7101fe32c83fcc5690c8d226c2b6 Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Fri, 8 Sep 2023 13:49:24 -0700 Subject: [PATCH] [js/common] prepare work for supporting webgpu IO binding implementation (#17465) ### Description This PR contains a few changes in /js/common/ to support a coming PR for a full implementation of webgpu IO binding. - allows pass-through if value is already a Tensor instance in return value of `handler.run()` called by `InferenceSession.run()` (inference-session-impl.ts). Specifically, onnxruntime-node and onnxruntime-react-native uses native bindings to generate a Tensor-like object so we need to create a real Tensor instance here; for onnxruntime-web the return value is already a Tensor instance. - adds new types for GPU buffer supported types: `'float32'|'int32'` -> `'float32'|'float16'|'int32'|'int64'|'uint32'|'bool'` - exposes types `GpuBufferDataTypes` together with `CpuPinnedDataTypes` and `TextureDataTypes` as exported --- js/common/lib/inference-session-impl.ts | 7 +++++- js/common/lib/tensor-factory-impl.ts | 8 +++---- js/common/lib/tensor-factory.ts | 31 +++++++------------------ js/common/lib/tensor-impl.ts | 12 ++++++---- js/common/lib/tensor.ts | 15 ++++++++++++ 5 files changed, 40 insertions(+), 33 deletions(-) diff --git a/js/common/lib/inference-session-impl.ts b/js/common/lib/inference-session-impl.ts index 2403b82ea8..06949b4a26 100644 --- a/js/common/lib/inference-session-impl.ts +++ b/js/common/lib/inference-session-impl.ts @@ -109,7 +109,12 @@ export class InferenceSession implements InferenceSessionInterface { const returnValue: {[name: string]: OnnxValue} = {}; for (const key in results) { if (Object.hasOwnProperty.call(results, key)) { - returnValue[key] = new Tensor(results[key].type, results[key].data, results[key].dims); + const result = results[key]; + if (result instanceof Tensor) { + returnValue[key] = result; + } else { + returnValue[key] = new Tensor(result.type, result.data, result.dims); + } } } return returnValue; diff --git a/js/common/lib/tensor-factory-impl.ts b/js/common/lib/tensor-factory-impl.ts index 926312e62c..7228c4a970 100644 --- a/js/common/lib/tensor-factory-impl.ts +++ b/js/common/lib/tensor-factory-impl.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import {GpuBufferDataTypes, OptionsDimensions, OptionsFormat, OptionsNormalizationParameters, OptionsTensorFormat, OptionsTensorLayout, TensorFromGpuBufferOptions, TensorFromImageBitmapOptions, TensorFromImageDataOptions, TensorFromImageElementOptions, TensorFromTextureOptions, TensorFromUrlOptions, TextureDataTypes} from './tensor-factory.js'; +import {OptionsDimensions, OptionsFormat, OptionsNormalizationParameters, OptionsTensorFormat, OptionsTensorLayout, TensorFromGpuBufferOptions, TensorFromImageBitmapOptions, TensorFromImageDataOptions, TensorFromImageElementOptions, TensorFromTextureOptions, TensorFromUrlOptions} from './tensor-factory.js'; import {Tensor} from './tensor-impl.js'; import {Tensor as TensorInterface} from './tensor.js'; @@ -239,7 +239,7 @@ export const tensorFromImage = async( /** * implementation of Tensor.fromTexture(). */ -export const tensorFromTexture = ( +export const tensorFromTexture = ( texture: TensorInterface.TextureType, options: TensorFromTextureOptions): Tensor => { const {width, height, download, dispose} = options; // Always assume RGBAF32. TODO: support different texture format @@ -250,7 +250,7 @@ export const tensorFromTexture = ( /** * implementation of Tensor.fromGpuBuffer(). */ -export const tensorFromGpuBuffer = ( +export const tensorFromGpuBuffer = ( gpuBuffer: TensorInterface.GpuBufferType, options: TensorFromGpuBufferOptions): Tensor => { const {dataType, dims, download, dispose} = options; return new Tensor({location: 'gpu-buffer', type: dataType ?? 'float32', gpuBuffer, dims, download, dispose}); @@ -259,6 +259,6 @@ export const tensorFromGpuBuffer = ( /** * implementation of Tensor.fromPinnedBuffer(). */ -export const tensorFromPinnedBuffer = >( +export const tensorFromPinnedBuffer = ( type: T, buffer: TensorInterface.DataTypeMap[T], dims?: readonly number[]): Tensor => new Tensor({location: 'cpu-pinned', type, data: buffer, dims: dims ?? [buffer.length]}); diff --git a/js/common/lib/tensor-factory.ts b/js/common/lib/tensor-factory.ts index 38d3106d56..6e19d7fb89 100644 --- a/js/common/lib/tensor-factory.ts +++ b/js/common/lib/tensor-factory.ts @@ -39,15 +39,10 @@ interface GpuResourceConstructorParameters { dispose?(): void; } -/** - * supported data types for constructing a tensor from a pinned CPU buffer - */ -export type CpuPinnedDataTypes = Exclude; - /** * represent the parameter for constructing a tensor from a pinned CPU buffer */ -export interface CpuPinnedConstructorParameters extends +export interface CpuPinnedConstructorParameters extends CommonConstructorParameters { /** * Specify the location of the data to be 'cpu-pinned'. @@ -59,15 +54,10 @@ export interface CpuPinnedConstructorParameters extends +export interface TextureConstructorParameters extends CommonConstructorParameters, GpuResourceConstructorParameters { /** * Specify the location of the data to be 'texture'. @@ -79,15 +69,10 @@ export interface TextureConstructorParameters extends +export interface GpuBufferConstructorParameters extends CommonConstructorParameters, GpuResourceConstructorParameters { /** * Specify the location of the data to be 'gpu-buffer'. @@ -203,11 +188,11 @@ export interface TensorFromUrlOptions extends OptionsDimensions, OptionResizedDi export interface TensorFromImageBitmapOptions extends OptionResizedDimensions, OptionsTensorFormat, OptionsTensorLayout, OptionsTensorDataType, OptionsNormalizationParameters {} -export interface TensorFromTextureOptions extends +export interface TensorFromTextureOptions extends Required, OptionsFormat, GpuResourceConstructorParameters/* TODO: add more */ {} -export interface TensorFromGpuBufferOptions extends Pick, - GpuResourceConstructorParameters { +export interface TensorFromGpuBufferOptions extends + Pick, GpuResourceConstructorParameters { /** * Describes the data type of the tensor. */ @@ -298,7 +283,7 @@ export interface TensorFactory { * * @returns a tensor object */ - fromTexture( + fromTexture( texture: Tensor.TextureType, options: TensorFromTextureOptions): TypedTensor<'float32'>; /** @@ -318,7 +303,7 @@ export interface TensorFactory { * * @returns a tensor object */ - fromGpuBuffer( + fromGpuBuffer( buffer: Tensor.GpuBufferType, options: TensorFromGpuBufferOptions): TypedTensor; /** diff --git a/js/common/lib/tensor-impl.ts b/js/common/lib/tensor-impl.ts index a3cee19a6a..e3e2b9c728 100644 --- a/js/common/lib/tensor-impl.ts +++ b/js/common/lib/tensor-impl.ts @@ -4,7 +4,7 @@ import {tensorToDataURL, tensorToImageData} from './tensor-conversion-impl.js'; import {TensorToDataUrlOptions, TensorToImageDataOptions} from './tensor-conversion.js'; import {tensorFromGpuBuffer, tensorFromImage, tensorFromPinnedBuffer, tensorFromTexture} from './tensor-factory-impl.js'; -import {CpuPinnedConstructorParameters, CpuPinnedDataTypes, GpuBufferConstructorParameters, GpuBufferDataTypes, TensorFromGpuBufferOptions, TensorFromImageBitmapOptions, TensorFromImageDataOptions, TensorFromImageElementOptions, TensorFromTextureOptions, TensorFromUrlOptions, TextureConstructorParameters} from './tensor-factory.js'; +import {CpuPinnedConstructorParameters, GpuBufferConstructorParameters, TensorFromGpuBufferOptions, TensorFromImageBitmapOptions, TensorFromImageDataOptions, TensorFromImageElementOptions, TensorFromTextureOptions, TensorFromUrlOptions, TextureConstructorParameters} from './tensor-factory.js'; import {checkBigInt, NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP, NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP, SupportedTypedArray, SupportedTypedArrayConstructors} from './tensor-impl-type-mapping.js'; import {calculateSize, tensorReshape} from './tensor-utils-impl.js'; import {Tensor as TensorInterface} from './tensor.js'; @@ -102,7 +102,8 @@ export class Tensor implements TensorInterface { break; } case 'gpu-buffer': { - if (type !== 'float32' && type !== 'int32') { + if ((type !== 'float32' && type !== 'float16' && type !== 'int32' && type !== 'int64' && type !== 'uint32' && + type !== 'bool')) { throw new TypeError(`unsupported type "${type}" to create tensor from gpu buffer`); } this.gpuBufferData = arg0.gpuBuffer; @@ -240,16 +241,17 @@ export class Tensor implements TensorInterface { return tensorFromImage(image, options); } - static fromTexture(texture: TensorTextureType, options: TensorFromTextureOptions<'float32'>): TensorInterface { + static fromTexture( + texture: TensorTextureType, options: TensorFromTextureOptions): TensorInterface { return tensorFromTexture(texture, options); } - static fromGpuBuffer( + static fromGpuBuffer( gpuBuffer: TensorGpuBufferType, options: TensorFromGpuBufferOptions): TensorInterface { return tensorFromGpuBuffer(gpuBuffer, options); } - static fromPinnedBuffer( + static fromPinnedBuffer( type: T, buffer: TensorInterface.DataTypeMap[T], dims?: readonly number[]): Tensor { return tensorFromPinnedBuffer(type, buffer, dims); } diff --git a/js/common/lib/tensor.ts b/js/common/lib/tensor.ts index 10071eda39..6c08d1fe8e 100644 --- a/js/common/lib/tensor.ts +++ b/js/common/lib/tensor.ts @@ -105,11 +105,21 @@ export declare namespace Tensor { type DataType = DataTypeMap[Type]; type ElementType = ElementTypeMap[Type]; + /** + * supported data types for constructing a tensor from a pinned CPU buffer + */ + export type CpuPinnedDataTypes = Exclude; + /** * type alias for WebGL texture */ export type TextureType = WebGLTexture; + /** + * supported data types for constructing a tensor from a WebGL texture + */ + export type TextureDataTypes = 'float32'; + /** * type alias for WebGPU buffer * @@ -122,6 +132,11 @@ export declare namespace Tensor { */ export type GpuBufferType = {size: number; mapState: 'unmapped' | 'pending' | 'mapped'}; + /** + * supported data types for constructing a tensor from a WebGPU buffer + */ + export type GpuBufferDataTypes = 'float32'|'float16'|'int32'|'int64'|'uint32'|'bool'; + /** * represent where the tensor data is stored */