diff --git a/js/common/lib/inference-session.ts b/js/common/lib/inference-session.ts index 71a5912df2..8c1e69a68c 100644 --- a/js/common/lib/inference-session.ts +++ b/js/common/lib/inference-session.ts @@ -192,6 +192,7 @@ export declare namespace InferenceSession { wasm: WebAssemblyExecutionProviderOption; webgl: WebGLExecutionProviderOption; xnnpack: XnnpackExecutionProviderOption; + webgpu: WebGpuExecutionProviderOption; webnn: WebNNExecutionProviderOption; nnapi: NnapiExecutionProviderOption; } @@ -233,6 +234,10 @@ export declare namespace InferenceSession { export interface XnnpackExecutionProviderOption extends ExecutionProviderOption { readonly name: 'xnnpack'; } + export interface WebGpuExecutionProviderOption extends ExecutionProviderOption { + readonly name: 'webgpu'; + preferredLayout?: 'NCHW'|'NHWC'; + } export interface WebNNExecutionProviderOption extends ExecutionProviderOption { readonly name: 'webnn'; deviceType?: 'cpu'|'gpu'; diff --git a/js/web/lib/wasm/session-options.ts b/js/web/lib/wasm/session-options.ts index 2659b47173..02ff229cc4 100644 --- a/js/web/lib/wasm/session-options.ts +++ b/js/web/lib/wasm/session-options.ts @@ -88,6 +88,21 @@ const setExecutionProviders = break; case 'webgpu': epName = 'JS'; + if (typeof ep !== 'string') { + const webgpuOptions = ep as InferenceSession.WebGpuExecutionProviderOption; + if (webgpuOptions?.preferredLayout) { + if (webgpuOptions.preferredLayout !== 'NCHW' && webgpuOptions.preferredLayout !== 'NHWC') { + throw new Error(`preferredLayout must be either 'NCHW' or 'NHWC': ${webgpuOptions.preferredLayout}`); + } + const keyDataOffset = allocWasmString('preferredLayout', allocs); + const valueDataOffset = allocWasmString(webgpuOptions.preferredLayout, allocs); + if (getInstance()._OrtAddSessionConfigEntry(sessionOptionsHandle, keyDataOffset, valueDataOffset) !== + 0) { + checkLastError( + `Can't set a session config entry: 'preferredLayout' - ${webgpuOptions.preferredLayout}.`); + } + } + } break; case 'wasm': case 'cpu': diff --git a/onnxruntime/core/providers/js/js_execution_provider.cc b/onnxruntime/core/providers/js/js_execution_provider.cc index ae33fb752f..444f50958e 100644 --- a/onnxruntime/core/providers/js/js_execution_provider.cc +++ b/onnxruntime/core/providers/js/js_execution_provider.cc @@ -624,7 +624,8 @@ std::unique_ptr RegisterKernels() { using namespace js; JsExecutionProvider::JsExecutionProvider(const JsExecutionProviderInfo& info) - : IExecutionProvider{kJsExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, 0), true} { + : IExecutionProvider{kJsExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, 0), true}, + preferred_data_layout_{info.data_layout} { } std::vector JsExecutionProvider::CreatePreferredAllocators() { diff --git a/onnxruntime/core/providers/js/js_execution_provider.h b/onnxruntime/core/providers/js/js_execution_provider.h index 091aa29046..39d43498c0 100644 --- a/onnxruntime/core/providers/js/js_execution_provider.h +++ b/onnxruntime/core/providers/js/js_execution_provider.h @@ -19,12 +19,21 @@ KernelCreateInfo BuildKernelCreateInfo(); } // namespace js -// placeholder for future use. no options currently struct JsExecutionProviderInfo { - JsExecutionProviderInfo() = default; - JsExecutionProviderInfo(const ProviderOptions& po) { + auto it = po.find("preferred_layout"); + if (it != po.end()) { + auto& value = it->second; + if (value == "NCHW") { + data_layout = DataLayout::NCHW; + } else if (value == "NHWC") { + data_layout = DataLayout::NHWC; + } + } } + + // JSEP default preferred layout is NHWC + DataLayout data_layout = DataLayout::NHWC; }; class JsExecutionProvider : public IExecutionProvider { @@ -39,7 +48,7 @@ class JsExecutionProvider : public IExecutionProvider { std::shared_ptr GetKernelRegistry() const override; std::unique_ptr GetDataTransfer() const override; - DataLayout GetPreferredLayout() const override { return DataLayout::NHWC; } + DataLayout GetPreferredLayout() const override { return preferred_data_layout_; } FusionStyle GetFusionStyle() const override { return FusionStyle::FilteredGraphViewer; } @@ -48,6 +57,7 @@ class JsExecutionProvider : public IExecutionProvider { bool ConcurrentRunSupported() const override { return false; } std::vector CreatePreferredAllocators() override; + DataLayout preferred_data_layout_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/session/provider_registration.cc b/onnxruntime/core/session/provider_registration.cc index 9326c6eaff..50c9f6681a 100644 --- a/onnxruntime/core/session/provider_registration.cc +++ b/onnxruntime/core/session/provider_registration.cc @@ -109,6 +109,10 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, #endif } else if (strcmp(provider_name, "JS") == 0) { #if defined(USE_JSEP) + std::string preferred_layout; + if (options->value.config_options.TryGetConfigEntry("preferredLayout", preferred_layout)) { + provider_options["preferred_layout"] = preferred_layout; + } options->provider_factories.push_back(JsProviderFactoryCreator::Create(provider_options)); #else status = create_not_supported_status();