From 59c0ec10cce6e63f9693255af9e785f8039e083c Mon Sep 17 00:00:00 2001
From: Julia Jiang <julia.jiang@amd.com>
Date: Fri, 23 Sep 2022 15:34:35 -0400
Subject: [PATCH] SWDEV-352687 - fix AlignedAllocMem in hmm

Change-Id: Icec74ac732a0a6ac2a2600f504aa1ec73717c0ca


[ROCm/clr commit: c9ce9b3300025c7af01066104fa5c6ee0d21010b]
---
 projects/clr/hipamd/src/hip_hmm.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_hmm.cpp b/projects/clr/hipamd/src/hip_hmm.cpp
index 652e4158c8..5b9dc5ed43 100644
--- a/projects/clr/hipamd/src/hip_hmm.cpp
+++ b/projects/clr/hipamd/src/hip_hmm.cpp
@@ -87,10 +87,9 @@ hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device,
   size_t offset = 0;
   amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
 
-  if (memObj == nullptr || (count  > (memObj->getSize() - offset))) {
+  if ((memObj != nullptr) && (count  > (memObj->getSize() - offset))) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-
   if (device != hipCpuDeviceId && (static_cast<size_t>(device) >= g_devices.size())) {
     HIP_RETURN(hipErrorInvalidDevice);
   }
@@ -99,6 +98,11 @@ hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device,
   amd::Device* dev = nullptr;
   bool cpu_access = false;
 
+  if ((memObj == nullptr) && (device != hipCpuDeviceId) &&
+      (!g_devices[device]->devices()[0]->info().hmmCpuMemoryAccessible_)) {
+    HIP_RETURN(hipErrorNotSupported);
+  }
+
   // Pick the specified stream or Null one from the provided device
   if (device == hipCpuDeviceId) {
     cpu_access = true;