rocr/aie: AIE Queue Processing

Change-Id: I681c971ba7229037ca85d5529838aa7bbe5820e2


[ROCm/ROCR-Runtime commit: e9cc839b2b]

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

@@ -100,6 +100,9 @@ public:
     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.
   uint32_t GetNumCols() const { return num_cols_; }
@@ -124,6 +127,9 @@ private:
                        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;

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_aie_aql_queue.h

@@ -49,7 +49,6 @@
 #include "core/inc/queue.h"
 #include "core/inc/runtime.h"
 #include "core/inc/signal.h"
-#include "core/util/locks.h"
 
 namespace rocr {
 namespace AMD {
@@ -131,6 +130,12 @@ private:
   /// @brief Base of the queue's ring buffer storage.
   void *ring_buf_ = nullptr;
 
+  /// @brief Called when the doorbell is rung to iterate over
+  /// all packets and submit them. Submissions is done by
+  // calling into the XdnaDriver.
+  hsa_status_t SubmitCmd(XdnaDriver& driver, void* queue_base, uint64_t read_dispatch_id,
+                         uint64_t write_dispatch_id);
+
   /// @brief Handle for an application context on the AIE device.
   ///
   /// Each user queue will have an associated context. This handle is assigned
@@ -154,4 +159,4 @@ private:
 } // namespace AMD
 } // namespace rocr
 
-#endif // header guard
+#endif  // HSA_RUNTIME_CORE_INC_AMD_HW_AQL_AIE_COMMAND_PROCESSOR_H_

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

@@ -45,9 +45,46 @@
 #include <memory>
 #include <unordered_map>
 
+#include "core/driver/xdna/uapi/amdxdna_accel.h"
+#include "core/inc/amd_aie_agent.h"
 #include "core/inc/driver.h"
 #include "core/inc/memory_region.h"
 
