From 29adfcd3c088e9ff177ebd2185afb5be3b2fe6af Mon Sep 17 00:00:00 2001
From: agunashe <ajay.gunashekar@amd.com>
Date: Wed, 8 Sep 2021 11:01:02 -0700
Subject: [PATCH] SWDEV-301069 - current device id saved while allocating
 memory

Change-Id: I57045119e7adf915074c547cbe76349a4cfd72d9
---
 hipamd/src/hip_hmm.cpp    |  5 +++++
 hipamd/src/hip_memory.cpp | 24 +++++++++++-------------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/hipamd/src/hip_hmm.cpp b/hipamd/src/hip_hmm.cpp
index f79599a07a..0fab19bbbb 100644
--- a/hipamd/src/hip_hmm.cpp
+++ b/hipamd/src/hip_hmm.cpp
@@ -226,6 +226,11 @@ 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);
+  //saves the current device id so that it can be accessed later
+  memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
+
   if (*ptr == nullptr) {
     return hipErrorMemoryAllocation;
   }
diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index 04363363ec..8053de43cb 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -280,6 +280,11 @@ 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);
+  //saves the current device id so that it can be accessed later
+  memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
+
   if (*ptr == nullptr) {
     size_t free = 0, total =0;
     hipMemGetInfo(&free, &total);
@@ -649,6 +654,10 @@ 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);
+  //saves the current device id so that it can be accessed later
+  memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
 
   if (*ptr == nullptr) {
     return hipErrorOutOfMemory;
@@ -2473,19 +2482,8 @@ hipError_t hipPointerGetAttributes(hipPointerAttribute_t* attributes, const void
     attributes->allocationFlags = memObj->getMemFlags() >> 16;
 
     amd::Context* memObjCtx = &memObj->getContext();
-    if (hip::host_device->asContext() == memObjCtx) {
-        attributes->device = ihipGetDevice();
-        HIP_RETURN(hipSuccess);
-    }
-    for (auto& ctx : g_devices) {
-      if (ctx->asContext() == memObjCtx) {
-        attributes->device = device;
-        HIP_RETURN(hipSuccess);
-      }
-      ++device;
-    }
-    LogPrintfError("Cannot find memory object context, memObjCtx: 0x%x \n", memObjCtx);
-    HIP_RETURN(hipErrorInvalidDevice);
+    attributes->device = memObj->getUserData().deviceId;
+    HIP_RETURN(hipSuccess);
   }
 
   LogPrintfError("Cannot get amd_mem_obj for ptr: 0x%x \n", ptr);