mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-12 17:57:38 +00:00
561aca97cf
<del> **This PR is based on a few prerequisites PRs. They are listed as below:** - #17465 - #17469 - #17470 - #17472 - #17473 - #17484 Please review the current change by only looking at commit e2e6623e673ec6de55a5c1f8edcbd3a46b535a89 and later. </del> ### Description This PR introduces WebGPU IO binding. This new feature allows onnxruntime-web users to use tensors created from GPU as model input/output so that a model inferencing can be done without unnecessary data copy between CPU and GPU for model input/output. ### Examples An E2E demo/example is being worked on. Following is some simple demo with code snippet. Let's first check today how we do: ```js // STEP.1 - create an inference session: const mySession = await ort.InferenceSession.create('./my_model.onnx', { executionProviders: ['webgpu'] }); // STEP.2 - create model input: (supposing myImageCpuData is a Float32Array) const feeds = { 'input_image:0': new ort.Tensor('float32', myImageCpuData, [1, 224, 224, 3]) }; // STEP.3 - run model const myResults = await mySession.run(feeds); // STEP.4 - get output data const myData = myResults['output_image:0'].data; // Float32Array ``` #### for inputs (GPU tensor): Now, with IO binding, you can create a tensor from a GPU buffer, and feed it to the model: ```js // new STEP.2.A - create model input from a GPU buffer: (supposing myInputGpuBuffer is a `GPUBuffer` object with input data) const feeds = { 'input_image:0': ort.Tensor.fromGpuBuffer(myInputGpuBuffer, { dataType: 'float32', dims: [1, 224, 224, 3] }) }; ``` ### for outputs (pre-allocated GPU tensor) you can also do that for output, **if you know the output shape**: ```js // new STEP.2.B - create model output from a GPU buffer: (supposing myOutputGpuBuffer is a pre-allocated `GPUBuffer` object) const fetches = { 'output_image:0': ort.Tensor.fromGpuBuffer(myOutputGpuBuffer, { dataType: 'float32', dims: [1, 512, 512, 3] }) }; // new STEP.3 - run model with pre-allocated output (fetches) const myResults = await mySession.run(feeds, fetches); ``` ### for outputs (specify location) if you do not know the output shape, you can specify the output location when creating the session: ```js // new STEP.1 - create an inference session with an option "preferredOutputLocation": const mySession = await ort.InferenceSession.create('./my_model.onnx', { executionProviders: ['webgpu'], preferredOutputLocation: "gpu-buffer" }); ``` if the model has multiple outputs, you can specify them seperately: ```js // new STEP.1 - create an inference session with an option "preferredOutputLocation": const mySession = await ort.InferenceSession.create('./my_model.onnx', { executionProviders: ['webgpu'], preferredOutputLocation: { "output_image:0": "gpu-buffer" } }); ``` now you don't need to prepare the `fetches` object and onnxruntime-web will prepare output data on the location that specified. #### read data when you get the output tensor, you can: ```js // get the gpu buffer object: const gpuBuffer = myOutputTensor.gpuBuffer; // GPUBuffer // get the CPU data asynchronizely const cpuData = await myOutputTensor.getData(); // get the CPU data asynchronizely and release the underlying GPU resources const cpuData = await myOutputTensor.getData(true); // dispose the tensor (release the underlying GPU resources). This tensor object will be invalid after dispose() is called. myOutputTensor.dispose(); ``` #### resource management JavaScript has GC so you don't need to worry about managing JavaScript objects. But there are 2 types of resources that are not managed by GC: - GPU buffer that used in tensors - Underlying ORT native resources To simplify, most of the unmanaged resources and handled inside ORT web. But there are a few resources that need users to manage: - All external GPU resources, including GPU buffers inside all tensors created by `Tensor.fromGpuBuffer()`, will not be managed by ORT. User should manage those GPU buffers themselves. - When a session is created with `preferredOutputLocation` == "gpu-buffer" specified in session options, and the corresponding output is not pre-allocated, user need to call the output tensor's `dispose()` or `getData(true)` to manually release the underlying GPU buffers. - ORT internal errors (including providing a pre-allocated output tensor with wrong type/dims) will invalidate the whole wasm memory and is not recoverable. An exception is thrown in this situation.
77 lines
2.4 KiB
TypeScript
77 lines
2.4 KiB
TypeScript
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
import {Env, InferenceSession, Tensor} from 'onnxruntime-common';
|
|
|
|
export type SerializableTensorMetadata =
|
|
[dataType: Tensor.Type, dims: readonly number[], data: Tensor.DataType, location: 'cpu'];
|
|
|
|
export type GpuBufferMetadata = {
|
|
gpuBuffer: Tensor.GpuBufferType;
|
|
download?: () => Promise<Tensor.DataTypeMap[Tensor.GpuBufferDataTypes]>;
|
|
dispose?: () => void;
|
|
};
|
|
|
|
export type UnserializableTensorMetadata =
|
|
[dataType: Tensor.Type, dims: readonly number[], data: GpuBufferMetadata, location: 'gpu-buffer']|
|
|
[dataType: Tensor.Type, dims: readonly number[], data: Tensor.DataType, location: 'cpu-pinned'];
|
|
|
|
export type TensorMetadata = SerializableTensorMetadata|UnserializableTensorMetadata;
|
|
|
|
export type SerializableSessionMetadata = [sessionHandle: number, inputNames: string[], outputNames: string[]];
|
|
|
|
export type SerializableModeldata = [modelDataOffset: number, modelDataLength: number];
|
|
|
|
interface MessageError {
|
|
err?: string;
|
|
}
|
|
|
|
interface MessageInitWasm extends MessageError {
|
|
type: 'init-wasm';
|
|
in ?: Env.WebAssemblyFlags;
|
|
}
|
|
|
|
interface MessageInitOrt extends MessageError {
|
|
type: 'init-ort';
|
|
in ?: Env;
|
|
}
|
|
|
|
interface MessageCreateSessionAllocate extends MessageError {
|
|
type: 'create_allocate';
|
|
in ?: {model: Uint8Array};
|
|
out?: SerializableModeldata;
|
|
}
|
|
|
|
interface MessageCreateSessionFinalize extends MessageError {
|
|
type: 'create_finalize';
|
|
in ?: {modeldata: SerializableModeldata; options?: InferenceSession.SessionOptions};
|
|
out?: SerializableSessionMetadata;
|
|
}
|
|
|
|
interface MessageCreateSession extends MessageError {
|
|
type: 'create';
|
|
in ?: {model: Uint8Array; options?: InferenceSession.SessionOptions};
|
|
out?: SerializableSessionMetadata;
|
|
}
|
|
|
|
interface MessageReleaseSession extends MessageError {
|
|
type: 'release';
|
|
in ?: number;
|
|
}
|
|
|
|
interface MessageRun extends MessageError {
|
|
type: 'run';
|
|
in ?: {
|
|
sessionId: number; inputIndices: number[]; inputs: SerializableTensorMetadata[]; outputIndices: number[];
|
|
options: InferenceSession.RunOptions;
|
|
};
|
|
out?: SerializableTensorMetadata[];
|
|
}
|
|
|
|
interface MesssageEndProfiling extends MessageError {
|
|
type: 'end-profiling';
|
|
in ?: number;
|
|
}
|
|
|
|
export type OrtWasmMessage = MessageInitWasm|MessageInitOrt|MessageCreateSessionAllocate|MessageCreateSessionFinalize|
|
|
MessageCreateSession|MessageReleaseSession|MessageRun|MesssageEndProfiling;
|