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

- Renamed all conflicting WebNN methods from `jsep*` to `webnn*`.
- WebNN doesn't need flush(), therefore it doesn't need to set
`jsepBackend`.

This PR addresses issue microsoft/webnn-developer-preview#78

Author: egalli

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);
       }
     }
 

js/web/lib/wasm/wasm-types.ts

@@ -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.
@@ -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<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]>;
@@ -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<MLContext>;
+    webnnCreateMLContext(optionsOrGpuDevice?: MLContextOptions | GPUDevice): Promise<MLContext>;
 
     /**
      * [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,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;
   }
 }
 

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<int64_t>(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);

onnxruntime/core/providers/webnn/builders/model.cc

@@ -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) {
@@ -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();

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<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()));
@@ -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);

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<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();
     }
   }
 

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<Tensor>(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<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);