From aedbad0109d330f97939c049681423bc4b08f41e Mon Sep 17 00:00:00 2001
From: Sarbojit Sarkar <sarbojit.sarkar@amd.com>
Date: Thu, 9 Dec 2021 07:31:03 +0000
Subject: [PATCH] SWDEV-314254 - Fix for hipMemcpy3D test crash

Change-Id: Iac70bfe0d351cfb5b56fefc9a6487d3f26f2b4ef
---
 rocclr/device/device.hpp         | 2 ++
 rocclr/device/rocm/rocmemory.hpp | 3 +++
 rocclr/platform/command.cpp      | 4 ++--
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/rocclr/device/device.hpp b/rocclr/device/device.hpp
index 7e4a64b449..ef2288819c 100644
--- a/rocclr/device/device.hpp
+++ b/rocclr/device/device.hpp
@@ -834,6 +834,8 @@ class Memory : public amd::HeapObject {
 
   virtual uint64_t virtualAddress() const { return 0; }
 
+  virtual uint64_t originalDeviceAddress() const { return 0; }
+
   //! Returns CPU pointer to HW state
   virtual const address cpuSrd() const { return nullptr; }
 
diff --git a/rocclr/device/rocm/rocmemory.hpp b/rocclr/device/rocm/rocmemory.hpp
index 3ec861f74d..91731207fe 100644
--- a/rocclr/device/rocm/rocmemory.hpp
+++ b/rocclr/device/rocm/rocmemory.hpp
@@ -102,6 +102,8 @@ class Memory : public device::Memory {
 
   virtual uint64_t virtualAddress() const override { return reinterpret_cast<uint64_t>(getDeviceMemory()); }
 
+  virtual uint64_t originalDeviceAddress() const { return virtualAddress(); }
+
   // Accessors for indirect map memory object
   amd::Memory* mapMemory() const { return mapMemory_; }
 
@@ -219,6 +221,7 @@ class Image : public roc::Memory {
 
   amd::Image* CopyImageBuffer() const { return copyImageBuffer_; }
 
+  virtual uint64_t originalDeviceAddress() const { return reinterpret_cast<uint64_t>(originalDeviceMemory_); }
  private:
   //! Disable copy constructor
   Image(const Buffer&);
diff --git a/rocclr/platform/command.cpp b/rocclr/platform/command.cpp
index 96c50245c8..f3e7793ef2 100644
--- a/rocclr/platform/command.cpp
+++ b/rocclr/platform/command.cpp
@@ -507,14 +507,14 @@ bool TwoMemoryArgsCommand::validatePeerMemory(){
     // explicit allow access is needed for P2P access
     device::Memory* mem1 = memory1_->getDeviceMemory(*srcDevices[0]);
     if (!mem1->getAllowedPeerAccess() && srcDevices.size() == 1) {
-      void* src = reinterpret_cast<void*>(mem1->virtualAddress());
+      void* src = reinterpret_cast<void*>(mem1->originalDeviceAddress());
       accessAllowed = srcDevices[0]->deviceAllowAccess(src);
       mem1->setAllowedPeerAccess(true);
     }
 
     device::Memory* mem2 = memory2_->getDeviceMemory(*dstDevices[0]);
     if (!mem2->getAllowedPeerAccess() && dstDevices.size() == 1) {
-      void* dst = reinterpret_cast<void*>(mem2->virtualAddress());
+      void* dst = reinterpret_cast<void*>(mem2->originalDeviceAddress());
       accessAllowed &= dstDevices[0]->deviceAllowAccess(dst);
       mem2->setAllowedPeerAccess(true);
     }