Skip to content

[WebNN/WebGPU JS] Fix shared Module methods overriding each other #23998

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fs-eire merged 2 commits into from
Mar 19, 2025
Merged

[WebNN/WebGPU JS] Fix shared Module methods overriding each other #23998

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a60aa10
[WebNN/WebGPU JS] Fix shared Module methods overriding each other
egalli Mar 18, 2025
d1118af
Renamed all WebNN jsep* methods to webnn*
egalli Mar 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 15 additions & 13 deletions js/web/lib/wasm/wasm-core-impl.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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.');
}
Expand All @@ -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);
Expand All @@ -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.');
}
Expand Down Expand Up @@ -722,6 +723,7 @@ export const run = async (
}

wasm.jsepOnRunStart?.(sessionHandle);
wasm.webnnOnRunStart?.(sessionHandle);

let errorCode: number;
if (!BUILD_DEFS.DISABLE_JSEP && ioBindingState) {
Expand Down Expand Up @@ -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.');
}
Expand All @@ -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);
},
},
Expand All @@ -915,7 +917,7 @@ export const run = async (
if (!keepOutputTensor) {
wasm._OrtReleaseTensor(tensor);
}
wasm.jsepOnRunEnd?.(sessionHandle);
wasm.webnnOnRunEnd?.(sessionHandle);
}
}

Expand Down
44 changes: 29 additions & 15 deletions js/web/lib/wasm/wasm-types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,31 +156,45 @@ 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.
* @param sessionId - specify the session ID.
* @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.
Expand All @@ -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,
Expand All @@ -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<undefined>;
webnnDownloadTensor: (tensorId: number, dstBuffer: ArrayBufferView | ArrayBuffer) => Promise<undefined>;
/**
* [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<Tensor.DataTypeMap[Tensor.MLTensorDataTypes]>;
Expand All @@ -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,
Expand All @@ -240,7 +254,7 @@ export declare namespace JSEP {
* @param optionsOrGpuDevice - specify the options or GPUDevice.
* @returns
*/
jsepCreateMLContext(optionsOrGpuDevice?: MLContextOptions | GPUDevice): Promise<MLContext>;
webnnCreateMLContext(optionsOrGpuDevice?: MLContextOptions | GPUDevice): Promise<MLContext>;

/**
* [exported from pre-jsep.js] Register a WebNN Constant operand from external data.
Expand All @@ -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,
Expand All @@ -265,28 +279,28 @@ 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.
* @param dataType - specify the data type.
* @param shape - specify the shape.
* @returns the MLTensor ID for the temporary MLTensor.
*/
jsepCreateTemporaryTensor: (sessionId: number, dataType: DataType, shape: readonly number[]) => Promise<number>;
webnnCreateTemporaryTensor: (sessionId: number, dataType: DataType, shape: readonly number[]) => Promise<number>;
/**
* [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;
}
}

Expand Down
4 changes: 2 additions & 2 deletions onnxruntime/core/providers/webnn/allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int64_t>(size);
Expand All @@ -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);
Expand Down
6 changes: 3 additions & 3 deletions onnxruntime/core/providers/webnn/builders/model.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,15 +157,15 @@ onnxruntime::common::Status Model::Compute(const InlinedHashMap<std::string, Onn

onnxruntime::common::Status Model::Dispatch(const InlinedHashMap<std::string, OnnxTensorData>& inputs,
const InlinedHashMap<std::string, OnnxTensorData>& 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();
for (const auto& dim : tensor.tensor_info.shape) {
uint32_t dim_val = SafeInt<uint32_t>(dim);
shape.call<void>("push", dim_val);
}
auto ml_tensor = jsepEnsureTensor(emscripten::val::undefined(), reinterpret_cast<intptr_t>(tensor.buffer), tensor.tensor_info.data_type, shape, true);
auto ml_tensor = webnnEnsureTensor(emscripten::val::undefined(), reinterpret_cast<intptr_t>(tensor.buffer), tensor.tensor_info.data_type, shape, true);
promises.call<void>("push", ml_tensor);
}
for (const auto& [_, tensor] : outputs) {
Expand All @@ -174,7 +174,7 @@ onnxruntime::common::Status Model::Dispatch(const InlinedHashMap<std::string, On
uint32_t dim_val = SafeInt<uint32_t>(dim);
shape.call<void>("push", dim_val);
}
auto ml_tensor = jsepEnsureTensor(emscripten::val::undefined(), reinterpret_cast<intptr_t>(tensor.buffer), tensor.tensor_info.data_type, shape, false);
auto ml_tensor = webnnEnsureTensor(emscripten::val::undefined(), reinterpret_cast<intptr_t>(tensor.buffer), tensor.tensor_info.data_type, shape, false);
promises.call<void>("push", ml_tensor);
}
auto ml_tensors = emscripten::val::global("Promise").call<emscripten::val>("all", promises).await();
Expand Down
16 changes: 8 additions & 8 deletions onnxruntime/core/providers/webnn/builders/model_builder.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int32_t>(data_offset),
static_cast<int32_t>(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<int32_t>(data_offset),
static_cast<int32_t>(tensor_byte_size),
wnn_builder_,
desc,
should_convert_int64_to_int32);
} else {
if (tensor.has_raw_data()) {
tensor_ptr = reinterpret_cast<std::byte*>(const_cast<char*>(tensor.raw_data().c_str()));
Expand Down Expand Up @@ -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<emscripten::val>("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);
Expand Down
6 changes: 3 additions & 3 deletions onnxruntime/core/providers/webnn/data_transfer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<intptr_t>(src_data), bytes);
EM_ASM({ Module.webnnUploadTensor($0, HEAPU8.subarray($1, $1 + $2)); }, dst_data, reinterpret_cast<intptr_t>(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<char*>(dst_data));
jsepDownloadTensor(reinterpret_cast<intptr_t>(src_data), subarray).await();
webnnDownloadTensor(reinterpret_cast<intptr_t>(src_data), subarray).await();
}
}

Expand Down
10 changes: 5 additions & 5 deletions onnxruntime/core/providers/webnn/webnn_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Tensor>(0);
ORT_ENFORCE(X != nullptr, "Memcpy: input tensor is null");
auto* Y = context->Output(0, X->Shape());
Expand All @@ -294,10 +294,10 @@ class WebNNMemcpy : public OpKernel {
shape.call<void>("push", SafeInt<uint32_t>(dim).Ref());
}

jsepEnsureTensor(emscripten::val::undefined(),
reinterpret_cast<intptr_t>(Y->MutableDataRaw()),
Y->GetElementType(),
shape, false)
webnnEnsureTensor(emscripten::val::undefined(),
reinterpret_cast<intptr_t>(Y->MutableDataRaw()),
Y->GetElementType(),
shape, false)
.await();

const auto* data_transfer = Info().GetDataTransferManager().GetDataTransfer(X->Location().device, Y->Location().device);
Expand Down
Loading