Qualcomm AI Engine Direct - Fix mem_handel register twice issue (#13410)

Summary:
- Insert registered handle in pre_registered_handles_ map to avoid
register multiple times for the same data_ptr

Background:
When running llama in lookahead mode using the same AR-N model for both
the prompt processor and token generator. The input and output are the
same, and the kv cache is shared between both components.. This causes a
"register twice" error message from QNN when a shared buffer (Smart
Mask) is used.

Error message:
```
[ERROR] [Qnn ExecuTorch]:  <E> Memory Handle duplicate found, matches Handle ID 0x2

[ERROR] [Qnn ExecuTorch]:  <E> Mem handle exists already

[ERROR] [Qnn ExecuTorch]:  <E> Failed to register memHandles

[ERROR] [Qnn ExecuTorch]:  <E> Failed to register memHandles

[ERROR] [Qnn ExecuTorch]:  <E> Failed to register mem with error 0x1f42
```

cc: @sxu , @haowhsu-quic

Author: shewu-quic

backends/qualcomm/runtime/QnnManager.cpp

@@ -291,7 +291,8 @@ Error QnnManager::RegisterCustomMem(
           data_ptr,
           unaligned_custom_mem_base,
           total_custom_mem_size,
-          tensor_offset) == Error::Ok,
+          tensor_offset,
+          info) == Error::Ok,
       Internal,
       "Fail to register to shared memory.");
 

backends/qualcomm/runtime/backends/QnnMemManager.cpp

@@ -56,13 +56,16 @@ Error QnnMemManager::RegisterIonMem(
   return Error::Ok;
 }
 
+// TODO: Find a better way to unify RegisterCustomMem and
+// PreRegisterCustomMemHandle
 Error QnnMemManager::RegisterCustomMem(
     const std::shared_ptr<TensorWrapper>& tensor_wrapper,
     int32_t mem_fd,
     void* mem_ptr,
     void* unaligned_custom_mem_base,
     size_t total_custom_mem_size,
-    size_t tensor_offset) {
+    size_t tensor_offset,
+    const CustomMemTensorInfo& info) {
   const QnnInterface& qnn_interface = implementation_.GetQnnInterface();
   Qnn_MemDescriptor_t descriptor = {
       {tensor_wrapper->GetRank(), tensor_wrapper->GetDims(), nullptr},
@@ -94,6 +97,7 @@ Error QnnMemManager::RegisterCustomMem(
     return Error::Internal;
   }
   tensor_wrapper->SetMemHandle(handle);
+  pre_registered_handles_.insert({info, handle});
   registered_map_.insert({handle, mem_ptr});
   if (log_level_ >= QnnExecuTorchLogLevel::kLogLevelInfo) {
     QNN_EXECUTORCH_LOG_INFO(

backends/qualcomm/runtime/backends/QnnMemManager.h

@@ -41,7 +41,8 @@ class QnnMemManager {
       void* mem_ptr,
       void* unaligned_custom_mem_base,
       size_t total_custom_mem_size,
-      size_t tensor_offset);
+      size_t tensor_offset,
+      const CustomMemTensorInfo& info);
 
   // Pre-register custom mem handle from SharedBuffer. Bring forward the
   // memHandle creating time from execution to initialization.
@@ -67,7 +68,9 @@ class QnnMemManager {
   const QnnImplementation& implementation_;
   QnnContext* context_;
   QnnExecuTorchLogLevel log_level_;
+  // Store the registered Qnn_MemHandle_t for de-registration
   std::unordered_map<Qnn_MemHandle_t, void*> registered_map_;
+  // Store the pre-registered custom mem handles
   std::unordered_map<CustomMemTensorInfo, void*> pre_registered_handles_;
   std::unordered_map<executorch::aten::ScalarType, Qnn_DataType_t>
       scalar_type_to_qnn_dtype_ = {