+/// @brief struct amdxdna_cmd_chain - Interpretation of data payload for
+/// ERT_CMD_CHAIN
+struct amdxdna_cmd_chain {
+  /// Number of commands in chain
+  __u32 command_count;
+  /// Index of last successfully submitted command in chain
+  __u32 submit_index;
+  /// Index of failing command if cmd status is not completed
+  __u32 error_index;
+  __u32 reserved[3];
+  /// Address of each command in chain
+  __u64 data[] __counted_by(command_count);
+};
+
+/// @brief struct amdxdna_cmd - Exec buffer command header format
+struct amdxdna_cmd {
+  union {
+    struct {
+      /// Current state of a command
+      __u32 state : 4;
+      __u32 unused : 6;
+      /// Extra CU masks in addition to mandatory mask
+      __u32 extra_cu_masks : 2;
+      /// Number of words in payload (data)
+      __u32 count : 11;
+      /// Opcode identifying specific command
+      __u32 opcode : 5;
+      __u32 reserved : 4;
+    };
+    __u32 header;
+  };
+  /// Count number of words representing packet payload
+  __u32 data[] __counted_by(count);
+};
+
 namespace rocr {
 namespace core {
 class Queue;
@@ -55,6 +92,40 @@ class Queue;
 
 namespace AMD {
 
+/// @brief: The number of arguments in the packet payload before we start passing operands
+constexpr uint32_t NON_OPERAND_COUNT = 6;
+
+// @brief: Used to transform an address into a device address
+constexpr uint32_t DEV_ADDR_BASE = 0x04000000;
+constexpr uint32_t DEV_ADDR_OFFSET_MASK = 0x02FFFFFF;
+
+/// @brief: The driver places a structure before each command in a command chain.
+/// Need to increase the size of the command by the size of this structure.
+/// In the following xdna driver source can see where this is implemented:
+/// Commit hash: eddd92c0f61592c576a500f16efa24eb23667c23
+/// https://github.com/amd/xdna-driver/blob/main/src/driver/amdxdna/aie2_msg_priv.h#L387-L391
+/// https://github.com/amd/xdna-driver/blob/main/src/driver/amdxdna/aie2_message.c#L637
+constexpr uint32_t CMD_COUNT_SIZE_INCREASE = 3;
+
+/// @brief: The size of an instruction in bytes
+constexpr uint32_t INSTR_SIZE_BYTES = 4;
+
+/// @brief: Index of command payload where the instruction sequence
+/// address is located
+constexpr uint32_t CMD_PKT_PAYLOAD_INSTRUCTION_SEQUENCE_IDX = 2;
+constexpr uint32_t CMD_PKT_PAYLOAD_INSTRUCTION_SEQUENCE_SIZE_IDX = 4;
+
+/// @brief Environment variable to define job submission timeout
+constexpr uint32_t DEFAULT_TIMEOUT_VAL = 50;
+
+/// @brief: Calculates the number of operands in a packet
+/// given the number of arguments in the packet
+/// @param: arg_count(Input), Number of arguments in the packet
+/// @return: uint32_t, The number of operands in the packet
+inline uint32_t GetOperandCount(uint32_t arg_count) {
+  return ((arg_count - NON_OPERAND_COUNT) / 2);
+}
+
 class XdnaDriver final : public core::Driver {
 public:
   XdnaDriver(std::string devnode_name);
@@ -68,6 +139,9 @@ public:
   hsa_status_t Init() override;
   hsa_status_t QueryKernelModeDriver(core::DriverQuery query) override;
 
+  std::unordered_map<uint32_t, void*>& GetHandleMappings();
+  std::unordered_map<void*, uint32_t>& GetAddrMappings();
+
   hsa_status_t GetAgentProperties(core::Agent &agent) const override;
   hsa_status_t
   GetMemoryProperties(uint32_t node_id,
@@ -84,7 +158,11 @@ public:
   hsa_status_t CreateQueue(core::Queue &queue) const override;
   hsa_status_t DestroyQueue(core::Queue &queue) const override;
 
-private:
+  // @brief Submits num_pkts packets in a command chain to the XDNA driver
+  hsa_status_t SubmitCmdChain(hsa_amd_aie_ert_packet_t* first_pkt, uint32_t num_pkts,
+                              uint32_t num_operands, uint32_t hw_ctx_handle);
+
+ private:
   hsa_status_t QueryDriverVersion();
   /// @brief Allocate device accesible heap space.
   ///
@@ -92,27 +170,55 @@ private:
   hsa_status_t InitDeviceHeap();
   hsa_status_t FreeDeviceHeap();
 
+  /// @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
+  hsa_status_t CreateCmd(uint32_t size, uint32_t* handle, amdxdna_cmd** cmd);
+
+  /// @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
+  hsa_status_t RegisterCmdBOs(uint32_t count, std::vector<uint32_t>& bo_args,
+                              std::vector<uint32_t>& bo_sizes, std::vector<uint64_t>& bo_addrs,
+                              hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload,
+                              const std::unordered_map<void*, uint32_t>& vmem_addr_mappings);
+
+  /// @brief Syncs all BOs referenced in bo_args
+  ///
+  /// @param bo_args vector containing handles of BOs to sync
+  hsa_status_t SyncBos(const std::vector<uint64_t>& bo_args, const std::vector<uint32_t>& bo_sizes);
+
+  /// @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
+  hsa_status_t ExecCmdAndWait(amdxdna_drm_exec_cmd* exec_cmd, uint32_t hw_ctx_handle);
+
   /// TODO: Remove this in the future and rely on the core Runtime
   /// object to track handle allocations. Using the VMEM API for mapping XDNA
   /// driver handles requires a bit more refactoring. So rely on the XDNA driver
   /// to manage some of this for now.
   std::unordered_map<uint32_t, void *> vmem_handle_mappings;
+  std::unordered_map<void*, uint32_t> vmem_addr_mappings;
 
   /// @brief Virtual address range allocated for the device heap.
   ///
   /// Allocate a large enough space so we can carve out the device heap in
-  /// this range and ensure it is aligned to 64MB. Currently, AIE2 supports
-  /// 48MB device heap and it must be aligned to 64MB.
+  /// this range and ensure it is aligned to 64MB. Currently, npu1 supports
+  /// 64MB device heap and it must be aligned to 64MB.
   void *dev_heap_parent = nullptr;
 
   /// @brief The aligned device heap.
   void *dev_heap_aligned = nullptr;
-  static constexpr size_t dev_heap_size = 48 * 1024 * 1024;
+  static constexpr size_t dev_heap_size = 64 * 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