Qualcomm AI Engine Direct - Enable zero copy feature

Summary:
- Add argument "shared_buffer" into compiler_spec, qnn_executor_runner and test scripts
  -  Actually, shared_buffer should be a runtime option since user are responsible to allocate memory for tensors on device. But it seems to have no way to set the runtime option to QnnBackend. Therefore, we put it to compile_spec for now.
- Implement SharedBuffer to allocate and free RPC memory
- Add QnnMemManger to register shared buffer for tensor
  - During exection time, we will register memory of tensor data for QNN. And we will deregister them during destruction time of QnnBackend
- Add two API `void* QnnExecuTorchAllocCustomMem(size_t bytes, size_t alignment)` and `void QnnExecuTorchFreeCustomMem(void* buffer_ptr)` to allocate RPC memory with SharedBuffer
  -  Users are responsible to allocate "enough" tensor bytes, and set alignment as MemoryAllocator::kDefaultAlignment. See runtime/core/memory_allocator.h.

Author: shewu-quic

backends/qualcomm/CMakeLists.txt

@@ -118,26 +118,28 @@ include_directories(
 #
 # declare targets
 #
+add_library(executorch_backend INTERFACE)
 add_library(qcir INTERFACE qcir_schema_output)
 add_library(qcir_utils STATIC)
-add_library(qnn_schema INTERFACE ${_qnn_schema__outputs})
-add_library(executorch_backend INTERFACE)
+add_library(qnn_backend STATIC)
+add_library(qnn_backend_cache STATIC)
+add_library(qnn_context STATIC)
+add_library(qnn_device STATIC)
 add_library(qnn_executorch_backend SHARED)
 add_library(qnn_executorch_header INTERFACE)
 add_library(qnn_executorch_logging STATIC)
-add_library(qnn_manager STATIC)
+add_library(qnn_factory STATIC)
 add_library(qnn_function_interface INTERFACE)
+add_library(qnn_graph STATIC)
+add_library(qnn_header INTERFACE)
 add_library(qnn_implementation STATIC)
+add_library(qnn_logger STATIC)
+add_library(qnn_manager STATIC)
+add_library(qnn_mem_manager STATIC)
+add_library(qnn_schema INTERFACE ${_qnn_schema__outputs})
 add_library(qnn_sys_function_interface INTERFACE)
 add_library(qnn_sys_implementation STATIC)
-add_library(qnn_logger STATIC)
-add_library(qnn_device STATIC)
-add_library(qnn_context STATIC)
-add_library(qnn_backend_cache STATIC)
-add_library(qnn_graph STATIC)
-add_library(qnn_backend STATIC)
-add_library(qnn_factory STATIC)
-add_library(qnn_header INTERFACE)
+add_library(shared_buffer STATIC)
 add_library(wrappers STATIC)
 add_library(utils STATIC)
 
@@ -214,6 +216,13 @@ target_link_libraries(qnn_graph
     qnn_implementation
     qnn_context
 )
+target_link_libraries(qnn_mem_manager
+    PRIVATE
+    qnn_executorch_logging
+    qnn_implementation
+    qnn_context
+)
+
 target_link_libraries(qnn_factory
     PUBLIC
     qnn_header
@@ -223,13 +232,15 @@ target_link_libraries(qnn_factory
     qnn_device
     qnn_context
     qnn_graph
+    qnn_mem_manager
 )
 target_link_libraries(qnn_manager
     PRIVATE
     qnn_factory
     wrappers
     qnn_schema
     utils
+    shared_buffer
 )
 target_link_libraries(qnn_executorch_backend
     PRIVATE
@@ -243,7 +254,11 @@ target_link_libraries(utils
     PRIVATE
     qnn_executorch_logging
 )
-
+target_link_libraries(shared_buffer
+    PRIVATE
+    qnn_executorch_logging
+    ${CMAKE_DL_LIBS}
+)
 #
 # add linker option
 #

backends/qualcomm/aot/wrappers/TensorWrapper.cpp

@@ -105,6 +105,7 @@ TensorWrapper::TensorWrapper(
 
 Error TensorWrapper::FillDataBuffer(const void* data, bool copy_data) {
   if (data != nullptr) {
+    QNN_VER_PTR(tensor_)->memType = QNN_TENSORMEMTYPE_RAW;
     QNN_VER_PTR(tensor_)->clientBuf.dataSize = bytes_;
     if (copy_data) {
       owned_data_ = std::make_unique<char[]>(bytes_);
@@ -144,6 +145,12 @@ Error TensorWrapper::SetName(const std::string& name) {
   return Error::Ok;
 }
 
+Error TensorWrapper::SetMemHandle(Qnn_MemHandle_t mem_handle) {
+  QNN_VER_PTR(tensor_)->memType = QNN_TENSORMEMTYPE_MEMHANDLE;
+  QNN_VER_PTR(tensor_)->memHandle = mem_handle;
+  return Error::Ok;
+}
+
 // base function for Create TensorWrapper
 std::shared_ptr<TensorWrapper> CreateTensorWrapper(
     const std::string& tensor_name,

backends/qualcomm/aot/wrappers/TensorWrapper.h

@@ -59,16 +59,38 @@ class TensorWrapper {
     return QNN_VER_PTR(tensor_)->type == QNN_TENSOR_TYPE_STATIC;
   };
 
-  const void* GetStaticTensorData() const {
-    return QNN_VER_PTR(tensor_)->clientBuf.data;
+  std::uint32_t* GetDims() const {
+    return QNN_VER_PTR(tensor_)->dimensions;
+  };
+
+  Qnn_DataType_t GetDataType() const {
+    return QNN_VER_PTR(tensor_)->dataType;
+  };
+
+  Qnn_MemHandle_t const GetMemHandle() {
+    return QNN_VER_PTR(tensor_)->memHandle;
+  };
+
+  Qnn_TensorMemType_t GetMemType() const {
+    return QNN_VER_PTR(tensor_)->memType;
   };
 
   std::string GetName() const {
     return qnn_tensor_name_;
   };
 
+  std::uint32_t GetRank() const {
+    return QNN_VER_PTR(tensor_)->rank;
+  };
+
+  const void* GetStaticTensorData() const {
+    return QNN_VER_PTR(tensor_)->clientBuf.data;
+  };
+
   Error SetName(const std::string& name);
 
+  Error SetMemHandle(Qnn_MemHandle_t mem_handle);
+
  private:
   // need this to handle QNN_TENSOR_ERROR_NAME_HASH_COLLISION
   std::string qnn_tensor_name_;

backends/qualcomm/runtime/CMakeLists.txt

@@ -47,3 +47,10 @@ target_sources(utils
     PRIVATE
     ${CMAKE_CURRENT_LIST_DIR}/Utils.cpp
 )
+
+# shared_buffer
+target_sources(shared_buffer
+    PRIVATE
+    ${CMAKE_CURRENT_LIST_DIR}/SharedBuffer.h
+    ${CMAKE_CURRENT_LIST_DIR}/SharedBuffer.cpp
+)

backends/qualcomm/runtime/QnnExecuTorch.h

@@ -31,6 +31,16 @@ typedef struct {
   }
 // clang-format on
 
+/// Allocate specific tensors (usually graph inputs and outputs) on shared
+/// memory. Users are responsible to allocate "enough" tensor bytes, and set
+/// alignment as MemoryAllocator::kDefaultAlignment.
+/// See runtime/core/memory_allocator.h. The function returns a valid pointer
+/// if allocation is successful.
+void* QnnExecuTorchAllocCustomMem(size_t bytes, size_t alignment);
+
+/// Free the allocated shared memory.
+void QnnExecuTorchFreeCustomMem(void* buffer_ptr);
+
 #ifdef __cplusplus
 }
 #endif // __cplusplus

backends/qualcomm/runtime/QnnExecuTorchBackend.cpp

@@ -188,19 +188,27 @@ Error QnnExecuTorchBackend::execute(
   std::vector<Qnn_Tensor_t> input_tensor_structs;
   std::vector<Qnn_Tensor_t> output_tensor_structs;
 
+  input_tensor_structs.reserve(input_tensors.size());
   for (int i = 0; i < input_tensors.size(); ++i) {
-    input_tensors[i]->FillDataBuffer(
-        args[i]->toTensor().const_data_ptr(), true /* copy_data */);
+    if (qnn_manager->RegisterMem(
+            args[i]->toTensor().mutable_data_ptr(), input_tensors[i]) !=
+        Error::Ok) {
+      input_tensors[i]->FillDataBuffer(
+          args[i]->toTensor().const_data_ptr(), true /* copy_data */);
+    }
     input_tensor_structs.push_back(input_tensors[i]->CloneTensorStruct());
   }
 
   int output_index = input_tensors.size();
   for (const auto& output_tensor : output_tensors) {
     // pos=0 limits the search to the prefix
     if (output_tensor->GetName().rfind("output_", 0) == 0) {
-      output_tensor->FillDataBuffer(
-          args[output_index]->toTensor().mutable_data_ptr(),
-          false /* copy_data */);
+      void* mutable_data_ptr =
+          args[output_index]->toTensor().mutable_data_ptr();
+      if (qnn_manager->RegisterMem(mutable_data_ptr, output_tensor) !=
+          Error::Ok) {
+        output_tensor->FillDataBuffer(mutable_data_ptr, false /* copy_data */);
+      }
       output_index++;
     }
     output_tensor_structs.push_back(output_tensor->CloneTensorStruct());

backends/qualcomm/runtime/QnnManager.cpp

@@ -6,9 +6,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 #include <executorch/backends/qualcomm/runtime/QnnManager.h>
+#include <executorch/backends/qualcomm/runtime/SharedBuffer.h>
 #include <executorch/backends/qualcomm/runtime/Utils.h>
 #include <executorch/backends/qualcomm/runtime/backends/QnnImplementation.h>
-
 #include <cstdlib>
 #include <cstring>
 #include <fstream>
@@ -51,7 +51,9 @@ QnnManager::QnnManager(
         "the size of qnn context binary: %d",
         qnn_executorch_context_binary.nbytes);
     QNN_EXECUTORCH_LOG_INFO(
-        "Is on-device graph construction: %d", options_->online_prepare());
+        "Is on-device graph construction: %d", options->online_prepare());
+    QNN_EXECUTORCH_LOG_INFO(
+        "Enable shared buffer: %d", options->shared_buffer());
   }
 
   if (library_path.empty()) {
@@ -79,6 +81,53 @@ Error QnnManager::LoadQnnLibrary() {
   return ret;
 }
 
+Error QnnManager::RegisterMem(
+    void* data_ptr,
+    const std::shared_ptr<TensorWrapper>& tensor_wrapper) {
+  SharedBuffer& shared_buffer_manager = SharedBuffer::GetSharedBufferManager();
+  // Not enable shared buffer
+  if (!options_->shared_buffer())
+    return Error::Internal;
+
+  if (backend_params_ptr_->qnn_mem_manager_ptr_ == nullptr) {
+    QNN_EXECUTORCH_LOG_WARN(
+        "Backend %s doesn't supported shared buffer.",
+        EnumNameQnnExecuTorchBackendType(
+            options_->backend_options()->backend_type()));
+    return Error::Internal;
+  }
+
+  if (!shared_buffer_manager.IsAllocated(data_ptr)) {
+    // It means two scenarios here:
+    // 1. the input and output partitioned graph
+    // 2. Actually, user doesn't allocate shared buffer with
+    // QnnExecuTorchAllocCustomMem API
+    return Error::Internal;
+  } else if (backend_params_ptr_->qnn_mem_manager_ptr_->IsRegistered(
+                 tensor_wrapper->GetMemHandle())) {
+    if (options_->log_level() >= QnnExecuTorchLogLevel::kLogLevelInfo)
+      QNN_EXECUTORCH_LOG_INFO(
+          "Tensor name %s has been registered shared memory.",
+          tensor_wrapper->GetName().c_str());
+    return Error::Ok;
+  }
+
+  int32_t mem_fd = SharedBuffer::GetSharedBufferManager().MemToFd(data_ptr);
+  if (mem_fd == -1) {
+    QNN_EXECUTORCH_LOG_WARN(
+        "Tensor name %s is failed to get file descriptor.",
+        tensor_wrapper->GetName().c_str());
+    return Error::Internal;
+  }
+  ET_CHECK_OR_RETURN_ERROR(
+      backend_params_ptr_->qnn_mem_manager_ptr_->RegisterMem(
+          tensor_wrapper, mem_fd) == Error::Ok,
+      Internal,
+      "Fail to register to shared memory.");
+
+  return Error::Ok;
+}
+
 Error QnnManager::Init() {
   ET_CHECK_OR_RETURN_ERROR(
       LoadQnnLibrary() == Error::Ok, Internal, "Fail to load Qnn library");
@@ -202,14 +251,6 @@ void QnnManager::Destroy() {
   qnn_loaded_backend_.TerminateAllBackends();
 }
 
-bool QnnManager::IsAvailable() {
-  return true;
-}
-
-bool QnnManager::IsOnlinePrepare() {
-  return options_->online_prepare();
-}
-
 bool QnnManager::IsNodeSupportedByBackend(
     std::vector<std::shared_ptr<OpWrapper>>& op_wrappers) {
   Qnn_ErrorHandle_t error = QNN_SUCCESS;
@@ -312,3 +353,14 @@ Error QnnManager::Compile(
 } // namespace qnn
 } // namespace executor
 } // namespace torch
+void* QnnExecuTorchAllocCustomMem(size_t bytes, size_t alignment) {
+  using torch::executor::qnn::SharedBuffer;
+  void* buffer_ptr =
+      SharedBuffer::GetSharedBufferManager().AllocMem(bytes, alignment);
+  return buffer_ptr;
+}
+
+void QnnExecuTorchFreeCustomMem(void* buffer_ptr) {
+  using torch::executor::qnn::SharedBuffer;
+  SharedBuffer::GetSharedBufferManager().FreeMem(buffer_ptr);
+}

backends/qualcomm/runtime/QnnManager.h

@@ -40,21 +40,29 @@ class QnnManager {
 
   void Destroy();
 
-  bool IsAvailable();
+  bool IsAvailable() {
+    return true;
+  }
+
+  bool IsOnlinePrepare() {
+    return options_->online_prepare();
+  }
 
   bool IsTensorDump() {
     return options_->tensor_dump_output_path()->size() > 0;
   }
 
-  bool IsOnlinePrepare();
-
   bool IsNodeSupportedByBackend(
       std::vector<std::shared_ptr<OpWrapper>>& op_wrappers);
 
   Error Compile(
       std::vector<std::shared_ptr<OpWrapper>>& op_wrappers,
       QnnExecuTorchContextBinary& qnn_executorch_context_binary);
 
+  Error RegisterMem(
+      void* data_ptr,
+      const std::shared_ptr<TensorWrapper>& tensor_wrapper);
+
   std::vector<std::shared_ptr<TensorWrapper>> GetGraphInputs() {
     return input_tensors_;
   }