SWDEV-532420 - Fix kokkos P2P copy failure with vmheap (#426)

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

[ROCm/clr commit: 22b1ca4d8c]

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();

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

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;