[WebGPU EP] Support device-free compile-only sessions for offline graph transformation - #29681
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
Please add comprehensive unit tests to validate the changes. |
055a270 to
3d355b3
Compare
Added in |
7137607 to
65d9401
Compare
65d9401 to
3484b75
Compare
78d0aa1 to
e62af83
Compare
|
I tested the compile-only WebGPU path on WebNN in Chromium locally, I found this PR caused a CPU-fallback regression for mobilenet. The reason is that moving plain-model serialization from I uploaded a new small commit e62af83 to fix. This commit restores the old behavior inside |
Rebuilding the entire graph seems like an extremely expensive operation (I'm a little surprised we do that when creating an EPContext model). Why do we need that for a non-EPContext model? AFAIK we don't do that if using |
I verified again that if I don't rebuild the entire graph, the transpose fallback issue happens when initialize the session during optimized_model_filepath is not our WebNN usage: I am debugging to try to get more details and see whether I can optimize the code. |
|
I did some debugging and compare the dumped model with rebuilding( So let's correct my earlier comment about "the saved model lost value_info on the layout-boundary Transpose nodes", the root cause should be transpose opset mismatch. Rebuilding the model from |
|
We should be preserving the opset import values from the original model. That said I wouldn't have expected the values in Graph::DomainToVersionMap to diverge from the original model, so I think the key thing to establish is where a different ONNX opset comes from. Does the original model use opset 21 or 26? |
WebNN uses model editor API to build a model using opset 21, see https://source.chromium.org/chromium/chromium/src/+/main:services/webnn/ort/model_editor.cc;l=45?q=model_editor I see it should be a bug in |
…ph transformation Enable a WebGPU EP session to run graph transformation (optimization + partitioning) and serialize the resulting model without creating a device or touching hardware. This supports offline / ahead-of-time model compilation in device-free or sandboxed host processes, where the optimized model is produced in one process and executed later in another. WebGPU compile-only mode (new provider option "ep.webgpuexecutionprovider.compileOnly"): * webgpu_provider_options.h, webgpu_provider_factory.cc: parse the option. * webgpu_context.cc/.h: in compile-only mode, skip Dawn adapter/device creation and all device-dependent initialization (queue, limits, features, adapter info, buffer/program managers, etc.), so the context can run graph transformation without a GPU. IsCompileOnly() exposes the state. Why: creating the Dawn device requires real hardware and is wasteful for a session that only produces a compiled model and never executes kernels. Stop-before-finalize, derived from an EP capability (no new public API): * execution_provider.h: add virtual IExecutionProvider::ShouldStopBeforeSessionStateFinalization() (default false); other EPs inherit the default and are unaffected. * webgpu_execution_provider.cc/.h: override it to return the WebGpuContext compile-only state. * inference_session.cc: after graph transformation, Resolve and EPContext model generation, if any registered EP requests it, return before session-state finalization (kernel creation, PrePack, initializer upload, memory planning). The session is left non-runnable; only the output/EPContext model is produced. Why: finalization creates kernels and allocates device memory, which both requires a device and is pure waste for a throwaway compile-only session. Deriving the stop from the EP's own compile-only state, rather than from a separate public session config key, adds no public API surface and makes the two properties inseparable: a compile-only WebGPU EP cannot be misconfigured into attempting finalization. It works identically whether the WebGPU EP is built in (--use_webgpu) or loaded as a plugin (--use_webgpu shared_lib), since both paths produce the same WebGpuExecutionProvider. Lazy allocator initialization: * allocator.cc/.h: GpuBufferAllocator computes mapped_at_creation on the first Alloc() instead of in the constructor. Why: the constructor previously queried the device (BufferManager::SupportsUMA), which is unavailable in compile-only mode. A compile-only session never allocates, so the query is deferred and never runs. Enforce disable_cpu_ep_fallback before the compile-only early-return: * inference_session.cc: move the "session.disable_cpu_ep_fallback" check to run right after graph transformation / Resolve, ahead of the stop-before-finalize early-return. Why: node-to-EP assignment is complete after partitioning, so the check has everything it needs. Running it before the early-return means the guarantee (no node silently falls back to the CPU EP) is enforced for compile-only sessions too, instead of being skipped. Virtual GPU device when hardware enumeration is unavailable: * ep/factory.cc/.h (adapter EP factory) and plugin_ep/ep_factory_webgpu.cc/.h (built-in EP factory): when no GPU OrtHardwareDevice is discovered and the environment sets "allow_virtual_devices=1", register a virtual GPU OrtEpDevice (vendor/device id 0, is_virtual=1) so the WebGPU EP can still be selected. * environment.cc, plugin_ep/ep_factory_internal.h, plugin_ep/ep_factory_internal_impl.h: plumb the allow_virtual_devices flag from the environment to the internal factory before GetSupportedDevices runs. Why: on hosts where OS device enumeration is unavailable (e.g. a sandboxed process under Win32k lockdown), ORT device discovery is skipped and no GPU device is found, so the WebGPU EP would not be selectable. The flag is plumbed rather than read from the OrtEnv singleton inside GetSupportedDevices because internal EPs are registered while the OrtEnv creation mutex is held; querying the singleton there would re-enter that lock and deadlock. Build: delay-load user32.dll (cmake/onnxruntime.cmake, cmake/onnxruntime_providers_webgpu.cmake). Why: the WebGPU EP pulls in user32.dll; delay-loading lets onnxruntime.dll load in processes with Win32k lockdown, where user32.dll is never actually needed because device discovery is skipped. Known limitation (optimization scope): Only basic-level graph transformations (default constant-folding / L1 passes plus layout transformation) are offloaded to the compile-only session. WebGPU EP-specific and higher-level fusions (ORT optimizer levels L2-L4) are not yet captured in the produced model. This is not specific to the WebGPU EP: it is a pre-existing property of the CompileModel (embed) API, which serializes the intermediate graph at the partitioning boundary -- before the L2-L4 optimizer passes run -- so those passes cannot contribute to the serialized output for any EP. The compile-only session therefore uses ORT_ENABLE_BASIC, since any higher level would run L2-L4 without benefit. The longer-term goal is to offload all optimization levels from the compiler process (by adding a post-L4 serialization sink), while the execution process loads the already-optimized model via CreateSessionFromArray with ORT_DISABLE_ALL and thus performs no graph transformation of its own.
A compile-only session does not produce its output via optimized_model_filepath, so it never needs the optimized-model save block in InferenceSession::Initialize(). Return right after skipping session-state finalization instead of falling through to that block and returning after it. This decouples the compile-only path from optimized_model_filepath entirely: the save block now runs only for non-compile-only sessions. Simplify the compile-only tests accordingly by dropping the artificial combination of compile_only and optimized_model_filepath (no real usage sets both): WebGpuCompileOnlySkipsFinalization now only asserts that Initialize() succeeds and Run() fails, and WebGpuVirtualDeviceCompileOnlyEndToEnd only asserts that the compile-only session constructs on the virtual device.
The Compile API can generate an output model even when no nodes are compiled into EPContext nodes (its kGenerateModel action, e.g. for a non-compiling EP such as WebGPU). That model was serialized during partitioning, before the Level2-Level4 optimizer loop, so it was always frozen at ORT_ENABLE_BASIC and the fusions the caller requested never reached the output. It is now emitted after the optimizer loop, so it reflects the graph optimized to the requested level. This is a general fix to the Compile API's model-generation path, not specific to any execution provider. It does not change the output of compiling EPs (OpenVINO/QNN/TensorRT), whose optimization lives inside the EPContext blob rather than in ORT's Level2+ passes. There are two output-model forms, and only the plain one is affected by Level2+: * EPContext model: a compiling EP produced EPContext nodes. Serialized during GraphPartitioner::Partition; its blobs are opaque to Level2+. * Plain optimized model: no nodes were compiled (kGenerateModel). This is a plain copy of the graph and is what the Level2+ passes change. Change: * graph_partitioner.h/.cc: Partition() now serializes only the EPContext form, gated by a new public AnyEpContextNodesProduced(). Removed the old defer_output_model_generation flag and GenerateOutputModel entry point. * inference_session.cc (TransformGraph): emits the plain form after the Level2-L4 loop and before the MemcpyTransformer (so device-placement copy nodes are not baked in), via a new epctx::BuildAndSaveOptimizedModel that honors action_if_no_compiled_nodes (kReturnError / kGenerateModel). * ep_context_utils.h/.cc: moved the shared serialization helpers here -- SaveModelProtoToLocation (writes to a buffer, a user write function, or a file) and GetValidatedEpContextPath. CreateEpContextModel reuses the former for its write step. Adds EP-agnostic tests (CPU EP) on bias_gelu_fusion.onnx, whose Level2-only com.microsoft.BiasGelu fusion is present at ORT_ENABLE_ALL and absent at ORT_ENABLE_BASIC, covering the three output locations (file, in-memory buffer, user write function).
Follow-up cleanups to the compile-only plain-model output path, no behavior change: - Rename emit_plain_optimized_model to emit_plain_compile_only_model so the flag is not confused with the optimized_model_filepath save path. - Move the level >= Level2 serialization point to after all graph transforms (including the training MemoryOptimizer), so a compile-only session's output matches what a plain session serializes to optimized_model_filepath (which happens after TransformGraph fully completes). - Document why a compile-only session intentionally does not return early from TransformGraph or Initialize: the disable_cpu_ep_fallback check must run on the fully transformed graph (including the unconditionally inserted Cast/Memcpy nodes), so returning early could let a session with CPU-assigned nodes slip through. - Drop a redundant nested "#if !defined(ORT_MINIMAL_BUILD)" guard around the second serialization; the whole function is already compiled only in non-minimal builds. - Reword the compile-only skip-finalization log message so it no longer claims an output model is always produced (a compile-only session without EPContext output configured emits nothing), and align the Partition() doc comment with the post-all-transforms serialization point. Also add tests covering the plain compile-only output through the public Compile API (Ort::ModelCompilationOptions + Ort::CompileModel) end-to-end for the file, in-memory buffer, and user-write-func output targets, asserting the emitted model is plain ONNX with no EPContext nodes, plus a test that a requested graph optimization level is reflected in the output. Guard the compile-only serialization tests that assert a contrib op (com.microsoft.BiasGelu) with !defined(DISABLE_CONTRIB_OPS).
A compile-only session skips session-state finalization and returns early from Initialize(), which also bypassed the session-creation telemetry and the end-of-initialization bookkeeping that the normal path emits (the profiler session_initialization event, each EP's OnSessionInitializationEnd(), and LogSessionCreationEnd()). Extract those into two helpers shared by both paths: - LogSessionCreationTelemetry(): the SessionCreation and initial EpDeviceUsage telemetry. - RecordSessionCreationEndTelemetry(): the profiler end event, OnSessionInitializationEnd() per EP, and LogSessionCreationEnd(). The compile-only early return now invokes both before returning, so this bookkeeping is no longer lost. The session is still left un-inited and non-runnable.
…mizedModel The compile-only (kGenerateModel) path serialized the optimized model via BuildAndSaveOptimizedModel, which dumped the live post-partition session graph directly (ToGraphProto on the in-memory model) with no rebuild and no Resolve(). That bare dump captured mid-transformation state: intermediate edges produced by the layout transform were missing value_info, so the layout-boundary Transpose nodes had untyped outputs. When a downstream execution session loaded such a model with graph optimization disabled (as required in the Chromium GPU process, which cannot run any graph transformation), it could not re-partition cleanly, so those Transpose nodes were stranded on the CPU EP and triggered a disable_cpu_ep_fallback error. The regression was introduced when plain-model serialization moved out of CreateEpContextModel (which rebuilt a fresh graph and called Resolve()) into BuildAndSaveOptimizedModel (which did not). This change makes BuildAndSaveOptimizedModel serialize the model exactly the way CreateEpContextModel used to: it rebuilds a fresh, self-consistent Model from the source graph -- same metadata, model path, schema registry, and domain-to-version map -- re-sets inputs/outputs via GetOrCreateNodeArg with their TypeAsProto, AddNodes every node, copies referenced initializers via graph_utils::MakeInitializerCopyIfNotExist, and calls Graph::Resolve() before serialization. Resolve() regenerates complete type/shape information and a canonical graph, so the emitted model is directly, self-consistently re-loadable by an execution session that performs no graph transformation. This is a plain optimized copy with no EPContext-node substitution. With this fix, the compiled model retains fully typed intermediate edges, so the WebGPU execution session assigns the Transpose nodes to the WebGPU EP instead of falling back to CPU.
…ound This reverts the rebuild + Resolve() workaround added in commit e62af83 ("Fix WebGPU CPU fallback by rebuilding + Resolve() in BuildAndSaveOptimizedModel") and fixes the actual root cause instead. Root cause: Model::LoadFromModelEditorApiModel populated model_proto_.opset_import by iterating the schema registry's domain_map and writing each version unconditionally. For the ONNX domain this overwrote the caller-supplied opset (e.g. 21) with the registry's latest/last-released version (e.g. 26), leaving the model-level opset_import inconsistent with the graph's domain-to-version map. That inconsistency caused the CPU fallback: serializing the live model emitted opset "" -> 26, so a layout-boundary Transpose resolved to since_version 25, which the WebGPU EP (kernels up to "since 24") could not claim, stranding the node on CPU and hitting disable_cpu_ep_fallback. Fix: split the loop in LoadFromModelEditorApiModel so it first merges only missing domains into domain_to_version, then populates model_proto_.opset_import from the merged map. The model-level opset imports now stay consistent with the graph and preserve the opsets explicitly provided via the Model Editor API. With the root cause fixed, the rebuild workaround in BuildAndSaveOptimizedModel is unnecessary, so it is reverted: the function again serializes the in-memory model directly.
e62af83 to
8877c42
Compare
|
I uploaded the 8877c42 which reverts the unnecessary
Now this PR can work well with no CPU fallback issues. |
|
The 'Compare' looks off. force-pushed the pr-webgpu-compile-only branch from e62af83 to 8877c42 Showing with 18,091 additions and 3,508 deletions. |
Overall diffs look okay though so I'll ignore that. |
I have to force push because my branch has conficts with the main branch, so I |
|
Thanks so much for the detailed review and your patience in getting this landed! |
Description
Enable a compile-only session to run graph transformation (optimization + partitioning) and serialize
the optimized model without creating a device or touching hardware. This supports offline /
ahead-of-time compilation, where the optimized model is produced in a device-free (or sandboxed)
process and executed later in a separate device-capable process. The serialized output reflects the
requested graph-optimization level, including L2–L4 fusions.
The feature targets WebGPU offline compilation and is driven by the Compile API's existing internal
session.compile_only(kOrtSessionOptionCompileOnly) — no new public config key. The WebGPU-specificparts are the device-free context, the no-op allocator, and the virtual device; two of the mechanisms are
EP-agnostic Compile-API improvements that benefit any EP: skipping session-state finalization (and returning
early) for a compile-only session, and capturing the requested optimization level in the generated model.
Changes
webgpu_provider_factory.cc,webgpu_context.cc/.h): in acompile-only session, skip Dawn adapter/device creation and all device-dependent init.
HasDevice()exposes whether a real device exists; the mechanical "device-free" concept stays separate from the
session concept "compile-only".
inference_session.cc): a compile-only session skipsFinalizeSessionState()(kernel creation, initializer upload, memory planning — all need a deviceand are wasted for a throwaway session) and returns before the
optimized_model_filepathsave block. Its output is produced during graph transformation, not via that block.
inference_session.cc,graph_partitioner.cc/.h,ep_context_utils.cc/.h): the plain (non-EPContext) optimized model isserialized after the L2–L4 optimizer passes, so the emitted graph reflects the requested level
(previously it was frozen at the partition boundary, before L2–L4). Serialization point is chosen by
level: before the loop for
< Level2, after all transforms for>= Level2. This is a generalCompile-API fix; it does not change the output of compiling EPs, whose optimization lives in the
EPContext blob.
allocator.cc/.h): every device-allocator constructionsite hands out
WebGpuNoOpAllocatorwhen there is no device, so ORT never builds a real allocatorwithout one.
ep/factory.cc/.h, internalWebGpuEpFactory): when no GPUOrtHardwareDeviceis discovered andallow_virtual_devices=1, register a virtual GPUOrtEpDeviceso the WebGPU EP stays selectable on hosts where device enumeration is unavailable (e.g. a sandbox).
user32.dll(cmake/onnxruntime.cmake,cmake/onnxruntime_providers_webgpu.cmake) so the EP loads whereuser32.dllis unavailable butnever actually needed (device discovery skipped).
Tests (
test/framework/inference_session_test.cc)Initialize()succeeds,Run()fails (non-runnable); no-opallocator handed out; virtual-device registration is deterministic with or without a GPU.
Ort::ModelCompilationOptions+Ort::CompileModel) for the file, in-memory buffer, and user-write-func targets: asserts the emittedmodel is plain ONNX with no EPContext nodes, and that the requested optimization level is reflected in
the output.
Motivation and Context
Since the WebNN Compiler process is sandboxed, graph compilation within it must be offline /
device-free: the untrusted model is parsed and graph-transformed in a low-privilege, device-free
process, and the resulting model is executed later in a separate device-ful process. This is a
security offload — it keeps untrusted-graph processing out of the high-privilege GPU process.