From 8ceb2ac5d6ed2132545227abcf1e2af3e61e0e88 Mon Sep 17 00:00:00 2001
From: Sourabh Betigeri <sourabh.betigeri@amd.com>
Date: Mon, 26 Sep 2022 22:40:28 -0700
Subject: [PATCH] SWDEV-357246 - Fixes to return with a error code instead of
 segfault in alloc failure case

Change-Id: Ie18d50f2f7ef3f4b2c0153152fca057232c1f8a2


[ROCm/clr commit: aeeab1c975da6e4635626849aa3acf94b903b9b2]
---
 projects/clr/hipamd/src/hip_hmm.cpp    |  3 +++
 projects/clr/hipamd/src/hip_memory.cpp | 16 +++++++++++-----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_hmm.cpp b/projects/clr/hipamd/src/hip_hmm.cpp
index 8a3cf7da16..652e4158c8 100644
--- a/projects/clr/hipamd/src/hip_hmm.cpp
+++ b/projects/clr/hipamd/src/hip_hmm.cpp
@@ -226,6 +226,9 @@ hipError_t ihipMallocManaged(void** ptr, size_t size, unsigned int align) {
   }
   size_t offset = 0; //this is ignored
   amd::Memory* memObj = getMemoryObject(*ptr, offset);
+  if (memObj == nullptr) {
+    return hipErrorMemoryAllocation;
+  }
   //saves the current device id so that it can be accessed later
   memObj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
 
diff --git a/projects/clr/hipamd/src/hip_memory.cpp b/projects/clr/hipamd/src/hip_memory.cpp
index 4be5e43e1c..75da3f0beb 100644
--- a/projects/clr/hipamd/src/hip_memory.cpp
+++ b/projects/clr/hipamd/src/hip_memory.cpp
@@ -43,13 +43,19 @@ amd::Memory* getMemoryObject(const void* ptr, size_t& offset, size_t size) {
     }
     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);
+      device::Memory* devMem = 
+              (memObj->getDeviceMemory(*memObj->getContext().devices()[0]));
+      if (devMem == nullptr) {
+        return nullptr;
       }
       else {
-        ShouldNotReachHere();
+        const void* devPtr = reinterpret_cast<void*>((devMem)->virtualAddress());
+        if (devPtr != nullptr) {
+          offset = reinterpret_cast<size_t>(ptr) - reinterpret_cast<size_t>(devPtr);
+        }
+        else {
+          return nullptr;
+        }
       }
     }
   } else {