rocr/aie: Add support for creating AIE queue context

Adds support for initialzing the XDNA driver so that a hardware context can be created for an AIE queue. Right now this initializes the device heap in the driver, gets the relevant tile parameters for the AIE agent, and creates a hardware context that backs the AIE queue. Change-Id: Ib90e1bc67a8637f6db3ff2bebe34677843796417 [ROCm/ROCR-Runtime commit: 931733d51a]
2024-08-19 15:47:34 +00:00
@@ -61,6 +61,8 @@ namespace AMD {
 KfdDriver::KfdDriver(std::string devnode_name)
    : core::Driver(core::DriverType::KFD, devnode_name) {}

+hsa_status_t KfdDriver::Init() { return HSA_STATUS_SUCCESS; }
+
 hsa_status_t KfdDriver::DiscoverDriver() {
  if (hsaKmtOpenKFD() == HSAKMT_STATUS_SUCCESS) {
    std::unique_ptr<Driver> kfd_drv(new KfdDriver("/dev/kfd"));
@@ -74,6 +76,10 @@ hsa_status_t KfdDriver::QueryKernelModeDriver(core::DriverQuery query) {
  return HSA_STATUS_SUCCESS;
 }

+hsa_status_t KfdDriver::GetAgentProperties(core::Agent &agent) const {
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t
 KfdDriver::GetMemoryProperties(uint32_t node_id,
                               core::MemoryRegion &mem_region) const {
@@ -230,7 +236,7 @@ hsa_status_t KfdDriver::FreeMemory(void *mem, size_t size) {
  return FreeKfdMemory(mem, size) ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR;
 }

-hsa_status_t KfdDriver::CreateQueue(core::Queue &queue) {
+hsa_status_t KfdDriver::CreateQueue(core::Queue &queue) const {
  return HSA_STATUS_SUCCESS;
 }

@@ -43,12 +43,15 @@
 #include "core/inc/amd_xdna_driver.h"

 #include <sys/ioctl.h>
+#include <sys/mman.h>

 #include <memory>
 #include <string>

+#include "core/inc/amd_aie_aql_queue.h"
 #include "core/inc/amd_memory_region.h"
 #include "core/inc/runtime.h"
+#include "core/util/utils.h"
 #include "uapi/amdxdna_accel.h"

 namespace rocr {
@@ -57,6 +60,8 @@ namespace AMD {
 XdnaDriver::XdnaDriver(std::string devnode_name)
    : core::Driver(core::DriverType::XDNA, devnode_name) {}

+XdnaDriver::~XdnaDriver() { FreeDeviceHeap(); }
+
 hsa_status_t XdnaDriver::DiscoverDriver() {
  const int max_minor_num(64);
  const std::string devnode_prefix("/dev/accel/accel");
@@ -67,6 +72,7 @@ hsa_status_t XdnaDriver::DiscoverDriver() {
    if (xdna_drv->Open() == HSA_STATUS_SUCCESS) {
      if (xdna_drv->QueryKernelModeDriver(
              core::DriverQuery::GET_DRIVER_VERSION) == HSA_STATUS_SUCCESS) {
+        static_cast<XdnaDriver *>(xdna_drv.get())->Init();
        core::Runtime::runtime_singleton_->RegisterDriver(xdna_drv);
        return HSA_STATUS_SUCCESS;
      } else {
@@ -78,6 +84,8 @@ hsa_status_t XdnaDriver::DiscoverDriver() {
  return HSA_STATUS_ERROR;
 }

+hsa_status_t XdnaDriver::Init() { return InitDeviceHeap(); }
+
 hsa_status_t XdnaDriver::QueryKernelModeDriver(core::DriverQuery query) {
  switch (query) {
  case core::DriverQuery::GET_DRIVER_VERSION:
@@ -88,6 +96,29 @@ hsa_status_t XdnaDriver::QueryKernelModeDriver(core::DriverQuery query) {
  return HSA_STATUS_SUCCESS;
 }

+hsa_status_t XdnaDriver::GetAgentProperties(core::Agent &agent) const {
+  if (agent.device_type() != core::Agent::DeviceType::kAmdAieDevice) {
+    return HSA_STATUS_ERROR_INVALID_AGENT;
+  }
+
+  auto &aie_agent(static_cast<AieAgent &>(agent));
+
+  amdxdna_drm_query_aie_metadata aie_metadata{0};
+  amdxdna_drm_get_info get_info_args{
+      .param = DRM_AMDXDNA_QUERY_AIE_METADATA,
+      .buffer_size = sizeof(aie_metadata),
+      .buffer = reinterpret_cast<uintptr_t>(&aie_metadata)};
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_GET_INFO, &get_info_args) < 0) {
+    return HSA_STATUS_ERROR;
+  }
+
+  aie_agent.SetNumCols(aie_metadata.cols);
+  aie_agent.SetNumCoreRows(aie_metadata.core.row_count);
+
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t
 XdnaDriver::GetMemoryProperties(uint32_t node_id,
                                core::MemoryRegion &mem_region) const {
@@ -105,11 +136,51 @@ hsa_status_t XdnaDriver::FreeMemory(void *mem, size_t size) {
  return HSA_STATUS_SUCCESS;
 }

-hsa_status_t XdnaDriver::CreateQueue(core::Queue &queue) {
+hsa_status_t XdnaDriver::CreateQueue(core::Queue &queue) const {
+  if (!AieAqlQueue::IsType(&queue)) {
+    return HSA_STATUS_ERROR_INVALID_QUEUE;
+  }
+
+  auto &aie_queue(static_cast<AieAqlQueue &>(queue));
+  auto &aie_agent(aie_queue.GetAgent());
+
+  // Currently we do not leverage QoS information.
+  amdxdna_qos_info qos_info{0};
+  amdxdna_drm_create_hwctx create_hwctx_args{
+      .ext = 0,
+      .ext_flags = 0,
+      .qos_p = reinterpret_cast<uintptr_t>(&qos_info),
+      .umq_bo = 0,
+      .log_buf_bo = 0,
+      // TODO: Make this configurable.
+      .max_opc = 0x800,
+      // This field is for the number of core tiles.
+      .num_tiles = aie_agent.GetNumCores(),
+      .mem_size = 0,
+      .umq_doorbell = 0};
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CREATE_HWCTX, &create_hwctx_args) < 0) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  aie_queue.SetHwCtxHandle(create_hwctx_args.handle);
+
  return HSA_STATUS_SUCCESS;
 }

 hsa_status_t XdnaDriver::DestroyQueue(core::Queue &queue) const {
+  if (!AieAqlQueue::IsType(&queue)) {
+    return HSA_STATUS_ERROR_INVALID_QUEUE;
+  }
+
+  auto &aie_queue(static_cast<AieAqlQueue &>(queue));
+  amdxdna_drm_destroy_hwctx destroy_hwctx_args{.handle =
+                                                   aie_queue.GetHwCtxHandle()};
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_DESTROY_HWCTX, &destroy_hwctx_args) < 0) {
+    return HSA_STATUS_ERROR;
+  }
+
  return HSA_STATUS_SUCCESS;
 }

@@ -128,5 +199,75 @@ hsa_status_t XdnaDriver::QueryDriverVersion() {
  return HSA_STATUS_SUCCESS;
 }

+hsa_status_t XdnaDriver::InitDeviceHeap() {
+  amdxdna_drm_create_bo create_bo_args{
+    .flags = 0,
+    .type = AMDXDNA_BO_DEV_HEAP,
+    ._pad = 0,
+    .vaddr = reinterpret_cast<uintptr_t>(nullptr),
+    .size = dev_heap_size,
+    .handle = 0};
+
+  amdxdna_drm_get_bo_info get_bo_info_args{0};
+  drm_gem_close close_bo_args{0};
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CREATE_BO, &create_bo_args) < 0) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  get_bo_info_args.handle = create_bo_args.handle;
+  // In case we need to close this BO to avoid leaks due to some error after
+  // creation.
+  close_bo_args.handle = create_bo_args.handle;
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_GET_BO_INFO, &get_bo_info_args) < 0) {
+    // Close the BO in the case we can't get info about it.
+    ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+    return HSA_STATUS_ERROR;
+  }
+
+  dev_heap_parent = mmap(0, dev_heap_align * 2 - 1, PROT_READ | PROT_WRITE,
+                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+
+  if (dev_heap_parent == MAP_FAILED) {
+    // Close the BO in the case when a mapping fails and we got a BO handle.
+    ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+    dev_heap_parent = nullptr;
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  void *addr_aligned(reinterpret_cast<void *>(
+      AlignUp(reinterpret_cast<uintptr_t>(dev_heap_parent), dev_heap_align)));
+
+  dev_heap_aligned =
+      mmap(addr_aligned, dev_heap_size, PROT_READ | PROT_WRITE,
+           MAP_SHARED | MAP_FIXED, fd_, get_bo_info_args.map_offset);
+
+  if (dev_heap_aligned == MAP_FAILED) {
+    // Close the BO in the case when a mapping fails and we got a BO handle.
+    ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+    // Unmap the dev_heap_parent.
+    dev_heap_aligned = nullptr;
+    FreeDeviceHeap();
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::FreeDeviceHeap() {
+  if (dev_heap_parent) {
+    munmap(dev_heap_parent, dev_heap_align * 2 - 1);
+    dev_heap_parent = nullptr;
+  }
+
+  if (dev_heap_aligned) {
+    munmap(dev_heap_aligned, dev_heap_size);
+    dev_heap_aligned = nullptr;
+  }
+
+  return HSA_STATUS_SUCCESS;
+}
+
 } // namespace AMD
 } // namespace rocr
@@ -86,18 +86,37 @@ public:
    return regions_;
  }

+  // AIE agent methods.
+  /// @brief Get the number of columns on this AIE agent.
+  int GetNumCols() const { return num_cols_; }
+  void SetNumCols(int num_cols) { num_cols_ = num_cols; }
+  /// @brief Get the number of core tile rows on this AIE agent.
+  int GetNumCoreRows() const { return num_core_rows_; }
+  void SetNumCoreRows(int num_core_rows) { num_core_rows_ = num_core_rows; }
+  /// @brief Get the number of core tiles on this AIE agent.
+  int GetNumCores() const { return num_cols_ * num_core_rows_; }
+
 private:
-  // @brief Query the driver to get the region list owned by this agent.
+  /// @brief Query the driver to get the region list owned by this agent.
  void InitRegionList();

+  /// @brief Query the driver to get properties for this AIE agent.
+  void GetAgentProperties();
+
  std::vector<const core::MemoryRegion *> regions_;

  const hsa_profile_t profile_ = HSA_PROFILE_BASE;
-  static const uint32_t maxQueues_ = 8;
+  static const uint32_t maxQueues_ = 1;
  static const uint32_t minAqlSize_ = 0x40;
  static const uint32_t maxAqlSize_ = 0x40;
  uint32_t max_queues_;
  uintptr_t device_heap_vaddr_ = 0;
+
+  /// @brief Number of columns in the AIE array.
+  int num_cols_ = 0;
+  /// @brief Number of rows of core tiles in the AIE array. Not all rows in a
+  /// column are cores. Some can be memory or shim tiles.
+  int num_core_rows_ = 0;
 };

 } // namespace AMD
@@ -55,7 +55,9 @@ namespace AMD {
 /// @brief Encapsulates HW AIE AQL Command Processor functionality. It
 /// provides the interface for things such as doorbells, queue read and
 /// write pointers, and a buffer.
-class AieAqlQueue : public core::Queue, public core::DoorbellSignal {
+class AieAqlQueue : public core::Queue,
+                    private core::LocalSignal,
+                    core::DoorbellSignal {
 public:
  static __forceinline bool IsType(core::Signal *signal) {
    return signal->IsType(&rtti_id_);
@@ -95,6 +97,13 @@ public:
  hsa_status_t GetInfo(hsa_queue_info_attribute_t attribute,
                       void *value) override;

+  // AIE-specific API
+  AieAgent &GetAgent() { return agent_; }
+  void SetHwCtxHandle(uint32_t hw_ctx_handle) {
+    hw_ctx_handle_ = hw_ctx_handle;
+  }
+  uint32_t GetHwCtxHandle() const { return hw_ctx_handle_; }
+
  // GPU-specific queue functions are unsupported.
  hsa_status_t GetCUMasking(uint32_t num_cu_mask_count,
                            uint32_t *cu_mask) override;
@@ -126,7 +135,18 @@ private:
                                       uint32_t node_id);
  core::SharedSignal *CreateSharedSignal(AieAgent *agent);

-  AieAgent *agent_;
+  AieAgent &agent_;
+
+  /// @brief Handle for an application context on the AIE device.
+  ///
+  /// Each user queue will have an associated context. This handle is assigned
+  /// by the driver on context creation.
+  ///
+  /// TODO: For now we support a single context that allocates all core tiles in
+  /// the array. In the future we can make the number of tiles configurable so
+  /// that multiple workloads with different core tile configurations can
+  /// execute on the AIE agent at the same time.
+  uint32_t hw_ctx_handle_ = std::numeric_limits<uint32_t>::max();
  /// Indicates if queue is active.
  std::atomic<bool> active_;
  static int rtti_id_;
@@ -66,7 +66,9 @@ public:

  static hsa_status_t DiscoverDriver();

+  hsa_status_t Init() override;
  hsa_status_t QueryKernelModeDriver(core::DriverQuery query) override;
+  hsa_status_t GetAgentProperties(core::Agent &agent) const override;
  hsa_status_t
  GetMemoryProperties(uint32_t node_id,
                      core::MemoryRegion &mem_region) const override;
@@ -75,7 +77,7 @@ public:
                              void **mem, size_t size,
                              uint32_t node_id) override;
  hsa_status_t FreeMemory(void *mem, size_t size) override;
-  hsa_status_t CreateQueue(core::Queue &queue) override;
+  hsa_status_t CreateQueue(core::Queue &queue) const override;
  hsa_status_t DestroyQueue(core::Queue &queue) const override;

 private:
@@ -58,10 +58,14 @@ class XdnaDriver : public core::Driver {
 public:
  XdnaDriver() = delete;
  XdnaDriver(std::string devnode_name);
+  ~XdnaDriver();

  static hsa_status_t DiscoverDriver();
+
+  hsa_status_t Init() override;
  hsa_status_t QueryKernelModeDriver(core::DriverQuery query) override;

+  hsa_status_t GetAgentProperties(core::Agent &agent) const override;
  hsa_status_t
  GetMemoryProperties(uint32_t node_id,
                      core::MemoryRegion &mem_region) const override;
@@ -70,11 +74,36 @@ public:
                              void **mem, size_t size,
                              uint32_t node_id) override;
  hsa_status_t FreeMemory(void *mem, size_t size) override;
-  hsa_status_t CreateQueue(core::Queue &queue) override;
+
+  /// @brief Creates a context on the AIE device for this queue.
+  /// @param queue Queue whose on-device context is being created.
+  /// @return hsa_status_t
+  hsa_status_t CreateQueue(core::Queue &queue) const override;
  hsa_status_t DestroyQueue(core::Queue &queue) const override;

 private:
  hsa_status_t QueryDriverVersion();
+  /// @brief Allocate device accesible heap space.
+  ///
+  /// Allocate and map a buffer object (BO) that the AIE device can access.
+  hsa_status_t InitDeviceHeap();
+  hsa_status_t FreeDeviceHeap();
+
+  /// @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.
+  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_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
@@ -74,6 +74,9 @@ class Driver {
  Driver(DriverType kernel_driver_type, std::string devnode_name);
  virtual ~Driver() = default;

+  /// @brief Initialize the driver's state after opening.
+  virtual hsa_status_t Init() = 0;
+
  /// @brief Query the kernel-model driver.
  /// @retval HSA_STATUS_SUCCESS if the kernel-model driver query was
  /// successful.
@@ -91,6 +94,13 @@ class Driver {
  /// @retval DriverVersionInfo containing the driver's version information.
  const DriverVersionInfo &Version() const { return version_; }

+  /// @brief Get the properties of a specific agent and initialize the agent
+  /// object.
+  /// @param agent Agent whose properties we're getting.
+  /// @retval HSA_STATUS_SUCCESS if the driver successfully returns the agent's
+  ///         properties.
+  virtual hsa_status_t GetAgentProperties(Agent &agent) const = 0;
+
  /// @brief Get the memory properties of a specific node.
  /// @param node_id Node ID of the agent
  /// @param[in, out] mem_region MemoryRegion object whose properties will be
@@ -113,7 +123,7 @@ class Driver {

  virtual hsa_status_t FreeMemory(void *mem, size_t size) = 0;

-  virtual hsa_status_t CreateQueue(Queue &queue) = 0;
+  virtual hsa_status_t CreateQueue(Queue &queue) const = 0;

  virtual hsa_status_t DestroyQueue(Queue &queue) const = 0;

@@ -53,6 +53,7 @@ AieAgent::AieAgent(uint32_t node)
                  core::Agent::DeviceType::kAmdAieDevice),
      max_queues_(core::Runtime::runtime_singleton_->flag().max_queues()) {
  InitRegionList();
+  GetAgentProperties();
 }

 AieAgent::~AieAgent() {
@@ -187,5 +188,10 @@ hsa_status_t AieAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,

 void AieAgent::InitRegionList() {}

+void AieAgent::GetAgentProperties() {
+  core::Runtime::runtime_singleton_->AgentDriver(driver_type)
+      .GetAgentProperties(*this);
+}
+
 } // namespace AMD
 } // namespace rocr
@@ -70,8 +70,8 @@ int AieAqlQueue::rtti_id_ = 0;

 AieAqlQueue::AieAqlQueue(AieAgent *agent, size_t req_size_pkts,
                         uint32_t node_id)
-    : Queue(0, 0), DoorbellSignal(CreateSharedSignal(agent)), agent_(agent),
-      active_(false) {
+    : Queue(0, 0), LocalSignal(0, false), DoorbellSignal(signal()),
+      agent_(*agent), active_(false) {
  amd_queue_.hsa_queue.doorbell_signal = Signal::Convert(this);
  amd_queue_.hsa_queue.size = 0x40;

@@ -80,13 +80,24 @@ AieAqlQueue::AieAqlQueue(AieAgent *agent, size_t req_size_pkts,
  signal_.kind = AMD_SIGNAL_KIND_DOORBELL;
  signal_.queue_ptr = &amd_queue_;
  active_ = true;
+
+  core::Runtime::runtime_singleton_->AgentDriver(agent_.driver_type)
+      .CreateQueue(*this);
 }

 AieAqlQueue::~AieAqlQueue() { Inactivate(); }

 hsa_status_t AieAqlQueue::Inactivate() {
  bool active(active_.exchange(false, std::memory_order_relaxed));
-  return HSA_STATUS_SUCCESS;
+  hsa_status_t status(HSA_STATUS_SUCCESS);
+
+  if (active) {
+    status = core::Runtime::runtime_singleton_->AgentDriver(agent_.driver_type)
+                 .DestroyQueue(*this);
+    hw_ctx_handle_ = std::numeric_limits<uint32_t>::max();
+  }
+
+  return status;
 }

 hsa_status_t AieAqlQueue::SetPriority(HSA_QUEUE_PRIORITY priority) {
@@ -176,7 +187,7 @@ hsa_status_t AieAqlQueue::GetInfo(hsa_queue_info_attribute_t attribute,
                                  void *value) {
  switch (attribute) {
  case HSA_AMD_QUEUE_INFO_AGENT:
-    *(reinterpret_cast<hsa_agent_t *>(value)) = agent_->public_handle();
+    *(reinterpret_cast<hsa_agent_t *>(value)) = agent_.public_handle();
    break;
  case HSA_AMD_QUEUE_INFO_DOORBELL_ID:
    // Hardware doorbell supports AQL semantics.