From 3d9d35a1f845ab485280819255eef85aa1e313b3 Mon Sep 17 00:00:00 2001
From: Jimbo <57198431+jiabaxie@users.noreply.github.com>
Date: Fri, 5 Sep 2025 10:31:20 -0400
Subject: [PATCH] SWDEV-553375 - Allow hipMemAllocationTypeUncached in
 hipMemGetAllocationGranularity (#847)

---
 projects/clr/hipamd/src/hip_vm.cpp            |  3 +-
 .../hip-tests/catch/unit/memory/hipMemVmm.cc  | 61 +++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/projects/clr/hipamd/src/hip_vm.cpp b/projects/clr/hipamd/src/hip_vm.cpp
index 37dbc8a219..48ccc9ca62 100644
--- a/projects/clr/hipamd/src/hip_vm.cpp
+++ b/projects/clr/hipamd/src/hip_vm.cpp
@@ -205,7 +205,8 @@ hipError_t hipMemGetAllocationGranularity(size_t* granularity, const hipMemAlloc
                                           hipMemAllocationGranularity_flags option) {
   HIP_INIT_API(hipMemGetAllocationGranularity, granularity, prop, option);
 
-  if (granularity == nullptr || prop == nullptr || prop->type != hipMemAllocationTypePinned ||
+  if (granularity == nullptr || prop == nullptr || (prop->type != hipMemAllocationTypePinned &&
+      prop->type != hipMemAllocationTypeUncached) ||
       prop->location.type != hipMemLocationTypeDevice || prop->location.id >= g_devices.size() ||
       (option != hipMemAllocationGranularityMinimum &&
        option != hipMemAllocationGranularityRecommended)) {
diff --git a/projects/hip-tests/catch/unit/memory/hipMemVmm.cc b/projects/hip-tests/catch/unit/memory/hipMemVmm.cc
index 73c062bde4..5872b6a65e 100644
--- a/projects/hip-tests/catch/unit/memory/hipMemVmm.cc
+++ b/projects/hip-tests/catch/unit/memory/hipMemVmm.cc
@@ -89,3 +89,64 @@ TEST_CASE("Unit_hipMemVmm_Basic") {
   HIP_CHECK(hipMemRelease(gaHandle));
   HIP_CHECK(hipMemAddressFree(reservedAddress, size));
 }
+
+/*
+    This testcase verifies HIP Mem VMM API basic scenario, but with Uncached memory -- supported
+    only on HIP
+ */
+
+#if HT_AMD
+TEST_CASE("Unit_hipMemVmm_Uncached") {
+  int vmm = 0;
+  HIP_CHECK(hipDeviceGetAttribute(&vmm, hipDeviceAttributeVirtualMemoryManagementSupported, 0));
+  INFO("hipDeviceAttributeVirtualMemoryManagementSupported: " << vmm);
+
+  if (vmm == 0) {
+    SUCCEED(
+        "GPU 0 doesn't support hipDeviceAttributeVirtualMemoryManagement "
+        "attribute. Hence skipping the testing with Pass result.\n");
+    return;
+  }
+
+  size_t granularity = 0;
+
+  hipMemAllocationProp memAllocationProp;
+  memAllocationProp.type = hipMemAllocationTypeUncached;
+  memAllocationProp.location.id = 0;
+  memAllocationProp.location.type = hipMemLocationTypeDevice;
+
+  HIP_CHECK(hipMemGetAllocationGranularity(&granularity, &memAllocationProp,
+                                           hipMemAllocationGranularityRecommended));
+
+  size_t size = 4 * 1024;
+  void* reservedAddress{nullptr};
+  HIP_CHECK(hipMemAddressReserve(&reservedAddress, size, granularity, nullptr, 0));
+
+  hipMemGenericAllocationHandle_t gaHandle{nullptr};
+  HIP_CHECK(hipMemCreate(&gaHandle, size, &memAllocationProp, 0));
+
+  HIP_CHECK(hipMemMap(reservedAddress, size, 0, gaHandle, 0));
+
+  hipDevice_t device;
+  HIP_CHECK(hipDeviceGet(&device, 0));
+  hipMemAccessDesc desc;
+  desc.location.type = hipMemLocationTypeDevice;
+  desc.location.id = device;
+  desc.flags = hipMemAccessFlagsProtReadWrite;
+  std::vector<char> values(size);
+  const char value = 1;
+
+  HIP_CHECK(hipMemSetAccess(reservedAddress, size, &desc, 1));
+  HIP_CHECK(hipMemset(reservedAddress, value, size));
+  HIP_CHECK(hipMemcpy(&values[0], reservedAddress, size, hipMemcpyDeviceToHost));
+
+  for (size_t i = 0; i < size; ++i) {
+    REQUIRE(values[i] == value);
+  }
+
+  HIP_CHECK(hipMemUnmap(reservedAddress, size));
+
+  HIP_CHECK(hipMemRelease(gaHandle));
+  HIP_CHECK(hipMemAddressFree(reservedAddress, size));
+}
+#endif
\ No newline at end of file