SWDEV-364576 - initialize device malloc heap state using blit kernel

Change-Id: I5d0172aff7d2c04b322a4d828b8a2b438158b80f

rocclr/device/blit.hpp

@@ -239,6 +239,12 @@ class BlitManager : public amd::HeapObject {
   //! Returns Xfer queue lock
   virtual amd::Monitor* lockXfer() const { return nullptr; }
 
+  virtual bool initHeap(device::Memory* heap_to_initialize,
+                        device::Memory* initial_blocks,
+                        uint heap_size,
+                        uint number_of_initial_blocks
+                        ) const = 0;
+
  protected:
   const Setup setup_;   //!< HW accelerated blit requested
   bool syncOperation_;  //!< Blit operations are synchronized

rocclr/device/device.cpp

@@ -755,6 +755,14 @@ bool Device::UpdateStackSize(uint64_t stackSize) {
   return true;
 }
 
+bool Device::UpdateInitialHeapSize(uint64_t initialHeapSize) {
+  if (initialHeapSize >= info().globalMemSize_) {
+    return false;
+  }
+  initial_heap_size_ = initialHeapSize;
+  return true;
+}
+
 char* Device::getExtensionString() {
   std::stringstream extStream;
   size_t size;

rocclr/device/device.hpp

@@ -1905,6 +1905,12 @@ class Device : public RuntimeObject {
   //! Sets the stack size of the device
   bool UpdateStackSize(uint64_t stackSize);
 
+  //! Returns initial heap size
+  uint64_t InitialHeapSize() const { return initial_heap_size_; }
+
+  //! Sets the heap size of the device
+  bool UpdateInitialHeapSize(uint64_t initialHeapSize);
+
   //! Does this device allow P2P access?
   bool P2PAccessAllowed() const { return (p2p_access_devices_.size() > 0) ? true : false; }
 
@@ -1964,7 +1970,8 @@ class Device : public RuntimeObject {
 
   amd::Memory* arena_mem_obj_;      //!< Arena memory object
   uint64_t stack_size_{1024};       //!< Device stack size
-
+  device::Memory* initial_heap_buffer_;   //!< Initial heap buffer
+  uint64_t initial_heap_size_{HIP_INITIAL_DM_SIZE};  //!< Initial device heap size
  private:
   const Isa *isa_;                //!< Device isa
   bool IsTypeMatching(cl_device_type type, bool offlineDevices);

rocclr/device/pal/palblit.cpp

@@ -2520,10 +2520,33 @@ bool KernelBlitManager::streamOpsWait(device::Memory& memory, uint64_t value, si
   return result;
 }
 
+// ================================================================================================
+bool KernelBlitManager::initHeap(device::Memory* heap_to_initialize, device::Memory* initial_blocks,
+                                 uint heap_size, uint number_of_initial_blocks) const {
+  bool result;
+  // Clear memory to 0 for device library logic and set
+  size_t globalWorkOffset[1] = {0};
+  size_t globalWorkSize[1] = {256};
+  size_t localWorkSize[1] = {256};
 
+  // Create ND range object for the kernel's execution
+  amd::NDRangeContainer ndrange(1, globalWorkOffset, globalWorkSize, localWorkSize);
+  uint blitType = InitHeap;
+  uint64_t management_heap_va = heap_to_initialize->virtualAddress();
+  uint64_t initial_heap_va = 0;
+  if (initial_blocks != nullptr) {
+    initial_heap_va = initial_blocks->virtualAddress();
+  }
+  setArgument(kernels_[blitType], 0, sizeof(cl_ulong), &management_heap_va);
+  setArgument(kernels_[blitType], 1, sizeof(cl_ulong), &initial_heap_va);
+  setArgument(kernels_[blitType], 2, sizeof(uint), &heap_size);
+  setArgument(kernels_[blitType], 3, sizeof(uint), &number_of_initial_blocks);
+  address parameters = kernels_[blitType]->parameters().values();
+  result = gpu().submitKernelInternal(ndrange, *kernels_[blitType], parameters);
+  synchronize();
 
-
-
+  return result;
+}
 
 bool KernelBlitManager::runScheduler(device::Memory& vqueue, device::Memory& params, uint paramIdx,
                                      uint threads) const {

rocclr/device/pal/palblit.hpp

@@ -178,28 +178,32 @@ class DmaBlitManager : public device::HostBlitManager {
                          ) const;
 
   //! Stream memory write operation - Write a 'value' at 'memory'.
-  virtual bool streamOpsWrite(device::Memory& memory, //!< Memory to write the 'value'
-                             uint64_t value,
-                             size_t offset,
-                             size_t sizeBytes
-  ) const {
+  virtual bool streamOpsWrite(device::Memory& memory,  //!< Memory to write the 'value'
+                              uint64_t value,
+                              size_t offset,
+                              size_t sizeBytes) const {
     assert(!"Unimplemented");
     return false;
-  };
+  }
 
   //! Stream memory ops- Waits for a 'value' at 'memory' and wait is released based on compare op.
-  virtual bool streamOpsWait(device::Memory& memory, //!< Memory contents to compare the 'value' against
+  virtual bool streamOpsWait(device::Memory& memory,  //!< Memory to compare the 'value' against
                              uint64_t value,
                              size_t offset,
                              size_t sizeBytes,
                              uint64_t flags,
-                             uint64_t mask
-  ) const {
+                             uint64_t mask) const {
     assert(!"Unimplemented");
     return false;
-  };
-
+  }
 
+  virtual bool initHeap(device::Memory* heap_to_initialize,
+                        device::Memory* initial_blocks,
+                        uint heap_size,
+                        uint number_of_initial_blocks) const {
+    assert(!"Unimplemented");
+    return false;
+  }
 
  protected:
   static constexpr uint MaxPinnedBuffers = 4;
@@ -271,6 +275,7 @@ class KernelBlitManager : public DmaBlitManager {
     GwsInit,
     StreamOpsWrite,
     StreamOpsWait,
+    InitHeap,
     BlitTotal,
   };
 
@@ -464,6 +469,11 @@ class KernelBlitManager : public DmaBlitManager {
                              uint64_t mask
   ) const;
 
+  virtual bool initHeap(device::Memory* heap_to_initialize,
+                        device::Memory* initial_blocks,
+                        uint heap_size,
+                        uint number_of_initial_blocks
+                        ) const;
 
  private:
   static constexpr size_t MaxXferBuffers = 2;
@@ -530,7 +540,7 @@ static const char* BlitName[KernelBlitManager::BlitTotal] = {
     "__amd_rocclr_copyBufferAligned", "__amd_rocclr_fillBufferAligned",
     "__amd_rocclr_fillImage", "__amd_rocclr_scheduler",
     "__amd_rocclr_gwsInit", "__amd_rocclr_streamOpsWrite",
-    "__amd_rocclr_streamOpsWait"
+    "__amd_rocclr_streamOpsWait", "__amd_rocclr_initHeap"
 };
 
 /*@}*/  // namespace pal

rocclr/device/pal/paldevice.cpp

@@ -2343,15 +2343,22 @@ void Device::ReleaseExclusiveGpuAccess(VirtualGPU& vgpu) const {
 
 // ================================================================================================
 void Device::HiddenHeapAlloc() {
-  auto HeapAlloc = [this]()->bool {
+  auto HeapAlloc = [this]() -> bool {
     // Allocate initial heap for device memory allocator
     static constexpr size_t HeapBufferSize = 128 * Ki;
     heap_buffer_ = createMemory(HeapBufferSize);
+    if (initial_heap_size_ != 0) {
+      initial_heap_size_ = amd::alignUp(initial_heap_size_, 2 * Mi);
+      initial_heap_buffer_ = createMemory(initial_heap_size_);
+    }
     if (heap_buffer_ == nullptr) {
       LogError("Heap buffer allocation failed!");
       return false;
     }
-    return true;
+    bool result = static_cast<const KernelBlitManager&>(xferMgr()).initHeap(
+        heap_buffer_, initial_heap_buffer_, HeapBufferSize, initial_heap_size_ / (2 * Mi));
+
+    return result;
   };
   std::call_once(heap_initialized_, HeapAlloc);
 }

rocclr/device/rocm/rocblit.cpp

@@ -2542,6 +2542,36 @@ bool KernelBlitManager::streamOpsWait(device::Memory& memory, uint64_t value, si
 
   return result;
 }
+
+// ================================================================================================
+bool KernelBlitManager::initHeap(device::Memory* heap_to_initialize, device::Memory* initial_blocks,
+                                 uint heap_size, uint number_of_initial_blocks) const {
+  bool result;
+  // Clear memory to 0 for device library logic and set
+  size_t globalWorkOffset[1] = {0};
+  size_t globalWorkSize[1] = {256};
+  size_t localWorkSize[1] = {256};
+
+  // Create ND range object for the kernel's execution
+  amd::NDRangeContainer ndrange(1, globalWorkOffset, globalWorkSize, localWorkSize);
+  uint blitType = InitHeap;
+  uint64_t management_heap_va = heap_to_initialize->virtualAddress();
+  uint64_t initial_heap_va = 0;
+  if (initial_blocks != nullptr) {
+    initial_heap_va = initial_blocks->virtualAddress();
+  }
+  setArgument(kernels_[blitType], 0, sizeof(cl_ulong), &management_heap_va);
+  setArgument(kernels_[blitType], 1, sizeof(cl_ulong), &initial_heap_va);
+  setArgument(kernels_[blitType], 2, sizeof(uint), &heap_size);
+  setArgument(kernels_[blitType], 3, sizeof(uint), &number_of_initial_blocks);
+  address parameters = captureArguments(kernels_[blitType]);
+  result = gpu().submitKernelInternal(ndrange, *kernels_[blitType], parameters, nullptr);
+  releaseArguments(parameters);
+  synchronize();
+
+  return result;
+}
+
 // ================================================================================================
 
 amd::Memory* DmaBlitManager::pinHostMemory(const void* hostMem, size_t pinSize,

rocclr/device/rocm/rocblit.hpp

@@ -180,28 +180,32 @@ class DmaBlitManager : public device::HostBlitManager {
                          ) const;
 
   //! Stream memory write operation - Write a 'value' at 'memory'.
-  virtual bool streamOpsWrite(device::Memory& memory, //!< Memory to write the 'value'
-                             uint64_t value,
-                             size_t offset,
-                             size_t sizeBytes
-  ) const {
+  virtual bool streamOpsWrite(device::Memory& memory,  //!< Memory to write the 'value'
+                              uint64_t value,
+                              size_t offset,
+                              size_t sizeBytes) const {
     assert(!"Unimplemented");
     return false;
   }
 
   //! Stream memory ops- Waits for a 'value' at 'memory' and wait is released based on compare op.
-  virtual bool streamOpsWait(device::Memory& memory, //!< Memory contents to compare the 'value' against
+  virtual bool streamOpsWait(device::Memory& memory,  //!< Memory to compare the 'value' against
                              uint64_t value,
                              size_t offset,
                              size_t sizeBytes,
                              uint64_t flags,
-                             uint64_t mask
-  ) const {
+                             uint64_t mask) const {
     assert(!"Unimplemented");
     return false;
   }
 
-
+  virtual bool initHeap(device::Memory* heap_to_initialize,
+                        device::Memory* initial_blocks,
+                        uint heap_size,
+                        uint number_of_initial_blocks) const {
+    assert(!"Unimplemented");
+    return false;
+  }
 
  protected:
   static constexpr uint MaxPinnedBuffers = 4;
@@ -292,6 +296,7 @@ class KernelBlitManager : public DmaBlitManager {
     BlitCopyImage1DA,
     BlitCopyImageToBuffer,
     BlitCopyBufferToImage,
+    InitHeap,
     BlitTotal
   };
 
@@ -512,6 +517,12 @@ class KernelBlitManager : public DmaBlitManager {
 
   virtual amd::Monitor* lockXfer() const { return &lockXferOps_; }
 
+  virtual bool initHeap(device::Memory* heap_to_initialize,
+                        device::Memory* initial_blocks,
+                        uint heap_size,
+                        uint number_of_initial_blocks
+                        ) const;
+
  private:
   static constexpr size_t MaxXferBuffers = 2;
   static constexpr uint TransferSplitSize = 1;
@@ -586,7 +597,7 @@ static const char* BlitName[KernelBlitManager::BlitTotal] = {
   "__amd_rocclr_copyBufferRectAligned", "__amd_rocclr_streamOpsWrite", "__amd_rocclr_streamOpsWait",
   "__amd_rocclr_scheduler", "__amd_rocclr_gwsInit", "__amd_rocclr_fillImage",
   "__amd_rocclr_copyImage", "__amd_rocclr_copyImage1DA", "__amd_rocclr_copyImageToBuffer",
-  "__amd_rocclr_copyBufferToImage"
+  "__amd_rocclr_copyBufferToImage", "__amd_rocclr_initHeap"
 };
 
 inline void KernelBlitManager::setArgument(amd::Kernel* kernel, size_t index,

rocclr/device/rocm/rocblitcl.cpp

@@ -27,8 +27,9 @@ const char* rocBlitLinearSourceCode = BLIT_KERNEL(
   // Extern
   extern void __amd_streamOpsWrite(__global uint*, __global ulong*, ulong, ulong);
 
-  extern void __amd_streamOpsWait(__global uint*,__global ulong*, ulong, ulong, ulong);
+  extern void __amd_streamOpsWait(__global uint*, __global ulong*, ulong, ulong, ulong);
 
+  extern void __ockl_dm_init_v1(ulong, ulong, uint, uint);
   // Implementation
   __kernel void __amd_rocclr_streamOpsWrite(__global uint* ptrInt, __global ulong* ptrUlong,
                                             ulong value, ulong sizeBytes) {
@@ -39,6 +40,12 @@ const char* rocBlitLinearSourceCode = BLIT_KERNEL(
                                            ulong value, ulong flags, ulong mask) {
     __amd_streamOpsWait(ptrInt, ptrUlong, value, flags, mask);
   }
+
+  __kernel void __amd_rocclr_initHeap(ulong heap_to_initialize, ulong initial_blocks,
+                                      uint heap_size, uint number_of_initial_blocks) {
+    __ockl_dm_init_v1(heap_to_initialize, initial_blocks, heap_size, number_of_initial_blocks);
+  }
+
 );
 
 const char* SchedulerSourceCode = BLIT_KERNEL(

rocclr/device/rocm/rocdevice.cpp

@@ -3209,19 +3209,22 @@ bool Device::IsValidAllocation(const void* dev_ptr, size_t size) const {
 
 // ================================================================================================
 void Device::HiddenHeapAlloc() {
-  auto HeapAllocZeroOut = [this]()->bool {
+  auto HeapAllocZeroOut = [this]() -> bool {
     // Allocate initial heap for device memory allocator
     static constexpr size_t HeapBufferSize = 128 * Ki;
     heap_buffer_ = createMemory(HeapBufferSize);
-    // Clear memory to 0 for device library logic
-    if ((heap_buffer_ == nullptr) ||
-        (HSA_STATUS_SUCCESS != hsa_amd_memory_fill(
-      reinterpret_cast<void*>(HeapBuffer()->virtualAddress()), 0,
-      HeapBufferSize / sizeof(uint32_t)))) {
+    if (initial_heap_size_ != 0) {
+      initial_heap_size_ = amd::alignUp(initial_heap_size_, 2 * Mi);
+      initial_heap_buffer_ = createMemory(initial_heap_size_);
+    }
+    if (heap_buffer_ == nullptr) {
       LogError("Heap buffer allocation failed!");
       return false;
     }
-    return true;
+    bool result = static_cast<const KernelBlitManager&>(xferMgr()).initHeap(
+        heap_buffer_, initial_heap_buffer_, HeapBufferSize, initial_heap_size_ / (2 * Mi));
+
+    return result;
   };
   std::call_once(heap_initialized_, HeapAllocZeroOut);
 }

rocclr/utils/flags.hpp

@@ -280,6 +280,8 @@ release(bool, HIP_USE_RUNTIME_UNBUNDLER, true,                                \
         "Force this to use Runtime code object unbundler.")                   \
 release(bool, HIPRTC_USE_RUNTIME_UNBUNDLER, false,                            \
         "Set this to true to force runtime unbundler in hiprtc.")             \
+release(size_t, HIP_INITIAL_DM_SIZE, 8 * Mi,                                  \
+        "Set initial heap size for device malloc.")                           \
 
 namespace amd {