rocr/aie: AIE Queue Processing

runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp

@@ -117,7 +117,8 @@ hsa_status_t XdnaDriver::GetAgentProperties(core::Agent &agent) const {
     return HSA_STATUS_ERROR;
   }
 
-  aie_agent.SetNumCols(aie_metadata.cols);
+  // Right now can only target N-1 columns
+  aie_agent.SetNumCols(aie_metadata.cols - 1);
   aie_agent.SetNumCoreRows(aie_metadata.core.row_count);
 
   return HSA_STATUS_SUCCESS;
@@ -147,7 +148,7 @@ XdnaDriver::AllocateMemory(const core::MemoryRegion &mem_region,
   }
 
   if (m_region.kernarg()) {
-    create_bo_args.type = AMDXDNA_BO_CMD;
+    create_bo_args.type = AMDXDNA_BO_SHMEM;
   } else {
     create_bo_args.type = AMDXDNA_BO_DEV;
   }
@@ -216,7 +217,7 @@ hsa_status_t XdnaDriver::CreateQueue(core::Queue &queue) const {
       // TODO: Make this configurable.
       .max_opc = 0x800,
       // This field is for the number of core tiles.
-      .num_tiles = aie_agent.GetNumCores(),
+      .num_tiles = static_cast<uint32_t>(aie_agent.GetNumCores()),
       .mem_size = 0,
       .umq_doorbell = 0};
 
