saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
onnxruntime/js/web/lib/wasm/binding/ort-wasm.d.ts
Yulong Wang 561aca97cf
[js/webgpu] support IO binding (#17480) 
<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.
2023-09-29 11:24:42 -07:00

178 lines
8.1 KiB
TypeScript
Raw Blame History

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
import type {Tensor} from 'onnxruntime-common';
export declare namespace JSEP {
type BackendType = unknown;
type AllocFunction = (size: number) => number;
type FreeFunction = (size: number) => number;
type UploadFunction = (dataOffset: number, gpuDataId: number, size: number) => void;
type DownloadFunction = (gpuDataId: number, dataOffset: number, size: number) => Promise<void>;
type CreateKernelFunction = (name: string, kernel: number, attribute: unknown) => void;
type ReleaseKernelFunction = (kernel: number) => void;
type RunFunction =
(kernel: number, contextDataOffset: number, sessionHandle: number, errors: Array<Promise<string|null>>) => number;
}
export interface OrtWasmModule extends EmscriptenModule {
// #region emscripten functions
stackSave(): number;
stackRestore(stack: number): void;
stackAlloc(size: number): number;
UTF8ToString(offset: number, maxBytesToRead?: number): string;
lengthBytesUTF8(str: string): number;
stringToUTF8(str: string, offset: number, maxBytes: number): void;
// #endregion
// #region ORT APIs
_OrtInit(numThreads: number, loggingLevel: number): number;
_OrtGetLastError(errorCodeOffset: number, errorMessageOffset: number): void;
_OrtCreateSession(dataOffset: number, dataLength: number, sessionOptionsHandle: number): number;
_OrtReleaseSession(sessionHandle: number): void;
_OrtGetInputOutputCount(sessionHandle: number, inputCountOffset: number, outputCountOffset: number): number;
_OrtGetInputName(sessionHandle: number, index: number): number;
_OrtGetOutputName(sessionHandle: number, index: number): number;
_OrtFree(stringHandle: number): void;
_OrtCreateTensor(
dataType: number, dataOffset: number, dataLength: number, dimsOffset: number, dimsLength: number,
dataLocation: number): number;
_OrtGetTensorData(tensorHandle: number, dataType: number, dataOffset: number, dimsOffset: number, dimsLength: number):
number;
_OrtReleaseTensor(tensorHandle: number): void;
_OrtCreateBinding(sessionHandle: number): number;
_OrtBindInput(bindingHandle: number, nameOffset: number, tensorHandle: number): Promise<number>;
_OrtBindOutput(bindingHandle: number, nameOffset: number, tensorHandle: number, location: number): number;
_OrtClearBoundOutputs(ioBindingHandle: number): void;
_OrtReleaseBinding(ioBindingHandle: number): void;
_OrtRunWithBinding(
sessionHandle: number, ioBindingHandle: number, outputCount: number, outputsOffset: number,
runOptionsHandle: number): Promise<number>;
_OrtRun(
sessionHandle: number, inputNamesOffset: number, inputsOffset: number, inputCount: number,
outputNamesOffset: number, outputCount: number, outputsOffset: number, runOptionsHandle: number): Promise<number>;
_OrtCreateSessionOptions(
graphOptimizationLevel: number, enableCpuMemArena: boolean, enableMemPattern: boolean, executionMode: number,
enableProfiling: boolean, profileFilePrefix: number, logId: number, logSeverityLevel: number,
logVerbosityLevel: number, optimizedModelFilePath: number): number;
_OrtAppendExecutionProvider(sessionOptionsHandle: number, name: number): number;
_OrtAddFreeDimensionOverride(sessionOptionsHandle: number, name: number, dim: number): number;
_OrtAddSessionConfigEntry(sessionOptionsHandle: number, configKey: number, configValue: number): number;
_OrtReleaseSessionOptions(sessionOptionsHandle: number): void;
_OrtCreateRunOptions(logSeverityLevel: number, logVerbosityLevel: number, terminate: boolean, tag: number): number;
_OrtAddRunConfigEntry(runOptionsHandle: number, configKey: number, configValue: number): number;
_OrtReleaseRunOptions(runOptionsHandle: number): void;
_OrtEndProfiling(sessionHandle: number): number;
// #endregion
// #region ORT Training APIs
_OrtTrainingLoadCheckpoint?(dataOffset: number, dataLength: number): number;
_OrtTrainingReleaseCheckpoint?(checkpointHandle: number): void;
_OrtTrainingCreateSession?
(sessionOptionsHandle: number, checkpointHandle: number, trainOffset: number, trainLength: number,
evalOffset: number, evalLength: number, optimizerOffset: number, optimizerLength: number): number;
_OrtTrainingLazyResetGrad?(trainingHandle: number): number;
_OrtTrainingRunTrainStep?
(trainingHandle: number, inputsOffset: number, inputCount: number, outputsOffset: number, outputCount: number,
runOptionsHandle: number): number;
_OrtTrainingOptimizerStep?(trainingHandle: number, runOptionsHandle: number): number;
_OrtTrainingEvalStep?
(trainingHandle: number, inputsOffset: number, inputCount: number, outputsOffset: number, outputCount: number,
runOptionsHandle: number): number;
_OrtTrainingGetParametersSize?(trainingHandle: number, paramSizeT: number, trainableOnly: boolean): number;
_OrtTrainingCopyParametersToBuffer?
(trainingHandle: number, parametersBuffer: number, parameterCount: number, trainableOnly: boolean): number;
_OrtTrainingCopyParametersFromBuffer?
(trainingHandle: number, parametersBuffer: number, parameterCount: number, trainableOnly: boolean): number;
_OrtTrainingReleaseSession?(trainingHandle: number): void;
// #endregion
// #region config
mainScriptUrlOrBlob?: string|Blob;
// #endregion
// #region JSEP
/**
* This is the entry of JSEP initialization. This function is called once when initializing ONNX Runtime.
* This function initializes WebGPU backend and registers a few callbacks that will be called in C++ code.
*/
jsepInit?
(backend: JSEP.BackendType, alloc: JSEP.AllocFunction, free: JSEP.FreeFunction, upload: JSEP.UploadFunction,
download: JSEP.DownloadFunction, createKernel: JSEP.CreateKernelFunction,
releaseKernel: JSEP.ReleaseKernelFunction, run: JSEP.RunFunction): void;
/**
* [exported from wasm] Specify a kernel's output when running OpKernel::Compute().
*
* @param context - specify the kernel context pointer.
* @param index - specify the index of the output.
* @param data - specify the pointer to encoded data of type and dims.
*/
_JsepOutput(context: number, index: number, data: number): number;
/**
* [exported from wasm] Get name of an operator node.
*
* @param kernel - specify the kernel pointer.
* @returns the pointer to a C-style UTF8 encoded string representing the node name.
*/
_JsepGetNodeName(kernel: number): number;
/**
* [exported from js_internal_api.js] Register a user GPU buffer for usage of a session's input or output.
*
* @param sessionId - specify the session ID.
* @param index - specify an integer to represent which input/output it is registering for. For input, it is the
* input_index corresponding to the session's inputNames. For output, it is the inputCount + output_index
* corresponding to the session's ouputNames.
* @param buffer - specify the GPU buffer to register.
* @param size - specify the original data size in byte.
* @returns the GPU data ID for the registered GPU buffer.
*/
jsepRegisterBuffer: (sessionId: number, index: number, buffer: GPUBuffer, size: number) => number;
/**
* [exported from js_internal_api.js] Unregister all user GPU buffers for a session.
*
* @param sessionId - specify the session ID.
*/
jsepUnregisterBuffers?: (sessionId: number) => void;
/**
* [exported from js_internal_api.js] Get the GPU buffer by GPU data ID.
*
* @param dataId - specify the GPU data ID
* @returns the GPU buffer.
*/
jsepGetBuffer: (dataId: number) => GPUBuffer;
/**
* [exported from js_internal_api.js] Create a function to be used to create a GPU Tensor.
*
* @param gpuBuffer - specify the GPU buffer
* @param size - specify the original data size in byte.
* @param type - specify the tensor type.
* @returns the generated downloader function.
*/
jsepCreateDownloader:
(gpuBuffer: GPUBuffer, size: number,
type: Tensor.GpuBufferDataTypes) => () => Promise<Tensor.DataTypeMap[Tensor.GpuBufferDataTypes]>;
// #endregion
}
declare const moduleFactory: EmscriptenModuleFactory<OrtWasmModule>;
export default moduleFactory;
Reference in a new issue View git blame Copy permalink