From 67af16bdf356008b10062d510deec090e60976bb Mon Sep 17 00:00:00 2001
From: German Andryeyev <German.Andryeyev@amd.com>
Date: Wed, 8 Feb 2023 17:55:08 -0500
Subject: [PATCH] SWDEV-381633 - Use blit manager on app's queue

Heap initialization used device queue, but it shoudl be used for
cooperative launches only. Heap initialization must use the same queue
as the current dispatch.

Change-Id: I856621bf82bbdeb1c2d0fbc4970e90d09af805cb
---
 rocclr/device/pal/paldevice.cpp   | 6 +++---
 rocclr/device/pal/paldevice.hpp   | 2 +-
 rocclr/device/pal/palkernel.cpp   | 2 +-
 rocclr/device/rocm/rocdevice.cpp  | 6 +++---
 rocclr/device/rocm/rocdevice.hpp  | 2 +-
 rocclr/device/rocm/rocvirtual.cpp | 2 +-
 6 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/rocclr/device/pal/paldevice.cpp b/rocclr/device/pal/paldevice.cpp
index 6c60ef0524..5755f34a81 100644
--- a/rocclr/device/pal/paldevice.cpp
+++ b/rocclr/device/pal/paldevice.cpp
@@ -2356,8 +2356,8 @@ void Device::ReleaseExclusiveGpuAccess(VirtualGPU& vgpu) const {
 }
 
 // ================================================================================================
-void Device::HiddenHeapAlloc() {
-  auto HeapAlloc = [this]() -> bool {
+void Device::HiddenHeapAlloc(const VirtualGPU& gpu) {
+  auto HeapAlloc = [this, &gpu]() -> bool {
     // Allocate initial heap for device memory allocator
     static constexpr size_t HeapBufferSize = 128 * Ki;
     heap_buffer_ = createMemory(HeapBufferSize);
@@ -2369,7 +2369,7 @@ void Device::HiddenHeapAlloc() {
       LogError("Heap buffer allocation failed!");
       return false;
     }
-    bool result = static_cast<const KernelBlitManager&>(xferMgr()).initHeap(
+    bool result = static_cast<const KernelBlitManager&>(gpu.blitMgr()).initHeap(
         heap_buffer_, initial_heap_buffer_, HeapBufferSize, initial_heap_size_ / (2 * Mi));
 
     return result;
diff --git a/rocclr/device/pal/paldevice.hpp b/rocclr/device/pal/paldevice.hpp
index 5b3b4af5e0..0226ea4005 100644
--- a/rocclr/device/pal/paldevice.hpp
+++ b/rocclr/device/pal/paldevice.hpp
@@ -629,7 +629,7 @@ class Device : public NullDevice {
 #endif
 #endif
   //! Allocates hidden heap for device memory allocations
-  void HiddenHeapAlloc();
+  void HiddenHeapAlloc(const VirtualGPU& gpu);
 
  private:
   static void PAL_STDCALL PalDeveloperCallback(void* pPrivateData, const Pal::uint32 deviceIndex,
diff --git a/rocclr/device/pal/palkernel.cpp b/rocclr/device/pal/palkernel.cpp
index e6e4fd2816..f425e3d7dc 100644
--- a/rocclr/device/pal/palkernel.cpp
+++ b/rocclr/device/pal/palkernel.cpp
@@ -362,7 +362,7 @@ hsa_kernel_dispatch_packet_t* HSAILKernel::loadArguments(VirtualGPU& gpu, const
       case amd::KernelParameterDescriptor::HiddenHeap:
         // Allocate hidden heap for HIP applications only
         if ((amd::IS_HIP) && (palDevice().HeapBuffer() == nullptr)) {
-          const_cast<Device&>(palDevice()).HiddenHeapAlloc();
+          const_cast<Device&>(palDevice()).HiddenHeapAlloc(gpu);
         }
         if (palDevice().HeapBuffer() != nullptr) {
           // Add heap pointer to the code
diff --git a/rocclr/device/rocm/rocdevice.cpp b/rocclr/device/rocm/rocdevice.cpp
index f9df3075cd..9ecaf22ef8 100644
--- a/rocclr/device/rocm/rocdevice.cpp
+++ b/rocclr/device/rocm/rocdevice.cpp
@@ -3243,8 +3243,8 @@ bool Device::IsValidAllocation(const void* dev_ptr, size_t size) const {
 }
 
 // ================================================================================================
-void Device::HiddenHeapAlloc() {
-  auto HeapAllocZeroOut = [this]() -> bool {
+void Device::HiddenHeapAlloc(const VirtualGPU& gpu) {
+  auto HeapAllocZeroOut = [this, &gpu]() -> bool {
     // Allocate initial heap for device memory allocator
     static constexpr size_t HeapBufferSize = 128 * Ki;
     heap_buffer_ = createMemory(HeapBufferSize);
@@ -3256,7 +3256,7 @@ void Device::HiddenHeapAlloc() {
       LogError("Heap buffer allocation failed!");
       return false;
     }
-    bool result = static_cast<const KernelBlitManager&>(xferMgr()).initHeap(
+    bool result = static_cast<const KernelBlitManager&>(gpu.blitMgr()).initHeap(
         heap_buffer_, initial_heap_buffer_, HeapBufferSize, initial_heap_size_ / (2 * Mi));
 
     return result;
diff --git a/rocclr/device/rocm/rocdevice.hpp b/rocclr/device/rocm/rocdevice.hpp
index 99d0976e13..d9eb4f204d 100644
--- a/rocclr/device/rocm/rocdevice.hpp
+++ b/rocclr/device/rocm/rocdevice.hpp
@@ -547,7 +547,7 @@ class Device : public NullDevice {
   bool IsValidAllocation(const void* dev_ptr, size_t size) const;
 
   //! Allocates hidden heap for device memory allocations
-  void HiddenHeapAlloc();
+  void HiddenHeapAlloc(const VirtualGPU& gpu);
 
  private:
   bool create();
diff --git a/rocclr/device/rocm/rocvirtual.cpp b/rocclr/device/rocm/rocvirtual.cpp
index 0976b18376..3b88b564ba 100644
--- a/rocclr/device/rocm/rocvirtual.cpp
+++ b/rocclr/device/rocm/rocvirtual.cpp
@@ -2958,7 +2958,7 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes,
         case amd::KernelParameterDescriptor::HiddenHeap:
           // Allocate hidden heap for HIP applications only
           if ((amd::IS_HIP) && (dev().HeapBuffer() == nullptr)) {
-            const_cast<Device&>(dev()).HiddenHeapAlloc();
+            const_cast<Device&>(dev()).HiddenHeapAlloc(*this);
           }
           if (dev().HeapBuffer() != nullptr) {
             // Add heap pointer to the code