olive changes to support QNN ABI EP

Author: Ronak

olive/common/ort_inference.py

@@ -53,6 +53,15 @@ def maybe_register_ep_libraries(ep_paths: dict[str, str]):
             if (Path(ort.__file__).parent / "capi" / builtin_library_name).exists():
                 ep_paths[provider] = builtin_library_name
 
+    # ABI Ep
+    if "QNNExecutionProvider" in ep_paths and ep_paths["QNNExecutionProvider"] is None:
+        try:
+            import onnxruntime_qnn as qnn_ep
+
+            ep_paths["QNNExecutionProvider"] = qnn_ep.get_library_path()
+        except ImportError:
+            logger.info("Failed to import onnxruntime_qnn")
+
     for ep_name, ep_path in ep_paths.items():
         if ep_path is None:
             continue

olive/passes/onnx/context_binary.py

@@ -76,16 +76,19 @@ def _run_for_config(
     ) -> Union[ONNXModelHandler, CompositeModelHandler]:
         # session created using providers argument so will use the ort.get_available_providers()
         # TODO(jambayk): consider switching to the new EP API for Windows
+        from onnxruntime import __version__ as OrtVersion
         from onnxruntime import get_available_providers
 
         # TODO(jambayk): validate and support other NPU EPs
         assert self.accelerator_spec.execution_provider == ExecutionProvider.QNNExecutionProvider, (
             "Only QNNExecutionProvider is supported for now."
         )
-        assert self.accelerator_spec.execution_provider in get_available_providers(), (
-            f"Execution provider {self.accelerator_spec.execution_provider} is not available. Available providers:"
-            f" {get_available_providers()}"
-        )
+
+        if version.parse(OrtVersion).release <= version.parse("1.23.2").release:
+            assert self.accelerator_spec.execution_provider in get_available_providers(), (
+                f"Execution provider {self.accelerator_spec.execution_provider} is not available. Available providers:"
+                f" {get_available_providers()}"
+            )
 
         result = self._run_single_target(model, config, output_model_path)
 
@@ -257,17 +260,22 @@ def _generate_context_binary(
         import onnxruntime as ort
         from onnxruntime import __version__ as OrtVersion
 
+        is_abi = (
+            "QNNExecutionProvider" not in ort.get_available_providers()
+            or version.parse(OrtVersion).release >= version.parse("1.25.0").release
+        )
+        logger.debug(" Using ABI EP: %s", str(is_abi))
+
         # prepare provider options
         provider_options = provider_options or {}
         if execution_provider == ExecutionProvider.QNNExecutionProvider:
             if str(device).lower() == "gpu":
                 provider_options["backend_path"] = "libQnnGpu.so" if platform.system() == "Linux" else "QnnGpu.dll"
                 update_llm_pipeline_genai_config_gpu_ctxbin(model_path)
             else:
-                if version.parse(OrtVersion).release < version.parse("1.22.0").release:
-                    provider_options["backend_path"] = "libQnnHtp.so" if platform.system() == "Linux" else "QnnHtp.dll"
-                    if share_ep_contexts:
-                        provider_options["enable_htp_weight_sharing"] = "1"
+                provider_options["backend_path"] = "libQnnHtp.so" if platform.system() == "Linux" else "QnnHtp.dll"
+                if share_ep_contexts:
+                    provider_options["enable_htp_weight_sharing"] = "1"
 
         # prepare session options
         session_options = session_options or {}
@@ -299,9 +307,40 @@ def _generate_context_binary(
         # create the inference session
         # requires regular onnxruntime package, not winml (not tested with winml)
         logger.debug("Creating context binary for model %s", str(model_path))
-        ort.InferenceSession(
-            model_path, sess_options=sess_options, providers=[execution_provider], provider_options=[provider_options]
-        )
+
+        if is_abi:
+            try:
+                import onnxruntime_qnn as qnn_ep
+
+                ep_lib_path = qnn_ep.get_library_path()
+                ep_registration_name = "QNNExecutionProvider"
+                ort.register_execution_provider_library(ep_registration_name, ep_lib_path)
+            except Exception as e:
+                if "already registered" in str(e):
+                    logger.debug(
+                        "Execution provider %s is already registered, skipping registration.", ep_registration_name
+                    )
+                else:
+                    raise
+            all_ep_devices = ort.get_ep_devices()
+            selected_ep_devices = [
+                ep_device for ep_device in all_ep_devices if ep_device.ep_name == ExecutionProvider.QNNExecutionProvider
+            ]
+
+            # Add QNN EP to session for abi ep
+            sess_options.add_provider_for_devices(selected_ep_devices, provider_options)
+            ort.InferenceSession(
+                model_path,
+                sess_options=sess_options,
+            )
+            ort.unregister_execution_provider_library(ep_registration_name)
+        else:
+            ort.InferenceSession(
+                model_path,
+                sess_options=sess_options,
+                providers=[execution_provider],
+                provider_options=[provider_options],
+            )
 
         assert output_model_path.exists(), f"Context binary not found at {output_model_path}"
 

olive/systems/utils/available_providers_runner.py

@@ -5,10 +5,13 @@
 # NOTE: Only onnxruntime and its dependencies can be imported in this file.
 import argparse
 import json
+import logging
 from pathlib import Path
 
 import onnxruntime as ort
 
+logger = logging.getLogger(__name__)
+
 
 def get_args(raw_args):
     parser = argparse.ArgumentParser(description="Get available execution providers")
@@ -19,10 +22,22 @@ def get_args(raw_args):
 
 def main(raw_args=None):
     args = get_args(raw_args)
-
+    available_eps = []
+    try:
+        import onnxruntime_qnn as qnn_ep
+
+        ep_lib_path = qnn_ep.get_library_path()
+        ep_registration_name = "QNNExecutionProvider"
+        ort.register_execution_provider_library(ep_registration_name, ep_lib_path)
+
+        # get available providers for ABI EP with ort 1.24 is broken. Hence the below hack
+        available_eps.append("QNNExecutionProvider")
+        ort.unregister_execution_provider_library(ep_registration_name)
+    except Exception as e:
+        logger.warning("Failed to register QNNExecutionProvider: %s", str(e))
     # get available execution providers
     # python environment system doesn't use EP registration yet
-    available_eps = ort.get_available_providers()
+    available_eps.extend(ort.get_available_providers())
 
     # save to json
     with Path(args.output_path).open("w") as f: