From a1597a358a9015412e9e337f9256e11142865888 Mon Sep 17 00:00:00 2001
From: David Yat Sin <77975354+dayatsin-amd@users.noreply.github.com>
Date: Fri, 22 Aug 2025 09:52:39 -0400
Subject: [PATCH] rocr: Expose flag to allocate uncached memory (#674)

Add new flag for clients to directly request uncached memory
---
 .../runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp  | 6 ++++++
 .../runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp        | 3 +++
 projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h | 5 ++++-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
index ba4700340d..7a1b80c775 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
@@ -241,6 +241,12 @@ KfdDriver::AllocateMemory(const core::MemoryRegion &mem_region,
             ? 1
             : kmt_alloc_flags.ui32.Uncached);
 
+  if (kmt_alloc_flags.ui32.Uncached) {
+    /* Uncached overwrites CoarseGrain and ExtendedCoherent */
+    kmt_alloc_flags.ui32.CoarseGrain = 0;
+    kmt_alloc_flags.ui32.ExtendedCoherent = 0;
+  }
+
   kmt_alloc_flags.ui32.ExecuteBlit =
     !!(alloc_flags & core::MemoryRegion::AllocateExecutableBlitKernelObject);
 
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 e3d2349e7f..47e3a2e3bb 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
@@ -881,6 +881,9 @@ hsa_status_t hsa_amd_memory_pool_allocate(hsa_amd_memory_pool_t memory_pool, siz
   if (flags & HSA_AMD_MEMORY_POOL_EXECUTABLE_FLAG)
     alloc_flag |= core::MemoryRegion::AllocateExecutable;
 
+  if (flags & HSA_AMD_MEMORY_POOL_UNCACHED_FLAG)
+    alloc_flag |= core::MemoryRegion::AllocateUncached;
+
 #ifdef SANITIZER_AMDGPU
   if (mem_region->owner()->device_type() == core::Agent::kAmdGpuDevice)
     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 ed44a44781..1c590946b3 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
@@ -1579,7 +1579,10 @@ typedef enum hsa_amd_memory_pool_flag_s {
    *  Allocates executable memory
    */
   HSA_AMD_MEMORY_POOL_EXECUTABLE_FLAG = (1 << 2),
-
+  /**
+   *  Allocates uncached memory
+   */
+  HSA_AMD_MEMORY_POOL_UNCACHED_FLAG = (1 << 3),
 } hsa_amd_memory_pool_flag_t;
 
 /**