Qualcomm AI Engine Direct - enable multiple graphs in single pte

Summary:
- support multiple graphs in single qnn context in runtime
- helper function in aot for generating multi-methods pte
- enable weight sharing mechanism on HTP
- test cases

Author: haowhsu-quic

backends/qualcomm/aot/ir/qcir.fbs

@@ -94,8 +94,13 @@ table Operator {
 }
 
 table Graph {
+    name: string;
     nodes: [Operator];
     tensors: [Tensor];
 }
 
-root_type Graph;
+table Context {
+    graphs: [Graph];
+}
+
+root_type Context;

backends/qualcomm/aot/ir/qcir_utils.cpp

@@ -161,9 +161,7 @@ flatbuffers::Offset<qcir::QuantizeParam> ToQuantizeParam(
       }
     } break;
     default:
-      QNN_EXECUTORCH_LOG_WARN(
-          "QNN_QUANTIZATION_ENCODING_UNDEFINED detected: %s",
-          QNN_VER_PTR(tensor)->name);
+      // encodings are not required if lowering with floating point precision
       break;
   }
   return CreateQuantizeParamDirect(
@@ -229,9 +227,7 @@ Qnn_QuantizeParams_t ToQuantizeParam(const tensor_type& tensor) {
           const_cast<int32_t*>(param->offsets()->data());
     } break;
     default:
-      QNN_EXECUTORCH_LOG_WARN(
-          "qcir::QuantizeType::UNDEFINED detected: %s",
-          tensor->name()->c_str());
+      // encodings are not required if lowering with floating point precision
       break;
   }
   return p;

backends/qualcomm/aot/python/PyQnnManagerAdaptor.cpp

@@ -30,15 +30,22 @@ PYBIND11_MODULE(PyQnnManagerAdaptor, m) {
   py::class_<PyQnnManager, std::shared_ptr<PyQnnManager>>(m, "QnnManager")
       .def(py::init<const py::bytes&>())
       .def(py::init<const py::bytes&, const py::bytes&>())
+      .def(py::init<const py::bytes&, const py::list&>())
       .def("Init", &PyQnnManager::Init)
       .def("IsNodeSupportedByBackend", &PyQnnManager::IsNodeSupportedByBackend)
-      .def("Compile", &PyQnnManager::Compile)
+      .def("Compile", py::overload_cast<>(&PyQnnManager::Compile))
+      .def(
+          "Compile",
+          py::overload_cast<
+              const std::string&,
+              std::vector<std::shared_ptr<OpWrapper>>&>(&PyQnnManager::Compile))
       .def("Destroy", &PyQnnManager::Destroy)
       .def("IsAvailable", &PyQnnManager::IsAvailable)
       .def("IsTensorDump", &PyQnnManager::IsTensorDump)
       .def("AllocateTensor", &PyQnnManager::AllocateTensor)
       .def("GetGraphInputs", &PyQnnManager::GetGraphInputs)
       .def("GetGraphOutputs", &PyQnnManager::GetGraphOutputs)
+      .def("GetGraphNames", &PyQnnManager::GetGraphNames)
       .def("GetSpillFillBufferSize", &PyQnnManager::GetSpillFillBufferSize);
 }
 } // namespace qnn

backends/qualcomm/aot/python/PyQnnManagerAdaptor.h

@@ -47,20 +47,100 @@ class PyQnnManager {
     qnn_manager_ = std::make_shared<QnnManager>(
         qnn_executorch_options, qnn_executorch_context_binary_);
   }
