Freeing the commands and the command chain created during dispatch

runtime/hsa-runtime/core/inc/amd_aie_agent.h

@@ -93,6 +93,9 @@ class AieAgent : public core::Agent {
     return system_allocator_;
   }
 
+  /// @brief Getter for the AIE system deallocator.
+  const std::function<void(void*)>& system_deallocator() const { return system_deallocator_; }
+
   // AIE agent methods.
   /// @brief Get the number of columns on this AIE agent.
   int GetNumCols() const { return num_cols_; }
@@ -117,6 +120,9 @@ class AieAgent : public core::Agent {
                        core::MemoryRegion::AllocateFlags flags)>
       system_allocator_;
 
+
+  std::function<void(void*)> system_deallocator_;
+
   const hsa_profile_t profile_ = HSA_PROFILE_BASE;
   const uint32_t min_aql_size_ = 0x40;
   const uint32_t max_aql_size_ = 0x40;

runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp

@@ -345,6 +345,8 @@ void AieAgent::InitAllocators() {
                    ? mem
                    : nullptr;
       };
+
+      system_deallocator_ = [](void* ptr) { core::Runtime::runtime_singleton_->FreeMemory(ptr); };
       break;
     }
   }

runtime/hsa-runtime/core/runtime/amd_aie_aql_queue.cpp

@@ -129,7 +129,12 @@ AieAqlQueue::AieAqlQueue(AieAgent *agent, size_t req_size_pkts,
       .CreateQueue(*this);
 }
 
-AieAqlQueue::~AieAqlQueue() { AieAqlQueue::Inactivate(); }
+AieAqlQueue::~AieAqlQueue() {
+  AieAqlQueue::Inactivate();
+  if (ring_buf_) {
+    agent_.system_deallocator()(ring_buf_);
+  }
+}
 
 hsa_status_t AieAqlQueue::Inactivate() {
   bool active(active_.exchange(false, std::memory_order_relaxed));
@@ -346,6 +351,8 @@ hsa_status_t AieAqlQueue::SubmitCmd(
       case HSA_AMD_AIE_ERT_START_CU: {
         std::vector<uint32_t> bo_args;
         std::vector<uint32_t> cmd_handles;
+        std::vector<uint32_t> cmd_sizes;
+        std::vector<amdxdna_cmd *> cmds;
 
         // Iterating over future packets and seeing how many contiguous HSA_AMD_AIE_ERT_START_CU
         // packets there are. All can be combined into a single chain.
@@ -391,6 +398,8 @@ hsa_status_t AieAqlQueue::SubmitCmd(
 
           // Keeping track of the handle
           cmd_handles.push_back(cmd_bo_handle);
+          cmds.push_back(cmd);
+          cmd_sizes.push_back(cmd_size);
         }
 
         // Creating a packet that contains the command chain
@@ -439,6 +448,19 @@ hsa_status_t AieAqlQueue::SubmitCmd(
         // Executing all commands in the command chain
         ExecCmdAndWait(&exec_cmd_0, hw_ctx_handle, fd);
 
+        // Unmapping and closing the cmd BOs
+        drm_gem_close close_bo_args{0};
+        for (int i = 0; i < cmd_handles.size(); i++) {
+          munmap(cmds[i], cmd_sizes[i]);
+          close_bo_args.handle = cmd_handles[i];
+          ioctl(fd, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+        }
+
+        // Unmapping and closing the cmd_chain BO
+        munmap(cmd_chain, cmd_chain_size);
+        close_bo_args.handle = cmd_chain_bo_handle;
+        ioctl(fd, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+
         // Syncing BOs after we execute the command
         if (SyncBos(bo_args, fd))
           return HSA_STATUS_ERROR;