Repository: microsoft/Olive
Pull request #1854: Basic support for ONNX auto EP selection introduced in onnxruntime v1.22.0
Author: login: shaahji
Created: 2025-05-13T21:48:11Z | Merged: 2025-05-14T19:52:45Z

## Basic support for ONNX auto EP selection introduced in onnxruntime v1.22.0

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [x] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [x] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this change to be included in the release notes.
- [ ] Is this PR including examples changes? If yes, please remember to update [example documentation](https://github.com/microsoft/Olive/blob/main/docs/source/examples.md) in a follow-up PR.

## (Optional) Issue link



Diff:
diff --git a/olive/common/ort_inference.py b/olive/common/ort_inference.py
index 6c13d1b4d4..34441ef72f 100644
--- a/olive/common/ort_inference.py
+++ b/olive/common/ort_inference.py
@@ -13,11 +13,13 @@
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import numpy as np
+from packaging import version
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
     from onnxruntime import InferenceSession, IOBinding
 
+    from olive.hardware.accelerator import Device
 
 logger = logging.getLogger(__name__)
 
@@ -26,11 +28,73 @@ class OrtSessionFallbackError(Exception):
     """Raised when the onnxruntime fallback happens."""
 
 
+def is_winml_installation() -> bool:
+    from onnxruntime import __version__ as OrtVersion
+
+    if version.parse(OrtVersion) >= version.parse("1.22.0"):
+        try:
+            from onnxruntime import winml  # noqa: F401 # pylint: disable=unused-import
+        except ImportError:
+            logger.info("onnxruntime-winml not installed")
+            return False
+        return True
+    return False
+
+
+def get_ort_hardware_device_type(device: Union["Device", str]):
+    if is_winml_installation():
+        from onnxruntime import OrtHardwareDeviceType
+
+        mapping = {
+            "cpu": OrtHardwareDeviceType.CPU,
+            "gpu": OrtHardwareDeviceType.GPU,
+            "npu": OrtHardwareDeviceType.NPU,
+        }
+        return mapping.get(device.lower())
+    return None
+
+
+def get_ort_execution_provider_device_policy(policy: str):
+    if is_winml_installation():
+        from onnxruntime import OrtExecutionProviderDevicePolicy
+
+        mapping = {
+            "default": OrtExecutionProviderDevicePolicy.DEFAULT,
+            "prefer_cpu": OrtExecutionProviderDevicePolicy.PREFER_CPU,
+            "prefer_npu": OrtExecutionProviderDevicePolicy.PREFER_NPU,
+            "prefer_gpu": OrtExecutionProviderDevicePolicy.PREFER_GPU,
+            "max_performance": OrtExecutionProviderDevicePolicy.MAX_PERFORMANCE,
+            "max_efficiency": OrtExecutionProviderDevicePolicy.MAX_EFFICIENCY,
+            "overall_power": OrtExecutionProviderDevicePolicy.MIN_OVERALL_POWER,
+        }
+        return mapping.get(policy.lower())
+    return None
+
+
+def initialize_inference_session_options_for_winml(
+    sess_options, device, providers, provider_options, provider_selection_policy=None
+):
+    import onnxruntime as ort
+
+    providers = providers or []
+    provider_options = provider_options or []
+    provider_options_by_ep = dict(zip(providers, provider_options))
+    ort_device_type = get_ort_hardware_device_type(device)
+    for ep_device in ort.get_ep_devices():
+        if ep_device.device.type == ort_device_type and ep_device.ep_name in provider_options_by_ep:
+            sess_options.add_provider_for_devices([ep_device], provider_options_by_ep.get(ep_device.ep_name) or {})
+
+    if provider_selection_policy:
+        provider_selection_policy = get_ort_execution_provider_device_policy(provider_selection_policy)
+        sess_options.set_provider_selection_policy(provider_selection_policy)
+
+
 # NOTE: `device_id` is only used internally for inference with Distributed ONNX models.
 # For regular ONNX models, the recommended way to specify the device is to set the environment variable
-# `CUDA_VISIBLE_DEVICES` before runnning a workflow.
+# `CUDA_VISIBLE_DEVICES` before running a workflow.
 def get_ort_inference_session(
     model_path: Union[Path, str],
+    device: Union["Device", str],
     inference_settings: dict[str, Any],
     use_ort_extensions: bool = False,
     device_id: Optional[int] = None,
@@ -119,11 +183,18 @@ def get_ort_inference_session(
     if len(providers) >= 1 and providers[0] == "DmlExecutionProvider":
         sess_options.enable_mem_pattern = False
 
+    sess_kwargs = {}
+    if is_winml_installation():
+        initialize_inference_session_options_for_winml(
+            sess_options, device, providers, provider_options, inference_settings.get("provider_selection_policy")
+        )
+    else:
+        sess_kwargs.update({"providers": providers, "provider_options": provider_options})
+
     # create session
-    session = ort.InferenceSession(
-        str(model_path), sess_options=sess_options, providers=providers, provider_options=provider_options
-    )
+    session = ort.InferenceSession(str(model_path), sess_options=sess_options, **sess_kwargs)
     check_ort_fallback(session, providers)
+
     # set tuning results for tunable operators (currently only for ROCM EP)
     tuning_op_result = inference_settings.get("tuning_op_result")
     if tuning_op_result:
diff --git a/olive/model/handler/mixin/onnx_ep.py b/olive/model/handler/mixin/onnx_ep.py
index 77bc89c99f..fc5de78a90 100644
--- a/olive/model/handler/mixin/onnx_ep.py
+++ b/olive/model/handler/mixin/onnx_ep.py
@@ -15,6 +15,8 @@ def is_valid_ep(self, filepath: str, ep: str = None):
         # It should be a bug for onnxruntime where the execution provider is not be fallback.
         import onnxruntime as ort
 
+        from olive.common.ort_inference import is_winml_installation
+
         sess_options = ort.SessionOptions()
         if self.use_ort_extensions:
             # register custom ops for onnxruntime_extensions
@@ -22,14 +24,26 @@ def is_valid_ep(self, filepath: str, ep: str = None):
 
             sess_options.register_custom_ops_library(get_library_path())
 
-        try:
-            ort.InferenceSession(filepath, sess_options, providers=[ep])
-        except Exception as e:  # pylint: disable=broad-except
+        valid = True
+        e = None
+        if is_winml_installation():
+            sess_options.add_provider_for_devices(
+                [ep_device for ep_device in ort.get_ep_devices() if ep_device.ep_name == ep], {}
+            )
+            valid = sess_options.has_providers()
+        else:
+            try:
+                ort.InferenceSession(filepath, sess_options, providers=[ep])
+            except Exception as ex:  # pylint: disable=broad-except
+                valid = False
+                e = ex
+
+        if not valid:
             logger.warning(
-                "Error: %s Olive will ignore this %(ep)s."
-                "Please make sure the environment with %(ep)s has the required dependencies.",
+                "Error: %s Olive will ignore this %s."
+                "Please make sure the environment with %s has the required dependencies.",
                 e,
-                ep=ep,
+                ep,
+                ep,
             )
-            return False
-        return True
+        return valid
diff --git a/olive/model/handler/onnx.py b/olive/model/handler/onnx.py
index 5c2e4699f8..803266cb2c 100644
--- a/olive/model/handler/onnx.py
+++ b/olive/model/handler/onnx.py
@@ -129,7 +129,7 @@ def prepare_session(
 
         try:
             return get_ort_inference_session(
-                self.model_path, inference_settings, self.use_ort_extensions, device_id, external_initializers
+                self.model_path, device, inference_settings, self.use_ort_extensions, device_id, external_initializers
             )
         except OrtSessionFallbackError as e:
             raise OliveEvaluationError(e) from e
diff --git a/olive/passes/onnx/context_binary.py b/olive/passes/onnx/context_binary.py
index 24154ef69a..1d29134d4e 100644
--- a/olive/passes/onnx/context_binary.py
+++ b/olive/passes/onnx/context_binary.py
@@ -10,7 +10,7 @@
 
 from packaging import version
 
-from olive.hardware.accelerator import AcceleratorSpec
+from olive.hardware.accelerator import AcceleratorSpec, Device
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
@@ -93,6 +93,7 @@ def _run_for_config(
             return self._generate_context_binary(
                 model_path=model.model_path,
                 output_model_path=resolve_onnx_path(output_model_path, f"{Path(model.model_path).stem}_ctx.onnx"),
+                device=self.accelerator_spec.accelerator_type,
                 **generate_kwargs,
             )
 
@@ -122,6 +123,7 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
                             new_ctx_model_name: component_models[ctx_model_name].model_path,
                             new_iter_model_name: component_models[iter_model_name].model_path,
                         },
+                        device=self.accelerator_spec.accelerator_type,
                         output_model_dir=output_dir,
                         generate_kwargs=generate_kwargs,
                         weight_sharing=config.weight_sharing,
@@ -150,6 +152,7 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
         new_component_models = self._generate_composite_binaries(
             model_paths_map={f"{name}_ctx": component.model_path for name, component in component_map.items()},
             output_model_dir=output_model_path,
+            device=self.accelerator_spec.accelerator_type,
             generate_kwargs=generate_kwargs,
             weight_sharing=config.weight_sharing,
         )
@@ -164,6 +167,7 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
     def _generate_composite_binaries(
         cls,
         model_paths_map: dict[str, str],
+        device: Union[Device, str],
         output_model_dir: str,
         generate_kwargs: dict[str, str],
         weight_sharing: bool = False,
@@ -188,6 +192,7 @@ def _generate_composite_binaries(
             new_models[model_name] = cls._generate_context_binary(
                 model_path=model_path,
                 output_model_path=Path(output_model_dir) / f"{model_name}.onnx",
+                device=device,
                 **generate_kwargs,
             )
 
@@ -197,6 +202,7 @@ def _generate_composite_binaries(
     def _generate_context_binary(
         model_path: str,
         output_model_path: Union[str, Path],
+        device: Union[Device, str],
         execution_provider: str,
         provider_options: Optional[dict] = None,
         session_options: Optional[dict] = None,
@@ -219,7 +225,9 @@ def _generate_context_binary(
         :param ignore_missing_cb_bin: Whether to ignore missing context binary files.
         :return: ONNXModelHandler for the generated context binary.
         """
-        from onnxruntime import InferenceSession, SessionOptions
+        import onnxruntime as ort
+
+        from olive.common.ort_inference import initialize_inference_session_options_for_winml, is_winml_installation
 
         # prepare provider options
         provider_options = provider_options or {}
@@ -239,7 +247,7 @@ def _generate_context_binary(
                 "ep.stop_share_ep_contexts": int(stop_share_ep_contexts),
             }
         )
-        sess_options = SessionOptions()
+        sess_options = ort.SessionOptions()
         for key, value in session_options.items():
             sess_options.add_session_config_entry(key, str(value))
 
@@ -257,12 +265,15 @@ def _generate_context_binary(
 
         # create the inference session
         logger.debug("Creating context binary for model %s", str(model_path))
-        InferenceSession(
-            model_path,
-            sess_options=sess_options,
-            providers=[execution_provider],
-            provider_options=[provider_options],
-        )
+
+        sess_kwargs = {}
+        if is_winml_installation():
+            initialize_inference_session_options_for_winml(
+                sess_options, device, [execution_provider], [provider_options or {}]
+            )
+        else:
+            sess_kwargs.update({"providers": [execution_provider], "provider_options": [provider_options]})
+        ort.InferenceSession(model_path, sess_options=sess_options, **sess_kwargs)
 
         assert output_model_path.exists(), f"Context binary not found at {output_model_path}"
 
diff --git a/olive/passes/onnx/optimum_merging.py b/olive/passes/onnx/optimum_merging.py
index 4116ebfe58..4f711cc892 100644
--- a/olive/passes/onnx/optimum_merging.py
+++ b/olive/passes/onnx/optimum_merging.py
@@ -44,7 +44,9 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
     def _run_for_config(
         self, model: CompositeModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:
-        import onnxruntime
+        import onnxruntime as ort
+
+        from olive.common.ort_inference import initialize_inference_session_options_for_winml, is_winml_installation
 
         assert len(model.model_components) == 2
 
@@ -75,11 +77,18 @@ def new_byte_size_func(_):
         olive_model = model_proto_to_olive_model(merged_model, output_model_path, config)
 
         # Doing a dry run of ORT allows us to remove the initializers that were orphaned by the merging step
-        sess_options = onnxruntime.SessionOptions()
-        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
         sess_options.optimized_model_filepath = output_model_path
 
         execution_provider = self.accelerator_spec.execution_provider
-        onnxruntime.InferenceSession(output_model_path, sess_options, providers=[execution_provider])
+        sess_kwargs = {}
+        if is_winml_installation():
+            initialize_inference_session_options_for_winml(
+                sess_options, self.accelerator_spec.accelerator_type, [execution_provider], [{}]
+            )
+        else:
+            sess_kwargs = {"providers": [execution_provider]}
+        ort.InferenceSession(output_model_path, sess_options, **sess_kwargs)
 
         return olive_model
diff --git a/olive/systems/isolated_ort/inference_runner.py b/olive/systems/isolated_ort/inference_runner.py
index 5d0edd1f4f..9debb856bf 100644
--- a/olive/systems/isolated_ort/inference_runner.py
+++ b/olive/systems/isolated_ort/inference_runner.py
@@ -49,6 +49,7 @@ def main(raw_args=None):
     # create session
     session = get_ort_inference_session(
         args.model_path,
+        config.get("device", "cpu"),
         config["inference_settings"],
         config.get("use_ort_extensions", False),
         external_initializers=np.load(args.external_initializers_path) if args.external_initializers_path else None,
diff --git a/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py b/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py
index 7d9e7815a3..d603c74821 100644
--- a/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py
+++ b/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py
@@ -234,7 +234,7 @@ def test_inference_runner_with_run(self, mock_wrapper_class, mock_get_session, t
 
         # assert
         mock_get_session.assert_called_once_with(
-            Path(model), config["inference_settings"], False, external_initializers=None
+            Path(model), "cpu", config["inference_settings"], False, external_initializers=None
         )
         mock_wrapper_class.assert_called_once_with(
             mock_get_session.return_value,


Reviews:
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T03:30:14Z
user: login: xiaoyu-work
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T03:55:24Z
user: login: shaahji
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T05:41:27Z
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T05:45:49Z
user: login: shaahji
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T05:52:31Z
user: login: xiaoyu-work
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-14T18:26:05Z

Review comments:
diff_hunk: @@ -26,11 +28,38 @@
     """Raised when the onnxruntime fallback happens."""
 
 
+def is_winml_installation() -> bool:
+    from onnxruntime import __version__ as OrtVersion
+
+    if version.parse(OrtVersion) >= version.parse("1.22.0"):
+        try:
+            from onnxruntime import winml
path: olive/common/ort_inference.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0611

Unused winml imported from onnxruntime (unused-import)
See [unused-import](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/unused-import.html).

[Show more details](https://github.com/microsoft/Olive/security/code-scanning/16527)
created_at: 2025-05-14T03:30:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -26,11 +28,38 @@
     """Raised when the onnxruntime fallback happens."""
 
 
+def is_winml_installation() -> bool:
+    from onnxruntime import __version__ as OrtVersion
+
+    if version.parse(OrtVersion) >= version.parse("1.22.0"):
+        try:
+            from onnxruntime import winml
path: olive/common/ort_inference.py
user: login: github-advanced-security[bot]
body: ## RUFF/F401

`onnxruntime.winml` imported but unused; consider using `importlib.util.find_spec` to test for availability.
See https://docs.astral.sh/ruff/rules/unused-import

[Show more details](https://github.com/microsoft/Olive/security/code-scanning/16526)
created_at: 2025-05-14T03:30:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -75,11 +77,23 @@ def new_byte_size_func(_):
         olive_model = model_proto_to_olive_model(merged_model, output_model_path, config)
 
         # Doing a dry run of ORT allows us to remove the initializers that were orphaned by the merging step
-        sess_options = onnxruntime.SessionOptions()
-        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
         sess_options.optimized_model_filepath = output_model_path
 
         execution_provider = self.accelerator_spec.execution_provider
-        onnxruntime.InferenceSession(output_model_path, sess_options, providers=[execution_provider])
+        sess_kwargs = {}
+        if is_winml_installation():
+            ort_device_type = get_ort_hardware_device_type(self.accelerator_spec.accelerator_type)
+            sess_options.add_provider_for_devices(
+                [
+                    ep_device for ep_device in ort.get_ep_devices()
+                    if ep_device.device.type == ort_device_type and ep_device.ep_name == execution_provider
+                ],
+                {},
+            )
+        else:
+            sess_kwargs = {"providers": [execution_provider]}
+        ort.InferenceSession(output_model_path, sess_options, **sess_kwargs)
path: olive/passes/onnx/optimum_merging.py
user: login: xiaoyu-work
body: There are 3 similar code blocks in this PR. Could they be abstracted into a common function?
created_at: 2025-05-14T03:55:24Z
author_association: COLLABORATOR
diff_hunk: @@ -75,11 +77,23 @@ def new_byte_size_func(_):
         olive_model = model_proto_to_olive_model(merged_model, output_model_path, config)
 
         # Doing a dry run of ORT allows us to remove the initializers that were orphaned by the merging step
-        sess_options = onnxruntime.SessionOptions()
-        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
         sess_options.optimized_model_filepath = output_model_path
 
         execution_provider = self.accelerator_spec.execution_provider
-        onnxruntime.InferenceSession(output_model_path, sess_options, providers=[execution_provider])
+        sess_kwargs = {}
+        if is_winml_installation():
+            ort_device_type = get_ort_hardware_device_type(self.accelerator_spec.accelerator_type)
+            sess_options.add_provider_for_devices(
+                [
+                    ep_device for ep_device in ort.get_ep_devices()
+                    if ep_device.device.type == ort_device_type and ep_device.ep_name == execution_provider
+                ],
+                {},
+            )
+        else:
+            sess_kwargs = {"providers": [execution_provider]}
+        ort.InferenceSession(output_model_path, sess_options, **sess_kwargs)
path: olive/passes/onnx/optimum_merging.py
user: login: shaahji
body: Consolidated.
created_at: 2025-05-14T05:41:27Z
author_association: COLLABORATOR
diff_hunk: @@ -26,11 +28,72 @@
     """Raised when the onnxruntime fallback happens."""
 
 
+def is_winml_installation() -> bool:
+    from onnxruntime import __version__ as OrtVersion
+
+    if version.parse(OrtVersion) >= version.parse("1.22.0"):
+        try:
+            from onnxruntime import winml   # noqa: F401 # pylint: disable=unused-import
path: olive/common/ort_inference.py
user: login: github-advanced-security[bot]
body: ## RUFF/I001

Import block is un-sorted or un-formatted.
See https://docs.astral.sh/ruff/rules/unsorted-imports

[Show more details](https://github.com/microsoft/Olive/security/code-scanning/16528)
created_at: 2025-05-14T05:45:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -26,11 +28,72 @@
     """Raised when the onnxruntime fallback happens."""
 
 
+def is_winml_installation() -> bool:
+    from onnxruntime import __version__ as OrtVersion
+
+    if version.parse(OrtVersion) >= version.parse("1.22.0"):
+        try:
+            from onnxruntime import winml   # noqa: F401 # pylint: disable=unused-import
path: olive/common/ort_inference.py
user: login: shaahji
body: Spurious warning
created_at: 2025-05-14T05:52:31Z
author_association: COLLABORATOR

PR comments:
user: login: xiaoyu-work
created_at: 2025-05-14T03:48:18Z
body: Can you elaborate more about this change?
author_association: COLLABORATOR
user: login: shaahji
created_at: 2025-05-14T05:45:08Z
body: > Can you elaborate more about this change?

WinML uses the auto ep selection feature from onnxruntime. It's a new feature introduced in v.1.22.0. As part of this feature, all available execution providers are automatically registered with onnxruntime and the _best_ is selected based on a number of policies that the user can choose. Also, the `providers` and `provider_options` params on `InferenceSession` are now officially deprecated.

This is a stop-gap solution to unblock WinML (and thus //build). This will be revisited post conference to cleanup and generalize the Ort inference session management to avoid multiple difference code paths.
author_association: COLLABORATOR
user: login: xiaoyu-work
created_at: 2025-05-14T19:18:38Z
body: /azp run
author_association: COLLABORATOR
user: login: azure-pipelines[bot]
created_at: 2025-05-14T19:18:56Z
body: <samp>
Azure Pipelines successfully started running 2 pipeline(s).<br>

</samp>
author_association: NONE