From b53648f8fe6d765bf063a33686434475cdb89599 Mon Sep 17 00:00:00 2001
From: David Yat Sin <David.YatSin@amd.com>
Date: Wed, 7 Feb 2024 17:36:41 +0000
Subject: [PATCH] Add support for contiguous memory allocations

Support contiguous physical memory allocation flag. Allocations with
this flag will have contiguous physical memory. This is dependent on KFD
support for this flag and the AllocateKfdMemory(..) function call will
fail when it is not supported.

Change-Id: I6c51c8b061f7b026fdcc2aa2c37c74ecc13d95b6


[ROCm/ROCR-Runtime commit: 9af225e1b1a2159ee0ef326c7749c0f8b6cb5e8b]
---
 .../runtime/hsa-runtime/core/inc/memory_region.h          | 1 +
 .../hsa-runtime/core/runtime/amd_memory_region.cpp        | 6 ++++++
 .../runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp      | 8 ++++++--
 .../rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h    | 6 +++++-
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/memory_region.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/memory_region.h
index 6cd127b894..66acf36362 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/memory_region.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/memory_region.h
@@ -104,6 +104,7 @@ class MemoryRegion : public Checked<0x9C961F19EE175BB3> {
     // Note: The node_id needs to be the node_id of the device even though this is allocating
     // system memory
     AllocateGTTAccess = (1 << 9),
+    AllocateContiguous = (1 << 10),  // Physically contiguous memory
   };
 
   typedef uint32_t AllocateFlags;
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp
index 2139ca0a67..40e1d6821c 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp
@@ -209,6 +209,12 @@ hsa_status_t MemoryRegion::AllocateImpl(size_t& size, AllocateFlags alloc_flags,
   kmt_alloc_flags.ui32.NoSubstitute = (alloc_flags & AllocatePinned ? 1 : kmt_alloc_flags.ui32.NoSubstitute);
 
   kmt_alloc_flags.ui32.GTTAccess = (alloc_flags & AllocateGTTAccess ? 1 : kmt_alloc_flags.ui32.GTTAccess);
+  if (IsLocalMemory()) {
+    // Allocate physically contiguous memory - AllocateKfdMemory function call will fail
+    // if this flag is not supported in KFD.
+    kmt_alloc_flags.ui32.Contiguous =
+        (alloc_flags & AllocateContiguous ? 1 : kmt_alloc_flags.ui32.Contiguous);
+  }
 
   // Only allow using the suballocator for ordinary VRAM.
   if (IsLocalMemory() && !kmt_alloc_flags.ui32.NoAddress) {
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp
index 60d2df0712..dbd1286735 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp
@@ -761,7 +761,7 @@ hsa_status_t hsa_amd_memory_pool_allocate(hsa_amd_memory_pool_t memory_pool, siz
   TRY;
   IS_OPEN();
 
-  if (size == 0 || ptr == NULL || (flags > HSA_AMD_MEMORY_POOL_PCIE_FLAG)) {
+  if (size == 0 || ptr == NULL) {
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
 
@@ -774,7 +774,11 @@ hsa_status_t hsa_amd_memory_pool_allocate(hsa_amd_memory_pool_t memory_pool, siz
 
   MemoryRegion::AllocateFlags alloc_flag = core::MemoryRegion::AllocateRestrict;
 
-  if (flags == HSA_AMD_MEMORY_POOL_PCIE_FLAG) alloc_flag |= core::MemoryRegion::AllocatePCIeRW;
+  if (flags & HSA_AMD_MEMORY_POOL_PCIE_FLAG)
+    alloc_flag |= core::MemoryRegion::AllocatePCIeRW;
+
+  if (flags & HSA_AMD_MEMORY_POOL_CONTIGUOUS_FLAG)
+    alloc_flag |= core::MemoryRegion::AllocateContiguous;
 
 #ifdef SANITIZER_AMDGPU
   alloc_flag |= core::MemoryRegion::AllocateAsan;
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h b/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h
index 40630e2ece..a162db37de 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h
@@ -1182,7 +1182,11 @@ typedef enum hsa_amd_memory_pool_flag_s {
    * connection. Atomic memory operations on these memory buffers are not
    * guaranteed to be visible at system scope.
    */
-  HSA_AMD_MEMORY_POOL_PCIE_FLAG = 1,
+  HSA_AMD_MEMORY_POOL_PCIE_FLAG = (1 << 0),
+  /**
+   *  Allocates physically contiguous memory
+   */
+  HSA_AMD_MEMORY_POOL_CONTIGUOUS_FLAG = (1 << 1),
 
 } hsa_amd_memory_pool_flag_t;