diff --git a/js/web/lib/wasm/wasm-core-impl.ts b/js/web/lib/wasm/wasm-core-impl.ts index 3979af7fa1ec9..bb532e0fbae74 100644 --- a/js/web/lib/wasm/wasm-core-impl.ts +++ b/js/web/lib/wasm/wasm-core-impl.ts @@ -309,12 +309,12 @@ export const createSession = async ( if (context) { wasm.currentContext = context as MLContext; } else if (gpuDevice) { - wasm.currentContext = await wasm.jsepCreateMLContext!(gpuDevice); + wasm.currentContext = await wasm.webnnCreateMLContext!(gpuDevice); } else { - wasm.currentContext = await wasm.jsepCreateMLContext!({ deviceType, powerPreference }); + wasm.currentContext = await wasm.webnnCreateMLContext!({ deviceType, powerPreference }); } } else { - wasm.currentContext = await wasm.jsepCreateMLContext!(); + wasm.currentContext = await wasm.webnnCreateMLContext!(); } break; } @@ -330,7 +330,7 @@ export const createSession = async ( // clear current MLContext after session creation if (wasm.currentContext) { - wasm.jsepRegisterMLContext!(sessionHandle, wasm.currentContext); + wasm.webnnRegisterMLContext!(sessionHandle, wasm.currentContext); wasm.currentContext = undefined; wasm.shouldTransferToMLTensor = true; } @@ -454,6 +454,7 @@ export const releaseSession = (sessionId: number): void => { } wasm.jsepOnReleaseSession?.(sessionId); + wasm.webnnOnReleaseSession?.(sessionId); wasm.webgpuOnReleaseSession?.(sessionId); inputNamesUTF8Encoded.forEach((buf) => wasm._OrtFree(buf)); @@ -520,7 +521,7 @@ export const prepareInputOutputTensor = async ( const mlTensor = tensor[2].mlTensor as MLTensor; dataByteLength = calculateTensorSizeInBytes(tensorDataTypeStringToEnum(dataType), dims)!; - const registerMLTensor = wasm.jsepRegisterMLTensor; + const registerMLTensor = wasm.webnnRegisterMLTensor; if (!registerMLTensor) { throw new Error('Tensor location "ml-tensor" is not supported without using WebNN.'); } @@ -540,7 +541,7 @@ export const prepareInputOutputTensor = async ( wasm.setValue(rawData + i * ptrSize, allocWasmString(data[i], allocs), '*'); } } else { - const isGraphInput = wasm.jsepIsGraphInput; + const isGraphInput = wasm.webnnIsGraphInput; if (dataType !== 'string' && isGraphInput) { const tensorNameUTF8 = wasm._OrtGetInputName(sessionId, index); const tensorName = wasm.UTF8ToString(tensorNameUTF8); @@ -549,8 +550,8 @@ export const prepareInputOutputTensor = async ( const dataTypeEnum = tensorDataTypeStringToEnum(dataType); dataByteLength = calculateTensorSizeInBytes(dataTypeEnum, dims)!; actualLocation = 'ml-tensor'; - const createTemporaryTensor = wasm.jsepCreateTemporaryTensor; - const uploadTensor = wasm.jsepUploadTensor; + const createTemporaryTensor = wasm.webnnCreateTemporaryTensor; + const uploadTensor = wasm.webnnUploadTensor; if (!createTemporaryTensor || !uploadTensor) { throw new Error('Tensor location "ml-tensor" is not supported without using WebNN.'); } @@ -722,6 +723,7 @@ export const run = async ( } wasm.jsepOnRunStart?.(sessionHandle); + wasm.webnnOnRunStart?.(sessionHandle); let errorCode: number; if (!BUILD_DEFS.DISABLE_JSEP && ioBindingState) { @@ -862,8 +864,8 @@ export const run = async ( ]); } } else if (preferredLocation === 'ml-tensor' && size > 0) { - const ensureTensor = wasm.jsepEnsureTensor; - const isInt64Supported = wasm.jsepIsInt64Supported; + const ensureTensor = wasm.webnnEnsureTensor; + const isInt64Supported = wasm.webnnIsInt64Supported; if (!ensureTensor || !isInt64Supported) { throw new Error('preferredLocation "ml-tensor" is not supported without using WebNN.'); } @@ -890,9 +892,9 @@ export const run = async ( dims, { mlTensor, - download: wasm.jsepCreateMLTensorDownloader!(dataOffset, type), + download: wasm.webnnCreateMLTensorDownloader!(dataOffset, type), dispose: () => { - wasm.jsepReleaseTensorId!(dataOffset); + wasm.webnnReleaseTensorId!(dataOffset); wasm._OrtReleaseTensor(tensor); }, }, @@ -915,7 +917,7 @@ export const run = async ( if (!keepOutputTensor) { wasm._OrtReleaseTensor(tensor); } - wasm.jsepOnRunEnd?.(sessionHandle); + wasm.webnnOnRunEnd?.(sessionHandle); } } diff --git a/js/web/lib/wasm/wasm-types.ts b/js/web/lib/wasm/wasm-types.ts index 6de54078af031..752bac28d7efb 100644 --- a/js/web/lib/wasm/wasm-types.ts +++ b/js/web/lib/wasm/wasm-types.ts @@ -156,12 +156,26 @@ export declare namespace JSEP { */ shouldTransferToMLTensor: boolean; + /** + * [exported from pre-jsep.js] Called when InferenceSession.run started. This function will be called before + * _OrtRun[WithBinding]() is called. + * @param sessionId - specify the session ID. + */ + webnnOnRunStart: (sessionId: number) => void; + /** + * [exported from pre-jsep.js] Release a session. This function will be called before _OrtReleaseSession() is + * called. + * @param sessionId - specify the session ID. + * @returns + */ + webnnOnReleaseSession: (sessionId: number) => void; + /** * [exported from pre-jsep.js] Called when InferenceSession.run finished. This function will be called after * _OrtRun[WithBinding]() is called. * @param sessionId - specify the session ID. */ - jsepOnRunEnd: (sessionId: number) => void; + webnnOnRunEnd: (sessionId: number) => void; /** * [exported from pre-jsep.js] Register MLContext for a session. @@ -169,18 +183,18 @@ export declare namespace JSEP { * @param context - specify the MLContext. * @returns */ - jsepRegisterMLContext: (sessionId: number, context: MLContext) => void; + webnnRegisterMLContext: (sessionId: number, context: MLContext) => void; /** * [exported from pre-jsep.js] Reserve a MLTensor ID attached to the current session. * @returns the MLTensor ID. */ - jsepReserveTensorId: () => number; + webnnReserveTensorId: () => number; /** * [exported from pre-jsep.js] Release an MLTensor ID from use and destroys underlying MLTensor if no longer in use. * @param tensorId - specify the MLTensor ID. * @returns */ - jsepReleaseTensorId: (tensorId: number) => void; + webnnReleaseTensorId: (tensorId: number) => void; /** * [exported from pre-jsep.js] Ensure that an MLTensor of a given type and shape exists for a MLTensor ID. * @param sessionId - specify the session ID or current active session ID if undefined. @@ -190,7 +204,7 @@ export declare namespace JSEP { * @param copyOld - specify whether to copy the old tensor if a new tensor was created. * @returns the MLTensor associated with the tensor ID. */ - jsepEnsureTensor: ( + webnnEnsureTensor: ( sessionId: number | undefined, tensorId: number, dataType: DataType, @@ -203,20 +217,20 @@ export declare namespace JSEP { * @param data - specify the data to upload. It can be a TensorProto::data_type or a WebNN MLOperandDataType. * @returns */ - jsepUploadTensor: (tensorId: number, data: Uint8Array) => void; + webnnUploadTensor: (tensorId: number, data: Uint8Array) => void; /** * [exported from pre-jsep.js] Download data from an MLTensor. * @param tensorId - specify the MLTensor ID. * @returns the downloaded data. */ - jsepDownloadTensor: (tensorId: number, dstBuffer: ArrayBufferView | ArrayBuffer) => Promise; + webnnDownloadTensor: (tensorId: number, dstBuffer: ArrayBufferView | ArrayBuffer) => Promise; /** * [exported from pre-jsep.js] Creates a downloader function to download data from an MLTensor. * @param tensorId - specify the MLTensor ID. * @param type - specify the data type. * @returns the downloader function. */ - jsepCreateMLTensorDownloader: ( + webnnCreateMLTensorDownloader: ( tensorId: number, type: Tensor.MLTensorDataTypes, ) => () => Promise; @@ -228,7 +242,7 @@ export declare namespace JSEP { * @param dimensions - specify the dimensions. * @returns the MLTensor ID for the external MLTensor. */ - jsepRegisterMLTensor: ( + webnnRegisterMLTensor: ( sessionId: number, tensor: MLTensor, onnxDataType: DataType, @@ -240,7 +254,7 @@ export declare namespace JSEP { * @param optionsOrGpuDevice - specify the options or GPUDevice. * @returns */ - jsepCreateMLContext(optionsOrGpuDevice?: MLContextOptions | GPUDevice): Promise; + webnnCreateMLContext(optionsOrGpuDevice?: MLContextOptions | GPUDevice): Promise; /** * [exported from pre-jsep.js] Register a WebNN Constant operand from external data. @@ -252,7 +266,7 @@ export declare namespace JSEP { * @param shouldConvertInt64ToInt32 - specify whether to convert int64 to int32. * @returns the WebNN Constant operand for the specified external data. */ - jsepRegisterMLConstant( + webnnRegisterMLConstant( externalFilePath: string, dataOffset: number, dataLength: number, @@ -265,14 +279,14 @@ export declare namespace JSEP { * [exported from pre-jsep.js] Register a WebNN graph input. * @param inputName - specify the input name. */ - jsepRegisterGraphInput: (inputName: string) => void; + webnnRegisterGraphInput: (inputName: string) => void; /** * [exported from pre-jsep.js] Check if a graph input is a WebNN graph input. * @param sessionId - specify the session ID. * @param inputName - specify the input name. * @returns whether the input is a WebNN graph input. */ - jsepIsGraphInput: (sessionId: number, inputName: string) => boolean; + webnnIsGraphInput: (sessionId: number, inputName: string) => boolean; /** * [exported from pre-jsep.js] Create a temporary MLTensor for a session. * @param sessionId - specify the session ID. @@ -280,13 +294,13 @@ export declare namespace JSEP { * @param shape - specify the shape. * @returns the MLTensor ID for the temporary MLTensor. */ - jsepCreateTemporaryTensor: (sessionId: number, dataType: DataType, shape: readonly number[]) => Promise; + webnnCreateTemporaryTensor: (sessionId: number, dataType: DataType, shape: readonly number[]) => Promise; /** * [exported from pre-jsep.js] Check if a session's associated WebNN Context supports int64. * @param sessionId - specify the session ID. * @returns whether the WebNN Context supports int64. */ - jsepIsInt64Supported: (sessionId: number) => boolean; + webnnIsInt64Supported: (sessionId: number) => boolean; } } diff --git a/onnxruntime/core/providers/webnn/allocator.cc b/onnxruntime/core/providers/webnn/allocator.cc index 9c5cd651e1f00..8cf5b8cd72a5c 100644 --- a/onnxruntime/core/providers/webnn/allocator.cc +++ b/onnxruntime/core/providers/webnn/allocator.cc @@ -16,7 +16,7 @@ void* WebNNTensorAllocator::Alloc(size_t size) { // We don't need to transfer the tensor to an MLTensor, so we don't need to allocate an MLTensor id. return nullptr; } - void* p = EM_ASM_PTR({ return Module.jsepReserveTensorId(); }); + void* p = EM_ASM_PTR({ return Module.webnnReserveTensorId(); }); allocations_[p] = size; stats_.num_allocs++; stats_.bytes_in_use += SafeInt(size); @@ -27,7 +27,7 @@ void WebNNTensorAllocator::Free(void* p) { if (p == nullptr) { return; } - EM_ASM({ Module.jsepReleaseTensorId($0); }, p); + EM_ASM({ Module.webnnReleaseTensorId($0); }, p); size_t size = allocations_[p]; stats_.bytes_in_use -= size; allocations_.erase(p); diff --git a/onnxruntime/core/providers/webnn/builders/model.cc b/onnxruntime/core/providers/webnn/builders/model.cc index 35964d85862e4..492e2f717e30e 100644 --- a/onnxruntime/core/providers/webnn/builders/model.cc +++ b/onnxruntime/core/providers/webnn/builders/model.cc @@ -157,7 +157,7 @@ onnxruntime::common::Status Model::Compute(const InlinedHashMap& inputs, const InlinedHashMap& outputs) { - auto jsepEnsureTensor = emscripten::val::module_property("jsepEnsureTensor"); + auto webnnEnsureTensor = emscripten::val::module_property("webnnEnsureTensor"); auto promises = emscripten::val::array(); for (const auto& [_, tensor] : inputs) { emscripten::val shape = emscripten::val::array(); @@ -165,7 +165,7 @@ onnxruntime::common::Status Model::Dispatch(const InlinedHashMap(dim); shape.call("push", dim_val); } - auto ml_tensor = jsepEnsureTensor(emscripten::val::undefined(), reinterpret_cast(tensor.buffer), tensor.tensor_info.data_type, shape, true); + auto ml_tensor = webnnEnsureTensor(emscripten::val::undefined(), reinterpret_cast(tensor.buffer), tensor.tensor_info.data_type, shape, true); promises.call("push", ml_tensor); } for (const auto& [_, tensor] : outputs) { @@ -174,7 +174,7 @@ onnxruntime::common::Status Model::Dispatch(const InlinedHashMap(dim); shape.call("push", dim_val); } - auto ml_tensor = jsepEnsureTensor(emscripten::val::undefined(), reinterpret_cast(tensor.buffer), tensor.tensor_info.data_type, shape, false); + auto ml_tensor = webnnEnsureTensor(emscripten::val::undefined(), reinterpret_cast(tensor.buffer), tensor.tensor_info.data_type, shape, false); promises.call("push", ml_tensor); } auto ml_tensors = emscripten::val::global("Promise").call("all", promises).await(); diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.cc b/onnxruntime/core/providers/webnn/builders/model_builder.cc index 661b2ad7056c2..ed6ab7d2d7115 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/model_builder.cc @@ -140,13 +140,13 @@ Status ModelBuilder::RegisterInitializers() { ORT_RETURN_IF_ERROR(utils::GetExternalDataInfo( tensor, graph_viewer_.ModelPath(), external_file_path, data_offset, tensor_byte_size)); - auto jsepRegisterMLConstant = emscripten::val::module_property("jsepRegisterMLConstant"); - operand = jsepRegisterMLConstant(emscripten::val(external_file_path), - static_cast(data_offset), - static_cast(tensor_byte_size), - wnn_builder_, - desc, - should_convert_int64_to_int32); + auto webnnRegisterMLConstant = emscripten::val::module_property("webnnRegisterMLConstant"); + operand = webnnRegisterMLConstant(emscripten::val(external_file_path), + static_cast(data_offset), + static_cast(tensor_byte_size), + wnn_builder_, + desc, + should_convert_int64_to_int32); } else { if (tensor.has_raw_data()) { tensor_ptr = reinterpret_cast(const_cast(tensor.raw_data().c_str())); @@ -288,7 +288,7 @@ Status ModelBuilder::RegisterModelInputOutput(const NodeArg& node_arg, bool is_i desc.set("dataType", emscripten::val("int32")); } wnn_operands_.insert(std::make_pair(name, wnn_builder_.call("input", name, desc))); - emscripten::val::module_property("jsepRegisterGraphInput")(name); + emscripten::val::module_property("webnnRegisterGraphInput")(name); input_names_.push_back(name); } else { output_names_.push_back(name); diff --git a/onnxruntime/core/providers/webnn/data_transfer.cc b/onnxruntime/core/providers/webnn/data_transfer.cc index 44e9bf9edf3d9..aa85277b72453 100644 --- a/onnxruntime/core/providers/webnn/data_transfer.cc +++ b/onnxruntime/core/providers/webnn/data_transfer.cc @@ -29,11 +29,11 @@ common::Status DataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const { const auto& dst_device = dst.Location().device; if (dst_device.Type() == OrtDevice::GPU) { - EM_ASM({ Module.jsepUploadTensor($0, HEAPU8.subarray($1, $1 + $2)); }, dst_data, reinterpret_cast(src_data), bytes); + EM_ASM({ Module.webnnUploadTensor($0, HEAPU8.subarray($1, $1 + $2)); }, dst_data, reinterpret_cast(src_data), bytes); } else { - auto jsepDownloadTensor = emscripten::val::module_property("jsepDownloadTensor"); + auto webnnDownloadTensor = emscripten::val::module_property("webnnDownloadTensor"); auto subarray = emscripten::typed_memory_view(bytes, static_cast(dst_data)); - jsepDownloadTensor(reinterpret_cast(src_data), subarray).await(); + webnnDownloadTensor(reinterpret_cast(src_data), subarray).await(); } } diff --git a/onnxruntime/core/providers/webnn/webnn_execution_provider.cc b/onnxruntime/core/providers/webnn/webnn_execution_provider.cc index 7410ff66add30..c527ba213e55b 100644 --- a/onnxruntime/core/providers/webnn/webnn_execution_provider.cc +++ b/onnxruntime/core/providers/webnn/webnn_execution_provider.cc @@ -284,7 +284,7 @@ class WebNNMemcpy : public OpKernel { explicit WebNNMemcpy(const OpKernelInfo& info) : OpKernel(info) {} Status Compute(OpKernelContext* context) const override { - auto jsepEnsureTensor = emscripten::val::module_property("jsepEnsureTensor"); + auto webnnEnsureTensor = emscripten::val::module_property("webnnEnsureTensor"); const auto* X = context->Input(0); ORT_ENFORCE(X != nullptr, "Memcpy: input tensor is null"); auto* Y = context->Output(0, X->Shape()); @@ -294,10 +294,10 @@ class WebNNMemcpy : public OpKernel { shape.call("push", SafeInt(dim).Ref()); } - jsepEnsureTensor(emscripten::val::undefined(), - reinterpret_cast(Y->MutableDataRaw()), - Y->GetElementType(), - shape, false) + webnnEnsureTensor(emscripten::val::undefined(), + reinterpret_cast(Y->MutableDataRaw()), + Y->GetElementType(), + shape, false) .await(); const auto* data_transfer = Info().GetDataTransferManager().GetDataTransfer(X->Location().device, Y->Location().device); diff --git a/onnxruntime/wasm/pre-jsep.js b/onnxruntime/wasm/pre-jsep.js index a11c6d741d110..cca8da0525fbe 100644 --- a/onnxruntime/wasm/pre-jsep.js +++ b/onnxruntime/wasm/pre-jsep.js @@ -97,41 +97,40 @@ Module["jsepInit"] = (name, params) => { // Functions called via emscripten::val::module_property need to be assigned by name so that the minifier doesn't // change the name. + const backend = params[0]; [ - Module.jsepBackend, - Module.jsepReserveTensorId, - Module.jsepReleaseTensorId, - Module["jsepEnsureTensor"], - Module.jsepUploadTensor, - Module["jsepDownloadTensor"], - ] = params; + Module.webnnReserveTensorId, + Module.webnnReleaseTensorId, + Module["webnnEnsureTensor"], + Module.webnnUploadTensor, + Module["webnnDownloadTensor"], + ] = params.slice(1); // This function is called from both JS and an EM_ASM block, it needs both a minifiable name and an explicit name. - Module["jsepReleaseTensorId"] = Module.jsepReleaseTensorId; - Module["jsepUploadTensor"] = Module.jsepUploadTensor; + Module["webnnReleaseTensorId"] = Module.webnnReleaseTensorId; + Module["webnnUploadTensor"] = Module.webnnUploadTensor; // Functions called from JS also need to have explicit names. - const backend = Module.jsepBackend; - Module["jsepOnRunStart"] = (sessionId) => { + Module["webnnOnRunStart"] = (sessionId) => { return backend["onRunStart"](sessionId); }; - Module["jsepOnRunEnd"] = backend["onRunEnd"].bind(backend); - Module["jsepRegisterMLContext"] = (sessionId, mlContext) => { + Module["webnnOnRunEnd"] = backend["onRunEnd"].bind(backend); + Module["webnnRegisterMLContext"] = (sessionId, mlContext) => { backend["registerMLContext"](sessionId, mlContext); }; - Module["jsepOnReleaseSession"] = (sessionId) => { + Module["webnnOnReleaseSession"] = (sessionId) => { backend["onReleaseSession"](sessionId); }; - Module["jsepCreateMLTensorDownloader"] = (tensorId, type) => { + Module["webnnCreateMLTensorDownloader"] = (tensorId, type) => { return backend["createMLTensorDownloader"](tensorId, type); }; - Module["jsepRegisterMLTensor"] = (sessionId, tensor, dataType, shape) => { + Module["webnnRegisterMLTensor"] = (sessionId, tensor, dataType, shape) => { return backend["registerMLTensor"](sessionId, tensor, dataType, shape); }; - Module["jsepCreateMLContext"] = (optionsOrGpuDevice) => { + Module["webnnCreateMLContext"] = (optionsOrGpuDevice) => { return backend["createMLContext"](optionsOrGpuDevice); }; - Module["jsepRegisterMLConstant"] = ( + Module["webnnRegisterMLConstant"] = ( externalFilePath, dataOffset, dataLength, @@ -149,9 +148,12 @@ Module["jsepInit"] = (name, params) => { shouldConvertInt64ToInt32, ); }; - Module['jsepRegisterGraphInput'] = backend['registerGraphInput'].bind(backend); - Module['jsepIsGraphInput'] = backend['isGraphInput'].bind(backend); - Module['jsepCreateTemporaryTensor'] = backend['createTemporaryTensor'].bind(backend); - Module['jsepIsInt64Supported'] = backend['isInt64Supported'].bind(backend); + Module["webnnRegisterGraphInput"] = + backend["registerGraphInput"].bind(backend); + Module["webnnIsGraphInput"] = backend["isGraphInput"].bind(backend); + + Module["webnnCreateTemporaryTensor"] = + backend["createTemporaryTensor"].bind(backend); + Module["webnnIsInt64Supported"] = backend["isInt64Supported"].bind(backend); } };