mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-24 19:43:35 +00:00
85cef0af8c
### Description This PR expands the graph capture capability to JS EP, which is similar to #16081. But for JS EP, we don't use the CUDA Graph, instead, we records all gpu commands and replay them, which removes most of the cpu overhead to avoid the the situation that gpu waiting for cpu. mobilenetv2-12 becomes 3.7ms from 6ms on NV 3090 and becomes 3.38ms from 4.58ms on Intel A770. All limitations are similar with CUDA EP: 1. Models with control-flow ops (i.e. If, Loop and Scan ops) are not supported. 2. Usage of graph capture is limited to models where-in all ops in the model can be partitioned to the JS EP or CPU EP and no memory copy between them. 3. Shapes of inputs/outputs cannot change across inference calls. 4. IObinding is required. The usage is like below: Method 1: specify outputs buffers explicitly. ``` const sessionOptions = { executionProviders: [ { name: "webgpu", }, ], enableGraphCapture: true, }; const session = await ort.InferenceSession.create('./models/mobilenetv2-12.onnx', sessionOptions); // prepare the inputBuffer/outputBuffer ... ... const feeds = { 'input': ort.Tensor.fromGpuBuffer(inputBuffer, { dataType: 'float32', dims }) }; const fetches = { 'output': ort.Tensor.fromGpuBuffer(outputBuffer, { dataType: 'float32', dims: [1, 1000] }) }; let results = await session.run(feeds, fetches); // The first run will begin to capture the graph. // update inputBuffer content ... ... results = = await session.run(feeds, fetches); // The 2ed run and after will directly call replay to execute the graph. ... ... session.release(); ``` Method 2: Don't specify outputs buffers explicitly. Internally, when graph capture is enabled, it will set all outputs location to 'gpu-buffer'. ``` const sessionOptions = { executionProviders: [ { name: "webgpu", }, ], enableGraphCapture: true, }; const session = await ort.InferenceSession.create('./models/mobilenetv2-12.onnx', sessionOptions); // prepare the inputBuffer ... ... const feeds = { 'input': ort.Tensor.fromGpuBuffer(inputBuffer, { dataType: 'float32', dims }) }; let results = await session.run(feeds); // The first run will begin to capture the graph. // update inputBuffer content ... ... results = = await session.run(feeds); // The 2ed run and after will directly call replay to execute the graph. ... ... session.release();
200 lines
9.1 KiB
TypeScript
200 lines
9.1 KiB
TypeScript
// 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;
|
|
type CaptureBeginFunction = () => void;
|
|
type CaptureEndFunction = () => void;
|
|
type ReplayFunction = () => void;
|
|
}
|
|
|
|
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): Promise<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;
|
|
|
|
_OrtTrainingGetModelInputOutputCount?
|
|
(trainingHandle: number, inputCount: number, outputCount: number, isEvalModel: boolean): number;
|
|
_OrtTrainingGetModelInputOutputName?
|
|
(trainingHandle: number, index: number, isInput: boolean, isEvalModel: boolean): number;
|
|
|
|
_OrtTrainingReleaseSession?(trainingHandle: number): void;
|
|
// #endregion
|
|
|
|
// #region config
|
|
numThreads?: number;
|
|
mainScriptUrlOrBlob?: string|Blob;
|
|
// #endregion
|
|
|
|
// #region external data API
|
|
mountExternalData?(externalDataFilePath: string, externalDataFileData: Uint8Array): void;
|
|
unmountExternalData?(): void;
|
|
// #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, captureBegin: JSEP.CaptureBeginFunction,
|
|
captureEnd: JSEP.CaptureEndFunction, replay: JSEP.ReplayFunction): 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] 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]>;
|
|
/**
|
|
* [exported from js_internal_api.js] Called when InferenceSession.run started. This function will be called before
|
|
* _OrtRun[WithBinding]() is called.
|
|
* @param sessionId - specify the session ID.
|
|
*/
|
|
jsepOnRunStart: (sessionId: number) => void;
|
|
/**
|
|
* [exported from js_internal_api.js] Release a session. This function will be called before _OrtReleaseSession() is
|
|
* called.
|
|
* @param sessionId - specify the session ID.
|
|
* @returns
|
|
*/
|
|
jsepOnReleaseSession: (sessionId: number) => void;
|
|
// #endregion
|
|
}
|
|
|
|
declare const moduleFactory: EmscriptenModuleFactory<OrtWasmModule>;
|
|
export default moduleFactory;
|