rocr/aie: Init mem regions for AIE agents

Change-Id: If180bdbcb3eb659f0d05a710526864494316d7a9


[ROCm/ROCR-Runtime commit: a851f73da5]

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

@@ -129,6 +129,51 @@ hsa_status_t
 XdnaDriver::AllocateMemory(const core::MemoryRegion &mem_region,
                            core::MemoryRegion::AllocateFlags alloc_flags,
                            void **mem, size_t size, uint32_t node_id) {
+  const MemoryRegion &m_region(static_cast<const MemoryRegion &>(mem_region));
+
+  amdxdna_drm_create_bo create_bo_args{0};
+  create_bo_args.size = size;
+
+  amdxdna_drm_get_bo_info get_bo_info_args{0};
+  drm_gem_close close_bo_args{0};
+
+  if (!m_region.IsSystem()) {
+    return HSA_STATUS_ERROR_INVALID_REGION;
+  }
+
+  if (m_region.kernarg()) {
+    create_bo_args.type = AMDXDNA_BO_CMD;
+  } else {
+    create_bo_args.type = AMDXDNA_BO_DEV;
+  }
+
+  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;
+  }
+
+  if (m_region.kernarg()) {
+    *mem = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd_,
+                get_bo_info_args.map_offset);
+    if (*mem == 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);
+      return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    }
+  } else {
+    *mem = reinterpret_cast<void *>(get_bo_info_args.vaddr);
+  }
+
   return HSA_STATUS_SUCCESS;
 }
 

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

@@ -86,6 +86,13 @@ public:
     return regions_;
   }
 
+  /// @brief Getter for the AIE system allocator.
+  const std::function<void *(size_t size, size_t align,
+                             core::MemoryRegion::AllocateFlags flags)> &
+  system_allocator() const {
+    return system_allocator_;
+  }
+
   // AIE agent methods.
   /// @brief Get the number of columns on this AIE agent.
   int GetNumCols() const { return num_cols_; }
@@ -99,18 +106,21 @@ public:
 private:
   /// @brief Query the driver to get the region list owned by this agent.
   void InitRegionList();
+  /// @brief Setup the memory allocators used by this agent.
+  void InitAllocators();
 
   /// @brief Query the driver to get properties for this AIE agent.
   void GetAgentProperties();
 
   std::vector<const core::MemoryRegion *> regions_;
+  std::function<void *(size_t size, size_t align,
+                       core::MemoryRegion::AllocateFlags flags)>
+      system_allocator_;
 
   const hsa_profile_t profile_ = HSA_PROFILE_BASE;
-  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;
+  const uint32_t min_aql_size_ = 0x40;
+  const uint32_t max_aql_size_ = 0x40;
+  const uint32_t max_queues_ = 1;
 
   /// @brief Number of columns in the AIE array.
   int num_cols_ = 0;

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