+  // used for loading multiple graphs in qcir
+  explicit PyQnnManager(const py::bytes& buffer, const py::list& qcirs)
+      : qnn_executorch_option_ptr_(buffer) {
+    auto qnn_executorch_options = GetQnnExecuTorchOptions(
+        qnn_executorch_option_ptr_.cast<std::string_view>().data());
+
+    // merge multiple qcirs into one context with multiple graphs
+    std::vector<flatbuffers::Offset<qcir::Graph>> graphs;
+    for (size_t i = 0; i < qcirs.size(); ++i) {
+      py::buffer_info info(py::buffer(qcirs[i].cast<py::bytes>()).request());
+      flatbuffers::Verifier verifier(
+          static_cast<const uint8_t* const>(info.ptr),
+          info.size * info.itemsize);
+
+      if (!qcir::VerifyContextBuffer(verifier)) {
+        QNN_EXECUTORCH_LOG_ERROR("Fail to verify qcir format");
+        return;
+      }
+      auto context = qcir::GetContext(info.ptr);
+      for (const auto& graph : *context->graphs()) {
+        std::vector<flatbuffers::Offset<qcir::Tensor>> tensors;
+        for (const auto tensor : *graph->tensors()) {
+          // flatbuffers::Offset<Tensor> ToTensor(
+          //   QnnTensor
+          //   ToTensor(flatbuffers::Vector<::flatbuffers::Offset<qcir::Tensor>>
+          //   tensor), flatbuffers::FlatBufferBuilder* builder);
+          tensors.emplace_back(ToTensor(ToTensor(tensor), &builder_));
+        }
+        std::vector<flatbuffers::Offset<qcir::Operator>> nodes;
+        for (const auto& node : *graph->nodes()) {
+          int32_t* inputs_ptr = const_cast<int32_t*>(node->inputs()->data());
+          int32_t* outputs_ptr = const_cast<int32_t*>(node->outputs()->data());
+          int32_t* params_ptr = const_cast<int32_t*>(node->params()->data());
+          std::vector<int32_t> inputs(
+              inputs_ptr, inputs_ptr + node->inputs()->size());
+          std::vector<int32_t> outputs(
+              outputs_ptr, outputs_ptr + node->outputs()->size());
+          std::vector<int32_t> params(
+              params_ptr, params_ptr + node->params()->size());
+          nodes.emplace_back(qcir::CreateOperatorDirect(
+              builder_,
+              node->name()->str().c_str(),
+              node->package_name()->str().c_str(),
+              node->type_name()->str().c_str(),
+              &inputs,
+              &outputs,
+              &params));
+        }
+        graphs.emplace_back(qcir::CreateGraphDirect(
+            builder_, graph->name()->str().c_str(), &nodes, &tensors));
+      }
+    }
+    auto context = qcir::CreateContextDirect(builder_, &graphs);
+    builder_.Finish(context);
+    qnn_executorch_context_binary_.buffer = builder_.GetBufferPointer();
+    qnn_executorch_context_binary_.nbytes = builder_.GetSize();
+    qnn_manager_ = std::make_shared<QnnManager>(
+        qnn_executorch_options, qnn_executorch_context_binary_);
+  }
 
   executorch::runtime::Error Init() {
     return qnn_manager_->Init();
   }
+
   bool IsNodeSupportedByBackend(
       std::vector<std::shared_ptr<OpWrapper>>& op_wrappers) {
     return qnn_manager_->IsNodeSupportedByBackend(op_wrappers);
   }
+
+  // this method is specific for compiling multi-graphs
+  py::array_t<char> Compile() {
+    if (qnn_manager_->CompileQcir() != Error::Ok) {
+      QNN_EXECUTORCH_LOG_ERROR("Fail to compile qcir");
+      return py::array_t<char>(0);
+    }
+
+    // generate context binary if compilation succeded
+    QnnExecuTorchContextBinary context_binary;
+    qnn_manager_->GetContextBinary(context_binary);
+    // allocate py::array (to pass the result of the C++ function to Python)
+    auto result = py::array_t<char>(context_binary.nbytes);
+    auto result_buffer = result.request();
+    char* result_ptr = (char*)result_buffer.ptr;
+    std::memcpy(result_ptr, context_binary.buffer, context_binary.nbytes);
+    return result;
+  }
+
   py::array_t<char> Compile(
+      const std::string& graph_name,
       std::vector<std::shared_ptr<OpWrapper>>& op_wrappers) {
     QnnExecuTorchContextBinary context_binary;
     flatbuffers::FlatBufferBuilder builder;
 
-    if (qnn_manager_->IsOnlinePrepare()) {
+    if (qnn_manager_->IsOnlinePrepare() || qnn_manager_->IsMultipleGraphs()) {
       std::vector<flatbuffers::Offset<qcir::Tensor>> tensors;
       std::unordered_map<void*, int> tensor_map;
 
@@ -126,14 +206,19 @@ class PyQnnManager {
             &outputs,
             &params));
       }
-      auto graph = qcir::CreateGraphDirect(builder, &operators, &tensors);
-      builder.Finish(graph);
+      auto graph = qcir::CreateGraphDirect(
+          builder, graph_name.c_str(), &operators, &tensors);
+      std::vector<flatbuffers::Offset<qcir::Graph>> graphs({graph});
+      auto context = qcir::CreateContextDirect(builder, &graphs);
+      builder.Finish(context);
       context_binary.buffer = builder.GetBufferPointer();
       context_binary.nbytes = builder.GetSize();
-    } else if (
-        qnn_manager_->Compile(op_wrappers, context_binary) !=
-        executorch::runtime::Error::Ok) {
-      return py::array_t<char>(0);
+    } else {
+      if (qnn_manager_->Compile(graph_name, op_wrappers) !=
+          executorch::runtime::Error::Ok) {
+        return py::array_t<char>(0);
+      }
+      qnn_manager_->GetContextBinary(context_binary);
     }
 
     // allocate py::array (to pass the result of the C++ function to
@@ -144,6 +229,7 @@ class PyQnnManager {
     std::memcpy(result_ptr, context_binary.buffer, context_binary.nbytes);
     return result;
   }
+
   void Destroy() {
     return qnn_manager_->Destroy();
   }
