From 7c92e51d68ac3f781d2fe6c770c2cd50f3c0a151 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Wed, 18 Aug 2021 17:58:00 +0000
Subject: [PATCH] SWDEV-298985 - Fix hipMemPrefetchAsync

Use the dst device specified to prefetch the data

Change-Id: Id3d2295e0ae500ba03031ef43473a2c29852e45f


[ROCm/clr commit: 037f161f2e73e1a464154f61b935f98fdba522ee]
---
 projects/clr/hipamd/src/hip_hmm.cpp | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_hmm.cpp b/projects/clr/hipamd/src/hip_hmm.cpp
index 4c4b93efd3..fd18b2b6f9 100644
--- a/projects/clr/hipamd/src/hip_hmm.cpp
+++ b/projects/clr/hipamd/src/hip_hmm.cpp
@@ -84,29 +84,33 @@ hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device,
   amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
   // Return error if count passed is more than the actual size allocated
   if (memObj && count  > (memObj->getSize() - offset)) {
-    return hipErrorInvalidValue;
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
+  if (device != hipCpuDeviceId && (static_cast<size_t>(device) >= g_devices.size())) {
+    HIP_RETURN(hipErrorInvalidDevice);
   }
 
   amd::HostQueue* queue = nullptr;
-  bool cpu_access = (device == hipCpuDeviceId) ? true : false;
+  amd::Device* dev = nullptr;
+  bool cpu_access = false;
 
   // Pick the specified stream or Null one from the provided device
-  if (stream != nullptr) {
-    queue = hip::getQueue(stream);
+  if (device == hipCpuDeviceId) {
+    cpu_access = true;
+    queue = (stream == nullptr) ? hip::getCurrentDevice()->NullStream() : hip::getQueue(stream);
   } else {
-    if (!cpu_access) {
-      queue = g_devices[device]->NullStream();
-    } else {
-      queue = hip::getCurrentDevice()->NullStream();
-    }
+    dev = g_devices[device]->devices()[0];
+    queue = (stream == nullptr) ? g_devices[device]->NullStream() : hip::getQueue(stream);
   }
+
   if (queue == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
 
   amd::Command::EventWaitList waitList;
   amd::SvmPrefetchAsyncCommand* command =
-      new amd::SvmPrefetchAsyncCommand(*queue, waitList, dev_ptr, count, cpu_access);
+      new amd::SvmPrefetchAsyncCommand(*queue, waitList, dev_ptr, count, dev, cpu_access);
   if (command == nullptr) {
     return hipErrorOutOfMemory;
   }