From 342e478e7dac1cfb3d7681809ec132995e6a85e6 Mon Sep 17 00:00:00 2001
From: David Yat Sin <David.YatSin@amd.com>
Date: Wed, 14 May 2025 18:07:46 +0000
Subject: [PATCH] rocr: Perform memcpy for small code-object loads

On large BAR systems, for small-sized code-objects, we get performance
using direct memcpy due to latencies when doing the blit-copy.


[ROCm/ROCR-Runtime commit: da2607024ba9fba6d034f5f1490d665847228f46]
---
 .../hsa-runtime/core/inc/amd_gpu_agent.h      | 10 ++++++++
 .../core/runtime/amd_loader_context.cpp       | 24 ++++++++++++-------
 .../runtime/hsa-runtime/core/util/flag.h      |  7 ++++++
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_gpu_agent.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
index f514dda78c..87bad25786 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
@@ -440,6 +440,16 @@ class GpuAgent : public GpuAgentInt {
   /// @brief Is large BAR support enabled for this GPU.
   __forceinline bool LargeBarEnabled() const { return large_bar_enabled_; }
 
+  /// @brief Force a WC flush on PCIe devices by doing a write and then read-back
+  __forceinline void PcieWcFlush(void *ptr, size_t size) const {
+    if (!xgmi_cpu_gpu_) {
+      _mm_sfence();
+      *((uint8_t*)ptr + size - 1) = *((uint8_t*)ptr + size - 1);
+      _mm_mfence();
+      auto readback = *reinterpret_cast<volatile uint8_t*>(ptr + size - 1);
+    }
+  }
+
   const size_t MAX_SCRATCH_APERTURE_PER_XCC = (1ULL << 32);
   size_t MaxScratchDevice() const { return properties_.NumXcc * MAX_SCRATCH_APERTURE_PER_XCC; }
 
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp
index 94f21eb9af..d4c950dff9 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp
@@ -350,17 +350,25 @@ bool RegionMemory::Freeze() {
   assert(this->Allocated() && nullptr != host_ptr_);
 
   core::Agent* agent = region_->owner();
-  if (agent != NULL && agent->device_type() == core::Agent::kAmdGpuDevice) {
-    if (HSA_STATUS_SUCCESS != agent->DmaCopy(ptr_, host_ptr_, size_)) {
-      return false;
-    }
+
+  const size_t& code_object_dmacopy_size =
+    core::Runtime::runtime_singleton_->flag().co_dmacopy_size();
+
+  const bool isGpuDevice = (agent->device_type() == core::Agent::kAmdGpuDevice);
+  const bool isLargeBarDisabled = isGpuDevice && !reinterpret_cast<AMD::GpuAgent*>(agent)->LargeBarEnabled();
+  const bool shouldDmaCopy = isGpuDevice && (isLargeBarDisabled || size_ > code_object_dmacopy_size);
+
+  if (shouldDmaCopy) {
+      if (HSA_STATUS_SUCCESS != agent->DmaCopy(ptr_, host_ptr_, size_)) return false;
   } else {
-    memcpy(ptr_, host_ptr_, size_);
+      memcpy(ptr_, host_ptr_, size_);
+      if (is_code_ && isGpuDevice)
+        reinterpret_cast<AMD::GpuAgent*>(agent)->PcieWcFlush(ptr_, size_);
   }
 
-  // Invalidate agent caches which may hold lines of the new allocation.
-  if (is_code_ && (region_->owner()->device_type() == core::Agent::kAmdGpuDevice))
-    ((AMD::GpuAgent*)region_->owner())->InvalidateCodeCaches(ptr_, size_);
+  // Invalidate agent caches if needed
+  if (is_code_ && isGpuDevice)
+      reinterpret_cast<AMD::GpuAgent*>(agent)->InvalidateCodeCaches(ptr_, size_);
 
   return true;
 }
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h
index 1495352068..40c5ee1007 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h
@@ -290,6 +290,9 @@ class Flag {
     // HSA_DTIF_ENABLED = 0 will disable DTIF backend.
     var = os::GetEnvVar("HSA_ENABLE_DTIF");
     enable_dtif_ = (var == "1") ? true : false;
+
+    var = os::GetEnvVar("HSA_CO_DMACOPY_SIZE");
+    co_dmacopy_size_ = var.empty() ? 1024*1024 : atoi(var.c_str());
   }
 
   void parse_masks(uint32_t maxGpu, uint32_t maxCU) {
@@ -402,6 +405,8 @@ class Flag {
 
   size_t pc_sampling_max_device_buffer_size() const { return pc_sampling_max_device_buffer_size_; }
 
+  size_t co_dmacopy_size() const { return co_dmacopy_size_; }
+
   bool dev_mem_queue_buf() const { return dev_mem_queue_buf_; }
 
   uint32_t signal_abort_timeout() const { return signal_abort_timeout_; }
@@ -475,6 +480,8 @@ class Flag {
 
   size_t pc_sampling_max_device_buffer_size_;
 
+  size_t co_dmacopy_size_;
+
   // Map GPU index post RVD to its default cu mask.
   std::map<uint32_t, std::vector<uint32_t>> cu_mask_;