From 8e8dc41cf0c6167bb0be49e8052a6f1de85a63ef Mon Sep 17 00:00:00 2001
From: "Manocha, Rahul" <rahul.manocha@amd.com>
Date: Mon, 7 Jul 2025 17:27:13 -0700
Subject: [PATCH] SWDEV-532420 - Fix kokkos P2P copy failure with vmheap (#426)

Co-authored-by: Rahul Manocha <rmanocha@amd.com>

[ROCm/clr commit: 22b1ca4d8c192f73240419b3c679720d526ed73e]
---
 projects/clr/rocclr/platform/memory.cpp |  6 +++---
 projects/clr/rocclr/platform/memory.hpp | 15 ++++++++-------
 projects/clr/rocclr/platform/vmheap.cpp |  3 ++-
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/projects/clr/rocclr/platform/memory.cpp b/projects/clr/rocclr/platform/memory.cpp
index b03f1a8b7d..24e6bcf567 100644
--- a/projects/clr/rocclr/platform/memory.cpp
+++ b/projects/clr/rocclr/platform/memory.cpp
@@ -103,11 +103,11 @@ Memory::Memory(Context& context, Type type, Flags flags, size_t size, void* svmP
   canBeCached_ = true;
 }
 
-Memory::Memory(Memory& parent, Flags flags, size_t origin, size_t size, Type type)
+Memory::Memory(Memory& parent, Flags flags, size_t origin, size_t size, Type type, Context* new_ctx)
     : numDevices_(0),
       deviceMemories_(NULL),
       destructorCallbacks_(NULL),
-      context_(parent.getContext()),
+      context_((new_ctx != nullptr) ? *new_ctx : parent.getContext()),
       parent_(&parent),
       type_((type == 0) ? parent.type_ : type),
       hostMemRef_(NULL),
@@ -122,7 +122,7 @@ Memory::Memory(Memory& parent, Flags flags, size_t origin, size_t size, Type typ
       svmHostAddress_(parent.getSvmPtr()),
       resOffset_(0),
       flagsEx_(0),
-      lockMemoryOps_(true)  /* Memory Ops Lock */ {
+      lockMemoryOps_(true) /* Memory Ops Lock */ {
   svmPtrCommited_ = parent.isSvmPtrCommited();
   canBeCached_ = true;
   parent_->retain();
diff --git a/projects/clr/rocclr/platform/memory.hpp b/projects/clr/rocclr/platform/memory.hpp
index 214090f048..fe08a5985b 100644
--- a/projects/clr/rocclr/platform/memory.hpp
+++ b/projects/clr/rocclr/platform/memory.hpp
@@ -238,11 +238,12 @@ class Memory : public amd::RuntimeObject {
          void* svmPtr = NULL, //!< svm host memory address, NULL if no SVM mem object
          size_t alignment = 0 //!< allocation addr alignment
   );
-  Memory(Memory& parent,  //!< Context object
-         Flags flags,     //!< Object's flags
-         size_t offset,   //!< Memory offset
-         size_t size,     //!< Memory size
-         Type type = 0    //!< Memory type
+  Memory(Memory& parent,             //!< Parent Mem obj
+         Flags flags,                //!< Object's flags
+         size_t offset,              //!< Memory offset
+         size_t size,                //!< Memory size
+         Type type = 0,              //!< Memory type
+         Context* context = nullptr  //!< Input context
   );
 
   //! Memory object destructor
@@ -449,8 +450,8 @@ class Buffer : public Memory {
  public:
   Buffer(Context& context, Flags flags, size_t size, void* svmPtr = NULL, size_t alignment = 0)
       : Memory(context, CL_MEM_OBJECT_BUFFER, flags, size, svmPtr, alignment) {}
-  Buffer(Memory& parent, Flags flags, size_t origin, size_t size)
-      : Memory(parent, flags, origin, size) {}
+  Buffer(Memory& parent, Flags flags, size_t origin, size_t size, Context* context = nullptr)
+      : Memory(parent, flags, origin, size, 0, context) {}
 
   bool create(void* initFrom = NULL,     //!< Pointer to the initialization data
               bool sysMemAlloc = false,  //!< Allocate device memory in system memory
diff --git a/projects/clr/rocclr/platform/vmheap.cpp b/projects/clr/rocclr/platform/vmheap.cpp
index cebfda35a9..37c59979a8 100644
--- a/projects/clr/rocclr/platform/vmheap.cpp
+++ b/projects/clr/rocclr/platform/vmheap.cpp
@@ -246,7 +246,8 @@ address VmHeap::Alloc(size_t size) {
   } else {
     return nullptr;
   }
-  auto memory = new (base_memory_->getContext()) Buffer(*base_memory_, 0, offset, size);
+  auto memory =
+      new (device_->context()) Buffer(*base_memory_, 0, offset, size, &device_->context());
   if (nullptr == memory || !memory->create(nullptr)) {
     FreeBlock(hb);
     return nullptr;