@@ -143,7 +143,8 @@ class MemoryRegion : public core::MemoryRegion {
   }
 
   __forceinline bool IsSystem() const {
-    return mem_props_.HeapType == HSA_HEAPTYPE_SYSTEM;
+    return ((mem_props_.HeapType == HSA_HEAPTYPE_SYSTEM) ||
+            (mem_props_.HeapType == HSA_HEAPTYPE_DEVICE_SVM));
   }
 
   __forceinline bool IsLDS() const {

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

@@ -42,17 +42,21 @@
 
 #include "core/inc/amd_aie_agent.h"
 
+#include <functional>
+
 #include "core/inc/amd_aie_aql_queue.h"
+#include "core/inc/amd_memory_region.h"
 #include "core/inc/driver.h"
+#include "core/inc/runtime.h"
 
 namespace rocr {
 namespace AMD {
 
 AieAgent::AieAgent(uint32_t node)
     : core::Agent(core::DriverType::XDNA, node,
-                  core::Agent::DeviceType::kAmdAieDevice),
-      max_queues_(core::Runtime::runtime_singleton_->flag().max_queues()) {
+                  core::Agent::DeviceType::kAmdAieDevice) {
   InitRegionList();
+  InitAllocators();
   GetAgentProperties();
 }
 
@@ -82,7 +86,8 @@ hsa_status_t AieAgent::IterateRegion(
 hsa_status_t AieAgent::IterateCache(hsa_status_t (*callback)(hsa_cache_t cache,
                                                              void *data),
                                     void *data) const {
-  return HSA_STATUS_SUCCESS;
+  // AIE has no caches.
+  return HSA_STATUS_ERROR_INVALID_CACHE;
 }
 
 hsa_status_t AieAgent::GetInfo(hsa_agent_info_t attribute, void *value) const {
@@ -117,13 +122,13 @@ hsa_status_t AieAgent::GetInfo(hsa_agent_info_t attribute, void *value) const {
     *reinterpret_cast<uint32_t *>(value) = 0;
     break;
   case HSA_AGENT_INFO_QUEUES_MAX:
-    *reinterpret_cast<uint32_t *>(value) = maxQueues_;
+    *reinterpret_cast<uint32_t *>(value) = max_queues_;
     break;
   case HSA_AGENT_INFO_QUEUE_MIN_SIZE:
-    *reinterpret_cast<uint32_t *>(value) = minAqlSize_;
+    *reinterpret_cast<uint32_t *>(value) = min_aql_size_;
     break;
   case HSA_AGENT_INFO_QUEUE_MAX_SIZE:
-    *reinterpret_cast<uint32_t *>(value) = maxAqlSize_;
+    *reinterpret_cast<uint32_t *>(value) = max_aql_size_;
     break;
   case HSA_AGENT_INFO_QUEUE_TYPE:
     *reinterpret_cast<hsa_queue_type32_t *>(value) = HSA_QUEUE_TYPE_SINGLE;
@@ -176,7 +181,7 @@ hsa_status_t AieAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
 
-  if (size < minAqlSize_ || size > maxAqlSize_) {
+  if (size < min_aql_size_ || size > max_aql_size_) {
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
 
@@ -186,12 +191,53 @@ hsa_status_t AieAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,
   return HSA_STATUS_SUCCESS;
 }
 
-void AieAgent::InitRegionList() {}
+void AieAgent::InitRegionList() {
+  /// TODO: Find a way to set the other memory properties in a reasonable way.
+  ///       This should be easier once the ROCt source is incorporated into the
+  ///       ROCr source. Since the AIE itself currently has no memory regions of
+  ///       its own all memory is just the system DRAM.
+
+  /// For allocating kernel arguments or other objects that only need
+  /// system memory.
+  HsaMemoryProperties sys_mem_props{
+      .HeapType = HSA_HEAPTYPE_SYSTEM,
+  };
+  /// For allocating memory for programmable device image (PDI) files. These
+  /// need to be mapped to the device so the hardware can access the PDIs.
+  HsaMemoryProperties dev_mem_props{
+      .HeapType = HSA_HEAPTYPE_DEVICE_SVM,
+  };
+  /// As of now the AIE devices support coarse-grain memory regions that require
+  /// explicit sync operations.
+  regions_.push_back(
+      new MemoryRegion(false, true, false, false, true, this, sys_mem_props));
+  regions_.push_back(
+      new MemoryRegion(false, false, false, false, true, this, dev_mem_props));
+}
 
 void AieAgent::GetAgentProperties() {
   core::Runtime::runtime_singleton_->AgentDriver(driver_type)
       .GetAgentProperties(*this);
 }
 
+void AieAgent::InitAllocators() {
+  for (const auto *region : regions()) {
+    const MemoryRegion *amd_mem_region(
+        static_cast<const MemoryRegion *>(region));
+    if (amd_mem_region->kernarg()) {
+      system_allocator_ =
+          [region](size_t size, size_t align,
+                   core::MemoryRegion::AllocateFlags alloc_flags) -> void * {
+        void *mem(nullptr);
+        return (core::Runtime::runtime_singleton_->AllocateMemory(
+                    region, size, alloc_flags, &mem) == HSA_STATUS_SUCCESS)
+                   ? mem
+                   : nullptr;
+      };
+      break;
+    }
+  }
+}
+
 } // namespace AMD
 } // namespace rocr

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp

@@ -208,6 +208,7 @@ hsa_status_t MemoryRegion::GetInfo(hsa_region_info_t attribute,
     case HSA_REGION_INFO_SEGMENT:
       switch (mem_props_.HeapType) {
         case HSA_HEAPTYPE_SYSTEM:
+        case HSA_HEAPTYPE_DEVICE_SVM:
         case HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE:
         case HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC:
           *((hsa_region_segment_t*)value) = HSA_REGION_SEGMENT_GLOBAL;
@@ -223,6 +224,7 @@ hsa_status_t MemoryRegion::GetInfo(hsa_region_info_t attribute,
     case HSA_REGION_INFO_GLOBAL_FLAGS:
       switch (mem_props_.HeapType) {
         case HSA_HEAPTYPE_SYSTEM:
+        case HSA_HEAPTYPE_DEVICE_SVM:
         case HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC:
         case HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE: {
           uint32_t ret = 0;
@@ -246,6 +248,7 @@ hsa_status_t MemoryRegion::GetInfo(hsa_region_info_t attribute,
     case HSA_REGION_INFO_ALLOC_MAX_SIZE:
       switch (mem_props_.HeapType) {
         case HSA_HEAPTYPE_SYSTEM:
+        case HSA_HEAPTYPE_DEVICE_SVM:
           *((size_t*)value) = max_sysmem_alloc_size_;
           break;
         case HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE:
@@ -260,6 +263,7 @@ hsa_status_t MemoryRegion::GetInfo(hsa_region_info_t attribute,
     case HSA_REGION_INFO_RUNTIME_ALLOC_ALLOWED:
       switch (mem_props_.HeapType) {
         case HSA_HEAPTYPE_SYSTEM:
+        case HSA_HEAPTYPE_DEVICE_SVM:
         case HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE:
         case HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC:
           *((bool*)value) = true;
@@ -272,6 +276,7 @@ hsa_status_t MemoryRegion::GetInfo(hsa_region_info_t attribute,
     case HSA_REGION_INFO_RUNTIME_ALLOC_GRANULE:
       switch (mem_props_.HeapType) {
         case HSA_HEAPTYPE_SYSTEM:
+        case HSA_HEAPTYPE_DEVICE_SVM:
         case HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE:
         case HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC:
           *((size_t*)value) = kPageSize_;
@@ -284,6 +289,7 @@ hsa_status_t MemoryRegion::GetInfo(hsa_region_info_t attribute,
     case HSA_REGION_INFO_RUNTIME_ALLOC_ALIGNMENT:
       switch (mem_props_.HeapType) {
         case HSA_HEAPTYPE_SYSTEM:
+        case HSA_HEAPTYPE_DEVICE_SVM:
         case HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE:
         case HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC:
           *((size_t*)value) = kPageSize_;