From 5b10098da9d41d57693ee8b2fd1390b60f335244 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Fri, 13 Nov 2020 14:09:36 +0000
Subject: [PATCH] SWDEV-258754 Fix seg fault hipMemset with device pointer
 mapped from hipHostRegister

Change-Id: Ifab66b67df172812ebb6eb25c2bac71821f4d614
---
 rocclr/hip_memory.cpp                         | 25 +++++++++++-----
 .../src/runtimeApi/memory/hipHostRegister.cpp | 29 +++++++++++++++++++
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/rocclr/hip_memory.cpp b/rocclr/hip_memory.cpp
index ca094cbe81..818cb23494 100755
--- a/rocclr/hip_memory.cpp
+++ b/rocclr/hip_memory.cpp
@@ -30,14 +30,23 @@
 amd::Memory* getMemoryObject(const void* ptr, size_t& offset) {
   amd::Memory *memObj = amd::MemObjMap::FindMemObj(ptr);
   if (memObj != nullptr) {
-    if (memObj->getSvmPtr() != nullptr) {
-      // SVM pointer
-      offset = reinterpret_cast<size_t>(ptr) - reinterpret_cast<size_t>(memObj->getSvmPtr());
-    } else if (memObj->getHostMem() != nullptr) {
-      // Prepinned memory
-      offset = reinterpret_cast<size_t>(ptr) - reinterpret_cast<size_t>(memObj->getHostMem());
-    } else {
-      ShouldNotReachHere();
+    const char* hostPtr = reinterpret_cast<const char*>(ptr);
+    const char* hostMem = reinterpret_cast<const char*>(memObj->getHostMem());
+    //Prepinned memory
+    if ((hostMem != nullptr) &&
+        (hostPtr >= hostMem && hostPtr <= (hostMem + memObj->getSize()))) {
+      offset = reinterpret_cast<size_t>(hostPtr) - reinterpret_cast<size_t>(hostMem);
+    }
+    else {
+      //SVM ptr or device ptr mapped from host
+      const void *devPtr = reinterpret_cast<void*>
+              (memObj->getDeviceMemory(*memObj->getContext().devices()[0])->virtualAddress());
+      if (devPtr != nullptr) {
+        offset = reinterpret_cast<size_t>(ptr) - reinterpret_cast<size_t>(devPtr);
+      }
+      else {
+        ShouldNotReachHere();
+      }
     }
   }
   return memObj;
diff --git a/tests/src/runtimeApi/memory/hipHostRegister.cpp b/tests/src/runtimeApi/memory/hipHostRegister.cpp
index 7957a50142..8000e80a1a 100644
--- a/tests/src/runtimeApi/memory/hipHostRegister.cpp
+++ b/tests/src/runtimeApi/memory/hipHostRegister.cpp
@@ -111,6 +111,35 @@ int main(int argc, char* argv[]) {
         delete [] Ad;
     }
 
+    if (p_tests & 0x3) {
+        float *A, **Ad;
+        int num_devices;
+        HIPCHECK(hipGetDeviceCount(&num_devices));
+        Ad = new float*[num_devices];
+        A = (float*)malloc(size);
+        HIPCHECK(hipHostRegister(A, size, 0));
+
+        for (int i = 0; i < N; i++) {
+            A[i] = float(1);
+        }
+
+        for (int i = 0; i < num_devices; i++) {
+            HIPCHECK(hipSetDevice(i));
+            HIPCHECK(hipHostGetDevicePointer((void**)&Ad[i], A, 0));
+        }
+
+        // Reference the registered device pointer Ad in hipMemset:
+        for (int i = 0; i < num_devices; i++) {
+            HIPCHECK(hipSetDevice(i));
+            HIPCHECK(hipMemset(Ad[i], 0, size));
+        }
+        HIPASSERT(A[10] == 0.0f);
+
+        HIPCHECK(hipHostUnregister(A));
+
+        free(A);
+        delete [] Ad;
+    }
 
     if (p_tests & 0x6) {
         // Sensitize HIP bug if device does not match where the memory was registered.