From 17807d78bd83527326cb0f693cb5a427b45b9067 Mon Sep 17 00:00:00 2001
From: Yiannis Papadopoulos <Yiannis.Papadopoulos@amd.com>
Date: Tue, 28 Jan 2025 00:34:46 +0000
Subject: [PATCH] rocr/aie: AIE agent memory pools correct size and user data
 pool

Change-Id: I831711a7d1cdc36cbc9ed30bd74d0dc984228ce7


[ROCm/ROCR-Runtime commit: 1d8a77db3481b9a1042138731f79c7c0586cb652]
---
 .../core/driver/xdna/amd_xdna_driver.cpp            |  9 ++++++++-
 .../runtime/hsa-runtime/core/inc/amd_xdna_driver.h  |  3 +++
 .../hsa-runtime/core/runtime/amd_aie_agent.cpp      | 13 +++++++++++--
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
index b92b66d09b..c37c969066 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
@@ -45,6 +45,7 @@
 #include <fcntl.h>
 #include <sys/ioctl.h>
 #include <sys/mman.h>
+#include <unistd.h>
 
 #include <memory>
 #include <string>
@@ -68,7 +69,7 @@ XdnaDriver::XdnaDriver(std::string devnode_name)
 
 hsa_status_t XdnaDriver::DiscoverDriver(std::unique_ptr<core::Driver>& driver) {
   const int max_minor_num(64);
-  const std::string devnode_prefix("/dev/accel/accel");
+  static const std::string devnode_prefix("/dev/accel/accel");
 
   for (int i = 0; i < max_minor_num; ++i) {
     auto tmp_driver = std::unique_ptr<Driver>(new XdnaDriver(devnode_prefix + std::to_string(i)));
@@ -86,6 +87,12 @@ hsa_status_t XdnaDriver::DiscoverDriver(std::unique_ptr<core::Driver>& driver) {
   return HSA_STATUS_ERROR;
 }
 
+uint64_t XdnaDriver::GetSystemMemoryByteSize() {
+  const long pagesize = sysconf(_SC_PAGESIZE);
+  const long page_count = sysconf(_SC_PHYS_PAGES);
+  return pagesize * page_count;
+}
+
 uint64_t XdnaDriver::GetDevHeapByteSize() {
   return dev_heap_size;
 }
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
index a170f72000..080bfc5202 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
@@ -133,6 +133,9 @@ public:
 
   static hsa_status_t DiscoverDriver(std::unique_ptr<core::Driver>& driver);
 
+  /// @brief Returns the size of the system memory heap in bytes.
+  static uint64_t GetSystemMemoryByteSize();
+
   /// @brief Returns the size of the dev heap in bytes.
   static uint64_t GetDevHeapByteSize();
 
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp
index 3c823b464f..d655d0fa2f 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp
@@ -282,25 +282,34 @@ void AieAgent::InitRegionList() {
   ///       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.
+  const uint64_t total_system_memory = XdnaDriver::GetSystemMemoryByteSize();
 
   /// For allocating kernel arguments or other objects that only need
   /// system memory.
   HsaMemoryProperties sys_mem_props = {};
   sys_mem_props.HeapType = HSA_HEAPTYPE_SYSTEM;
+  sys_mem_props.SizeInBytes = total_system_memory;
+
+  /// For any other allocation, e.g., buffers.
+  HsaMemoryProperties other_mem_props = {};
+  other_mem_props.HeapType = HSA_HEAPTYPE_SYSTEM;
+  other_mem_props.SizeInBytes = total_system_memory;
 
   /// 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 = {};
-  dev_mem_props.HeapType = HSA_HEAPTYPE_DEVICE_SVM,
+  dev_mem_props.HeapType = HSA_HEAPTYPE_DEVICE_SVM;
   dev_mem_props.SizeInBytes = XdnaDriver::GetDevHeapByteSize();
 
   /// As of now the AIE devices support coarse-grain memory regions that require
   /// explicit sync operations.
-  regions_.reserve(2);
+  regions_.reserve(3);
   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));
+  regions_.push_back(new MemoryRegion(false, false, false, false, true, this,
+                                      other_mem_props));
 }
 
 void AieAgent::GetAgentProperties() {