Moved _qnn_compile functionality to qnn_compiler.py

Signed-off-by: Shubham Agrawal <quic_shubhagr@quicinc.com>

Author: Shubham Agrawal

QEfficient/base/modeling_qeff.py

@@ -24,7 +24,6 @@
 from QEfficient.compile.qnn_compiler import compile as qnn_compile
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import constants, dump_qconfig
-from QEfficient.utils._utils import load_json
 from QEfficient.utils.cache import QEFF_HOME, to_hashable
 
 logger = logging.getLogger(__name__)
@@ -248,19 +247,6 @@ def _compile(
                 - convert_to_fp16=True -> -convert-to-fp16
 
         """
-        if enable_qnn:
-            return self._qnn_compile(
-                onnx_path,
-                compile_dir,
-                specializations=specializations,
-                custom_io=custom_io,
-                mdp_ts_num_devices=mdp_ts_num_devices,
-                num_cores=compiler_options.get("aic_num_cores", 16),
-                mxfp6_matmul=compiler_options.get("mxfp6_matmul", False),
-                mxint8_kv_cache=mxint8_kv_cache,
-                qnn_config=qnn_config,
-            )
-
         if onnx_path is None and self.onnx_path is None:
             self.export()
 
@@ -269,6 +255,22 @@ def _compile(
         qpc_path = compile_dir / "qpc"
         if not onnx_path.is_file():
             raise FileNotFoundError(f"ONNX file not found at: {onnx_path}")
+
+        if enable_qnn:
+            self.qpc_path = qnn_compile(
+                onnx_path=onnx_path,
+                qpc_base_path=compile_dir,
+                specializations=specializations,
+                custom_io=custom_io,
+                device_group=list(range(mdp_ts_num_devices)),
+                num_cores=compiler_options.get("aic_num_cores", 16),
+                mxfp6=compiler_options.get("mxfp6_matmul", False),
+                mxint8=mxint8_kv_cache,
+                qnn_config=qnn_config,
+            )
+
+            return self.qpc_path
+
         command = constants.COMPILER + [f"-m={onnx_path}"]
         if mdp_ts_json_path := compiler_options.pop("mdp_ts_json_path", None):
             mdp_ts_num_devices = None
@@ -363,96 +365,3 @@ def _compile(
         self.qpc_path = qpc_path
 
         return qpc_path
-
-    @dump_qconfig
-    def _qnn_compile(
-        self,
-        onnx_path: Optional[str] = None,
-        compile_dir: Optional[str] = None,
-        *,
-        custom_io: Optional[Dict[str, str]] = None,
-        specializations: Optional[List[Dict[str, int]]] = None,
-        mdp_ts_num_devices: int = 1,
-        num_cores: int = 16,
-        mxfp6_matmul: bool = False,
-        mxint8_kv_cache: bool = False,
-        qnn_config: Optional[str] = None,
-    ) -> str:
-        """
-        Interface for QNN compiler
-
-        Args:
-            :onnx_path (str): Onnx file to compile
-            :compile_dir (str): Directory path to compile the qpc. A suffix is added to the directory path to avoid reusing same qpc for different parameters.
-            :custom_io (dict): Custom IO to specify the input and outputs in different formats than default
-            :specializations (list): List of specializations to compile for
-            :mdp_ts_num_devices (int): Number of devices to partition to use Multi-Device Partitioning with tensor-slicing.
-            :num_cores (int): Number of cores used to compile the model.
-            :mxfp6_matmul (bool, optional): Whether to use ``mxfp6`` compression for weights. ``Defaults to True``.
-            :mxint8_kv_cache (bool, optional): Whether to use ``mxint8`` compression for KV cache. ``Defaults to False``.
-            :qnn_config (str): Path of QNN Config parameters file. ``Defaults to None.``
-        """
-        if onnx_path is None and self.onnx_path is None:
-            self.export()
-
-        onnx_path = Path(onnx_path or self.onnx_path)
-        compile_dir = Path(compile_dir or onnx_path.parent)
-        qpc_path = compile_dir / "qpc"
-        if not onnx_path.is_file():
-            raise FileNotFoundError(f"ONNX file not found at: {onnx_path}")
-
-        compile_hash = hashlib.sha256(to_hashable("qnn"))
-
-        if specializations is not None:
-            compile_hash.update(to_hashable(specializations))
-
-        if custom_io is not None:
-            compile_hash.update(to_hashable(custom_io))
-
-        if qnn_config is not None:
-            qnn_config_values = load_json(qnn_config)
-            compile_hash.update(to_hashable(qnn_config_values))
-
-        if mdp_ts_num_devices > 1:
-            compile_hash.update(to_hashable({"mdp_ts_num_devices": mdp_ts_num_devices}))
-
-        compile_hash.update(to_hashable({"num_cores": num_cores}))
-        compile_hash.update(to_hashable({"mxfp6_matmul": mxfp6_matmul}))
-        compile_hash.update(to_hashable({"mxint8_kv_cache": mxint8_kv_cache}))
-
-        # Check if already compiled
-        compile_hash = compile_hash.hexdigest()[:16]
-        qpc_path = qpc_path.with_name(qpc_path.name + "-" + compile_hash)
-        if qpc_path.is_dir():
-            if (qpc_path / "programqpc.bin").is_file():
-                self.qpc_path = qpc_path
-                return qpc_path
-            # Probably compilation failure last time, delete directory to start over
-            shutil.rmtree(qpc_path)
-
-        # Write specializations.json file
-        if specializations is not None:
-            specializations_json = compile_dir / "specializations.json"
-            with open(specializations_json, "w") as fp:
-                json.dump(
-                    {"specializations": [{k: str(v) for k, v in spec.items()} for spec in specializations]},
-                    fp,
-                    indent=4,
-                )
-
-        qnn_compile(
-            onnx_path=onnx_path,
-            qpc_base_path=compile_dir,
-            num_cores=num_cores,
-            device_group=list(range(mdp_ts_num_devices)),
-            mxfp6=mxfp6_matmul,
-            mxint8=mxint8_kv_cache,
-            qnn_config=qnn_config,
-            qnn_binary_dir=qpc_path,
-            specializations=specializations,
-            custom_io=custom_io,
-        )
-
-        self.qpc_path = qpc_path
-
-        return qpc_path

QEfficient/compile/compile_helper.py

@@ -196,6 +196,7 @@ def compile(
         qpc_path = qnn_compile(
             onnx_path=onnx_path,
             qpc_base_path=qpc_path,
+            qnn_binary_dir=os.path.join(qpc_path, "qpcs"),
             num_cores=num_cores,
             mxfp6=mxfp6,
             mxint8=mxint8,

QEfficient/compile/qnn_compiler.py

@@ -5,11 +5,14 @@
 #
 # -----------------------------------------------------------------------------
 
+import hashlib
+import json
 import os
 import shutil
 from typing import Dict, List, Optional
 
 from QEfficient.utils._utils import create_json, execute_command, load_json
+from QEfficient.utils.cache import to_hashable
 from QEfficient.utils.constants import QnnConstants
 from QEfficient.utils.generate_qnn_network_specialization_config import (
     generate_data_format_config,
@@ -384,6 +387,47 @@ def compile(
 
     prefill_only = True if len(specializations) == 1 else False
 
+    if qnn_binary_dir is None:
+        compile_hash = hashlib.sha256(to_hashable("qnn"))
+
+        if specializations is not None:
+            compile_hash.update(to_hashable(specializations))
+
+        if custom_io is not None:
+            compile_hash.update(to_hashable(custom_io))
+
+        if qnn_config is not None:
+            qnn_config_values = load_json(qnn_config)
+            compile_hash.update(to_hashable(qnn_config_values))
+
+        if device_group is not None:
+            compile_hash.update(to_hashable({"device_group": device_group}))
+
+        compile_hash.update(to_hashable({"num_cores": num_cores}))
+        compile_hash.update(to_hashable({"mxfp6": mxfp6}))
+        compile_hash.update(to_hashable({"mxint8": mxint8}))
+
+        # Check if already compiled
+        compile_hash = compile_hash.hexdigest()[:16]
+
+        qnn_binary_dir = qpc_base_path / "qpc"
+        qnn_binary_dir = qnn_binary_dir.with_name(qnn_binary_dir.name + "-" + compile_hash)
+        if qnn_binary_dir.is_dir():
+            if (qnn_binary_dir / "programqpc.bin").is_file():
+                return qnn_binary_dir
+            # Probably compilation failure last time, delete directory to start over
+            shutil.rmtree(qnn_binary_dir)
+
+        # Write specializations.json file
+        if specializations is not None:
+            specializations_json = qpc_base_path / "specializations.json"
+            with open(specializations_json, "w") as fp:
+                json.dump(
+                    {"specializations": [{k: str(v) for k, v in spec.items()} for spec in specializations]},
+                    fp,
+                    indent=4,
+                )
+
     qnn_obj = QNN(
         onnx_path=onnx_path,
         qpc_base_path=qpc_base_path,