@@ -316,6 +317,16 @@ hsa_status_t XdnaDriver::InitDeviceHeap() {
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t XdnaDriver::GetHandleMappings(std::unordered_map<uint32_t, void*> &vmem_handle_mappings) {
+  vmem_handle_mappings = this->vmem_handle_mappings;
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::GetFd(int &fd) {
+  fd = fd_;
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t XdnaDriver::FreeDeviceHeap() {
   if (dev_heap_parent) {
     munmap(dev_heap_parent, dev_heap_align * 2 - 1);

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/inc/amd_aie_aql_queue.h

@@ -49,7 +49,35 @@
 #include "core/inc/queue.h"
 #include "core/inc/runtime.h"
 #include "core/inc/signal.h"
-#include "core/util/locks.h"
+
+/*
+ * Interpretation of the beginning of data payload for ERT_CMD_CHAIN in
+ * amdxdna_cmd. The rest of the payload in amdxdna_cmd is cmd BO handles.
+ */
+struct amdxdna_cmd_chain {
+  __u32 command_count;
+  __u32 submit_index;
+  __u32 error_index;
+  __u32 reserved[3];
+  __u64 data[] __counted_by(command_count);
+};
+
+
+/* Exec buffer command header format */
+struct amdxdna_cmd {
+  union {
+    struct {
+      __u32 state : 4;
+      __u32 unused : 6;
+      __u32 extra_cu_masks : 2;
+      __u32 count : 11;
+      __u32 opcode : 5;
+      __u32 reserved : 4;
+    };
+    __u32 header;
+  };
+  __u32 data[]  __counted_by(count);
+};
 
 namespace rocr {
 namespace AMD {
@@ -71,7 +99,7 @@ class AieAqlQueue : public core::Queue,
 
   AieAqlQueue() = delete;
   AieAqlQueue(AieAgent *agent, size_t req_size_pkts, uint32_t node_id);
-  ~AieAqlQueue();
+  ~AieAqlQueue() override;
 
   hsa_status_t Inactivate() override;
   hsa_status_t SetPriority(HSA_QUEUE_PRIORITY priority) override;
@@ -100,7 +128,7 @@ class AieAqlQueue : public core::Queue,
                        void *value) override;
 
   // AIE-specific API
-  AieAgent &GetAgent() { return agent_; }
+  AieAgent &GetAgent() const { return agent_; }
   void SetHwCtxHandle(uint32_t hw_ctx_handle) {
     hw_ctx_handle_ = hw_ctx_handle;
   }
@@ -116,7 +144,7 @@ class AieAqlQueue : public core::Queue,
                   hsa_fence_scope_t releaseFence = HSA_FENCE_SCOPE_NONE,
                   hsa_signal_t *signal = NULL) override;
 
-  uint32_t queue_id_ = INVALID_QUEUEID;
+  uint64_t queue_id_ = INVALID_QUEUEID;
   /// @brief ID of AIE device on which this queue has been mapped.
   uint32_t node_id_ = std::numeric_limits<uint32_t>::max();
   /// @brief Queue size in bytes.
@@ -131,6 +159,44 @@ class AieAqlQueue : public core::Queue,
   /// @brief Base of the queue's ring buffer storage.
   void *ring_buf_ = nullptr;
 
+  static hsa_status_t SubmitCmd(
+      uint32_t hw_ctx_handle, int fd, void *queue_base,
+      uint64_t read_dispatch_id, uint64_t write_dispatch_id,
+      std::unordered_map<uint32_t, void *> &vmem_handle_mappings);
+
+  /// @brief Creates a command BO and returns a pointer to the memory and
+  //          the corresponding handle
+  ///
+  /// @param size size of memory to allocate
+  /// @param handle A pointer to the BO handle
+  /// @param cmd A pointer to the buffer
+  static hsa_status_t CreateCmd(uint32_t size, uint32_t *handle,
+                                amdxdna_cmd **cmd, int fd);
+
+  /// @brief Adds all BOs in a command packet payload to a vector
+  ///         and replaces the handles with a virtual address
+  ///
+  /// @param count Number of entries in the command
+  /// @param bo_args A pointer to a vector that contains all bo handles
+  /// @param cmd_pkt_payload A pointer to the payload of the command
+  static void RegisterCmdBOs(
+      uint32_t count, std::vector<uint32_t> &bo_args,
+      hsa_amd_aie_ert_start_kernel_data_t *cmd_pkt_payload,
+      std::unordered_map<uint32_t, void *> &vmem_handle_mappings);
+
+  /// @brief Syncs all BOs referenced in bo_args
+  ///
+  /// @param bo_args vector containing handles of BOs to sync
+  static hsa_status_t SyncBos(std::vector<uint32_t> &bo_args, int fd);
+
+  /// @brief Executes a command and waits for its completion
+  ///
+  /// @param exec_cmd Structure containing the details of the command to execute
+  /// @param hw_ctx_handle the handle of the hardware context to run this
+  /// command
+  static hsa_status_t ExecCmdAndWait(amdxdna_drm_exec_cmd *exec_cmd,
+                                     uint32_t hw_ctx_handle, int fd);
+
   /// @brief Handle for an application context on the AIE device.
   ///
   /// Each user queue will have an associated context. This handle is assigned
@@ -154,4 +220,4 @@ class AieAqlQueue : public core::Queue,
 } // namespace AMD
 } // namespace rocr
 
-#endif // header guard
+#endif // HSA_RUNTIME_CORE_INC_AMD_HW_AQL_AIE_COMMAND_PROCESSOR_H_

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

@@ -47,6 +47,7 @@
 
 #include "core/inc/driver.h"
 #include "core/inc/memory_region.h"
+#include "core/driver/xdna/uapi/amdxdna_accel.h"
 
 namespace rocr {
 namespace core {
@@ -69,6 +70,9 @@ class XdnaDriver : public core::Driver {
   hsa_status_t Init() override;
   hsa_status_t QueryKernelModeDriver(core::DriverQuery query) override;
 
+  hsa_status_t GetHandleMappings(std::unordered_map<uint32_t, void*> &vmem_handle_mappings);
+  hsa_status_t GetFd(int &fd);
+
   hsa_status_t GetAgentProperties(core::Agent &agent) const override;
   hsa_status_t
   GetMemoryProperties(uint32_t node_id,
@@ -110,10 +114,6 @@ class XdnaDriver : public core::Driver {
   void *dev_heap_aligned = nullptr;
   static constexpr size_t dev_heap_size = 48 * 1024 * 1024;
   static constexpr size_t dev_heap_align = 64 * 1024 * 1024;
-
-  /// @brief DRM buffer object handle for the device heap. Assigned by the
-  ///        kernel-mode driver.
-  uint32_t dev_heap_handle = 0;
 };
 
 } // namespace AMD

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

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