@@ -156,28 +242,36 @@ class PyQnnManager {
     return qnn_manager_->IsTensorDump();
   }
 
-  executorch::runtime::Error AllocateTensor() {
-    return qnn_manager_->AllocateTensor();
+  executorch::runtime::Error AllocateTensor(const std::string& graph_name) {
+    return qnn_manager_->AllocateTensor(graph_name);
   }
 
-  py::list GetGraphInputs() {
+  py::list GetGraphInputs(const std::string& graph_name) {
     py::list ret;
     for (const std::shared_ptr<TensorWrapper>& input :
-         qnn_manager_->GetGraphInputs()) {
+         qnn_manager_->GetGraphInputs(graph_name)) {
       ret.append(PyQnnTensorWrapper(input));
     }
     return ret;
   }
 
-  py::list GetGraphOutputs() {
+  py::list GetGraphOutputs(const std::string& graph_name) {
     py::list ret;
     for (const std::shared_ptr<TensorWrapper>& output :
-         qnn_manager_->GetGraphOutputs()) {
+         qnn_manager_->GetGraphOutputs(graph_name)) {
       ret.append(PyQnnTensorWrapper(output));
     }
     return ret;
   }
 
+  py::list GetGraphNames() {
+    py::list ret;
+    for (const std::string& graph_name : qnn_manager_->GetGraphNames()) {
+      ret.append(graph_name);
+    }
+    return ret;
+  }
+
   uint64_t GetSpillFillBufferSize() {
     return qnn_manager_->GetSpillFillBufferSize();
   }
@@ -188,6 +282,7 @@ class PyQnnManager {
   const py::bytes qnn_executorch_option_ptr_;
   QnnExecuTorchContextBinary qnn_executorch_context_binary_;
   std::shared_ptr<QnnManager> qnn_manager_;
+  flatbuffers::FlatBufferBuilder builder_;
 };
 } // namespace qnn
 } // namespace backends

backends/qualcomm/partition/qnn_partitioner.py

@@ -10,9 +10,9 @@
 import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager
 import torch
 from executorch.backends.qualcomm.builders import node_visitor
+from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER
-from executorch.backends.qualcomm.utils.utils import generate_qnn_executorch_option
 
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import (
@@ -32,6 +32,7 @@
     not_supported_operator,
     to_be_implemented_operator,
 )
+from .utils import generate_qnn_executorch_option
 
 
 class QnnOperatorSupport(OperatorSupportBase):
@@ -63,7 +64,11 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
             )
             return False
 
-        if node.target in allow_list_operator:
+        if (
+            node.target in allow_list_operator
+            # bypass if custom op appears
+            or OpContextLoader.namespace == node.target.namespace
+        ):
             return True
 
         if (

backends/qualcomm/partition/utils.py

@@ -0,0 +1,22 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List
+
+from executorch.backends.qualcomm.utils.constants import QCOM_QNN_COMPILE_SPEC
+
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+
+def generate_qnn_executorch_option(
+    compiler_specs: List[CompileSpec],
+) -> bytes:
+    for compiler_spec in compiler_specs:
+        if compiler_spec.key == QCOM_QNN_COMPILE_SPEC:
+            qnn_compile_spec_buffer = compiler_spec.value
+        else:
+            raise ValueError(f"unknown compiler spec key value: {compiler_spec.key}")
+    return qnn_compile_spec_buffer

backends/qualcomm/qnn_preprocess.py

@@ -19,7 +19,7 @@
 from executorch.backends.qualcomm._passes.layout_transform import LayoutTransform
 from executorch.backends.qualcomm.builders.node_visitor import get_node_visitors
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
-from executorch.backends.qualcomm.utils.utils import generate_qnn_executorch_option
+from executorch.backends.qualcomm.partition.utils import generate_qnn_executorch_option
 from executorch.exir.backend.backend_details import (
     BackendDetails,
     CompileSpec,
@@ -83,7 +83,7 @@ def preprocess(
                     )
                     try:
                         context_loader_target = eval(
-                            f"torch.ops.{OpContextLoader.namespace}.{node.name}.default",
+                            f"torch.ops.{OpContextLoader.namespace}.{node.target.__name__}",
                             globals().update(torch.__dict__),
                         )
                         assert node.target == context_loader_target, err_msg
@@ -104,7 +104,8 @@ def preprocess(
             else:
                 raise RuntimeError(f"{node.op} is not supported in Qnn")
         qnn_context_binary = qnn_manager.Compile(
-            [py_op_wrapper.GetOpWrapper() for py_op_wrapper in py_op_wrapper_list]
+            qnn_manager.GetGraphNames()[0],
+            [py_op_wrapper.GetOpWrapper() for py_op_wrapper in py_op_wrapper_list],
         )
         assert len(qnn_context_binary) != 0, "Failed to generate Qnn context binary."
         qnn_manager.Destroy()