From 361dec890cb2a85f79b918fa6391ad668a5c4085 Mon Sep 17 00:00:00 2001
From: Ajay <ajay.gunashekar@amd.com>
Date: Mon, 27 Sep 2021 16:38:30 -0400
Subject: [PATCH] SWDEV-301069 - save device id during hipHostRegister

Change-Id: I789b2252fa1a131ed1924b47718e355cb4fc9037
---
 hipamd/src/hip_hmm.cpp    |  4 ++--
 hipamd/src/hip_memory.cpp | 11 +++++++----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/hipamd/src/hip_hmm.cpp b/hipamd/src/hip_hmm.cpp
index 550daac0e0..73697c04c7 100644
--- a/hipamd/src/hip_hmm.cpp
+++ b/hipamd/src/hip_hmm.cpp
@@ -226,12 +226,12 @@ hipError_t ihipMallocManaged(void** ptr, size_t size, unsigned int align) {
   // allocation in the device driver
   *ptr = amd::SvmBuffer::malloc(ctx, CL_MEM_SVM_FINE_GRAIN_BUFFER | CL_MEM_ALLOC_HOST_PTR,
                                 size, (align == 0) ? dev.info().memBaseAddrAlign_ : align);
-  size_t offset = 0; //this is ignored
-  amd::Memory* memObj = getMemoryObject(*ptr, offset);
 
   if (*ptr == nullptr) {
     return hipErrorMemoryAllocation;
   }
+  size_t offset = 0; //this is ignored
+  amd::Memory* memObj = getMemoryObject(*ptr, offset);
   //saves the current device id so that it can be accessed later
   memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
 
diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index 340b18c5f2..e44cc42ed0 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -280,8 +280,6 @@ hipError_t ihipMalloc(void** ptr, size_t sizeBytes, unsigned int flags)
 
   *ptr = amd::SvmBuffer::malloc(*amdContext, flags, sizeBytes, amdContext->devices()[0]->info().memBaseAddrAlign_,
               useHostDevice ? curDevContext->svmDevices()[0] : nullptr);
-  size_t offset = 0; //this is ignored
-  amd::Memory* memObj = getMemoryObject(*ptr, offset);
 
   if (*ptr == nullptr) {
     size_t free = 0, total =0;
@@ -289,6 +287,8 @@ hipError_t ihipMalloc(void** ptr, size_t sizeBytes, unsigned int flags)
     LogPrintfError("Allocation failed : Device memory : required :%zu | free :%zu | total :%zu \n", sizeBytes, free, total);
     return hipErrorOutOfMemory;
   }
+  size_t offset = 0; //this is ignored
+  amd::Memory* memObj = getMemoryObject(*ptr, offset);
   //saves the current device id so that it can be accessed later
   memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
   return hipSuccess;
@@ -653,12 +653,12 @@ hipError_t ihipMallocPitch(void** ptr, size_t* pitch, size_t width, size_t heigh
 
   *ptr = amd::SvmBuffer::malloc(*hip::getCurrentDevice()->asContext(), 0, sizeBytes,
                                 device->info().memBaseAddrAlign_);
-  size_t offset = 0; //this is ignored
-  amd::Memory* memObj = getMemoryObject(*ptr, offset);
 
   if (*ptr == nullptr) {
     return hipErrorOutOfMemory;
   }
+  size_t offset = 0; //this is ignored
+  amd::Memory* memObj = getMemoryObject(*ptr, offset);
   //saves the current device id so that it can be accessed later
   memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
 
@@ -990,6 +990,9 @@ hipError_t hipHostRegister(void* hostPtr, size_t sizeBytes, unsigned int flags)
     }
 
     amd::MemObjMap::AddMemObj(hostPtr, mem);
+    if (mem != nullptr) {
+      mem->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
+    }
     HIP_RETURN(hipSuccess);
   } else {
     HIP_RETURN_DURATION(ihipMalloc(&hostPtr, sizeBytes, flags), hostPtr);