Quacomm AI Engine Direct - Enable FCB feature to Qualcomm HTP

backends/qualcomm/aot/python/PyQnnManagerAdaptor.cpp

@@ -239,6 +239,10 @@ PYBIND11_MODULE(PyQnnManagerAdaptor, m) {
               const std::vector<std::string>&,
               std::vector<std::vector<std::shared_ptr<OpWrapper>>>&>(
               &PyQnnManager::Compile))
+      .def("CreateDlc", &PyQnnManager::CreateDlc)
+      .def("CompileToDlc", &PyQnnManager::CompileToDlc)
+      .def("GetDlcBinary", &PyQnnManager::GetDlcBinary)
+      .def("FreeDlc", &PyQnnManager::FreeDlc)
       .def("Destroy", &PyQnnManager::Destroy)
       .def("DestroyContext", &PyQnnManager::DestroyContext)
       .def("IsAvailable", &PyQnnManager::IsAvailable)

backends/qualcomm/aot/python/PyQnnManagerAdaptor.h

@@ -245,7 +245,7 @@ class PyQnnManager {
     return qnn_manager_->IsNodeSupportedByBackend(op_wrappers);
   }
 
-  py::array_t<char> Compile(
+  py::bytes Compile(
       const std::vector<std::string>& graph_names,
       std::vector<std::vector<std::shared_ptr<OpWrapper>>>& op_wrappers) {
     QnnExecuTorchContextBinary binary_info;
@@ -254,23 +254,55 @@ class PyQnnManager {
       if (qnn_manager_->Compile(graph_names[i], op_wrappers[i]) !=
           executorch::runtime::Error::Ok) {
         QNN_EXECUTORCH_LOG_ERROR("Fail to compile QNN graph");
-        return py::array_t<char>(0);
+        return py::bytes("", 0);
       }
     }
     auto qnn_executorch_options = GetQnnExecuTorchOptions(
         qnn_executorch_option_ptr_.cast<std::string_view>().data());
     if (qnn_executorch_options->saver() ||
         qnn_manager_->GetContextBinary(binary_info) !=
             executorch::runtime::Error::Ok) {
-      return py::array_t<char>(0);
+      return py::bytes("", 0);
     }
 
-    // allocate py::array (to pass the result of the C++ function to Python)
-    auto result = py::array_t<char>(binary_info.nbytes);
-    auto result_buffer = result.request();
-    char* result_ptr = (char*)result_buffer.ptr;
-    std::memcpy(result_ptr, binary_info.buffer, binary_info.nbytes);
-    return result;
+    return py::bytes(
+        reinterpret_cast<const char*>(binary_info.buffer), binary_info.nbytes);
+  }
+  py::int_ CreateDlc() {
+    void* handle = nullptr;
+    if (qnn_manager_->CreateDlc(handle) != Error::Ok) {
+      throw std::runtime_error("Failed to create QNN DLC");
+    }
+    return reinterpret_cast<uintptr_t>(handle);
+  }
+
+  void CompileToDlc(
+      const std::vector<std::string>& graph_names,
+      std::vector<std::vector<std::shared_ptr<OpWrapper>>>& op_wrappers,
+      uintptr_t dlc_handle) {
+    for (uint32_t i = 0; i < graph_names.size(); ++i) {
+      if (qnn_manager_->Compile(graph_names[i], op_wrappers[i]) != Error::Ok) {
+        throw std::runtime_error("Failed to compile QNN graph");
+      }
+    }
+    if (qnn_manager_->AddContextToDlc(reinterpret_cast<void*>(dlc_handle)) !=
+        Error::Ok) {
+      throw std::runtime_error("Failed to add QNN context to DLC");
+    }
+  }
+
+  py::bytes GetDlcBinary(uintptr_t dlc_handle) {
+    std::vector<uint8_t> binary;
+    if (qnn_manager_->GetDlcBinary(
+            reinterpret_cast<void*>(dlc_handle), binary) != Error::Ok) {
+      throw std::runtime_error("Failed to get QNN DLC binary");
+    }
+    return py::bytes(
+        reinterpret_cast<const char*>(binary.data()), binary.size());
+  }
+
+  void FreeDlc(uintptr_t dlc_handle) {
+    qnn_manager_->FreeDlc(reinterpret_cast<void*>(dlc_handle));
   }
 
   void Destroy() {

backends/qualcomm/aot/wrappers/OpWrapper.h

@@ -116,6 +116,22 @@ class OpWrapper final {
     }
     return raw_params;
   }
+
+  void ResetGraphState() {
+    for (const auto& tensor_wrapper : input_tensors_) {
+      tensor_wrapper->ResetGraphState();
+    }
+    for (const auto& tensor_wrapper : output_tensors_) {
+      tensor_wrapper->ResetGraphState();
+    }
+    for (const auto& param : params_) {
+      auto* tensor_param = dynamic_cast<TensorParamWrapper*>(param.get());
+      if (tensor_param != nullptr) {
+        tensor_param->GetTensorWrapper()->ResetGraphState();
+      }
+    }
+  }
+
   Qnn_OpConfig_t GetOpConfig();
 
  private:

backends/qualcomm/aot/wrappers/TensorWrapper.h

@@ -60,6 +60,11 @@ class TensorWrapper {
     created_ = true;
   }
 
+  void ResetGraphState() {
+    created_ = false;
+    QNN_TENSOR_VER_PTR(tensor_)->id = 0u;
+  }
+
   // Return true if the tensor is static:
   bool IsTensorStatic() const {
     return QNN_TENSOR_VER_PTR(tensor_)->type == QNN_TENSOR_TYPE_STATIC;

backends/qualcomm/export_utils.py

@@ -16,7 +16,7 @@
 import tempfile
 import types
 from dataclasses import dataclass, fields
-from typing import Callable, List, Optional, Set, Tuple, Union
+from typing import Callable, List, Optional, Sequence, Set, Tuple, Union
 
 import numpy as np
 import torch
@@ -676,10 +676,18 @@ def make_quantizer(
     is_qat=False,
     submodule_qconfig_list: Optional[List[Tuple[Callable, ModuleQConfig]]] = None,
     backend=QnnExecuTorchBackendType.kHtpBackend,
-    soc_model="SM8750",
+    soc_model: Union[str, QcomChipset, Sequence[Union[str, QcomChipset]]] = "SM8750",
     eps=None,
 ):
-    quantizer = QnnQuantizer(backend=backend, soc_model=getattr(QcomChipset, soc_model))
+    def to_chipset(model: Union[str, QcomChipset]) -> QcomChipset:
+        return model if isinstance(model, QcomChipset) else getattr(QcomChipset, model)
+
+    soc_models = (
+        [to_chipset(model) for model in soc_model]
+        if not isinstance(soc_model, (str, QcomChipset))
+        else to_chipset(soc_model)
+    )
+    quantizer = QnnQuantizer(backend=backend, soc_model=soc_models)
     quantizer.add_custom_quant_annotations(custom_annotations)
     quantizer.set_default_quant_config(
         quant_dtype,

backends/qualcomm/partition/qnn_partitioner.py

@@ -74,9 +74,15 @@ def __init__(
         # checker (e.g. the LPAI fallback pass) are distinguishable in the logs.
         self.phase = phase
         self.nodes_to_wrappers = defaultdict(dict)
-        self.qnn_manager = get_current_qnn_manager(
-            python_options.backend_options.backend_type, compiler_specs
+        target_socs = (
+            [target.soc_info.soc_model for target in python_options.fcb_options.targets]
+            if python_options.fcb_options is not None
+            else [python_options.soc_info.soc_model]
         )
+        self.qnn_managers = [
+            get_current_qnn_manager(compiler_specs, soc_model)
+            for soc_model in target_socs
+        ]
 
     def is_node_supported(self, _, node: torch.fx.Node) -> bool:
         if node.op != "call_function" or node.target in not_supported_operator:
@@ -128,8 +134,10 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
             op_wrapper_list.append(op_wrapper)
 
         if op_wrapper is not None:
-            supported = self.qnn_manager.IsNodeSupportedByBackend(
-                [op_wrapper.GetOpWrapper() for op_wrapper in op_wrapper_list]
+            wrappers = [op.GetOpWrapper() for op in op_wrapper_list]
+            supported = all(
+                manager.IsNodeSupportedByBackend(wrappers)
+                for manager in self.qnn_managers
             )
 
         self.nodes_to_wrappers.clear()