From ce7b62d15c6e8926718e5e1b642a38188ba6922e Mon Sep 17 00:00:00 2001
From: Saleel Kudchadker <Saleel.Kudchadker@amd.com>
Date: Wed, 28 Feb 2024 23:08:47 +0000
Subject: [PATCH] SWDEV-443760 - Enable device kernel args for MI300

- Enable Device kernel args for MI300* for now.
- Fix a perf issue which impacts graph instantiate when dev kernel args
are enabled.

Change-Id: I962e58fd9d8dd1a8db95e601cb03a8e9c7bac97f


[ROCm/clr commit: 68f40f78dd4326ff8c9ecaf20316a5e103d298de]
---
 projects/clr/rocclr/device/rocm/rocsettings.cpp | 9 ++++++++-
 projects/clr/rocclr/device/rocm/rocvirtual.cpp  | 9 +++++----
 projects/clr/rocclr/utils/flags.hpp             | 2 +-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/projects/clr/rocclr/device/rocm/rocsettings.cpp b/projects/clr/rocclr/device/rocm/rocsettings.cpp
index 53729ece0a..3f479ad833 100644
--- a/projects/clr/rocclr/device/rocm/rocsettings.cpp
+++ b/projects/clr/rocclr/device/rocm/rocsettings.cpp
@@ -165,7 +165,9 @@ bool Settings::create(bool fullProfile, uint32_t gfxipMajor, uint32_t gfxipMinor
     barrier_value_packet_ = true;
   }
 
-  if (gfxipMajor >= 9) {
+  // Enable device kernel args for MI300* for now
+  if (gfxipMajor == 9 && gfxipMinor == 4 &&
+      (gfxStepping == 0 || gfxStepping == 1 || gfxStepping == 2)) {
     device_kernel_args_ = HIP_FORCE_DEV_KERNARG;
   }
 
@@ -178,6 +180,7 @@ bool Settings::create(bool fullProfile, uint32_t gfxipMajor, uint32_t gfxipMinor
        imageBufferWar_ = GPU_IMAGE_BUFFER_WAR;
      }
   }
+
   if (!flagIsDefault(GPU_ENABLE_WAVE32_MODE)) {
     enableWave32Mode_ = GPU_ENABLE_WAVE32_MODE;
   }
@@ -229,6 +232,10 @@ void Settings::override() {
   if (!flagIsDefault(ROC_USE_FGS_KERNARG)) {
     fgs_kernel_arg_ = ROC_USE_FGS_KERNARG;
   }
+
+  if (!flagIsDefault(HIP_FORCE_DEV_KERNARG)) {
+    device_kernel_args_ = HIP_FORCE_DEV_KERNARG;
+  }
 }
 }  // namespace roc
 
diff --git a/projects/clr/rocclr/device/rocm/rocvirtual.cpp b/projects/clr/rocclr/device/rocm/rocvirtual.cpp
index 1d2da0a562..52ef4de469 100644
--- a/projects/clr/rocclr/device/rocm/rocvirtual.cpp
+++ b/projects/clr/rocclr/device/rocm/rocvirtual.cpp
@@ -3210,12 +3210,13 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes,
     }
 
     const auto pcieKernargs = !dev().isXgmi() && dev().settings().device_kernel_args_;
-
     address argBuffer = hidden_arguments;
+    bool isGraphCapture = vcmd != nullptr && vcmd->getCapturingState();
+
     // Find all parameters for the current kernel
     if (!kernel.parameters().deviceKernelArgs() || gpuKernel.isInternalKernel()) {
       // Allocate buffer to hold kernel arguments
-      if(vcmd != nullptr && vcmd->getCapturingState()) {
+      if (isGraphCapture) {
         argBuffer = vcmd->getKernArgOffset();
       } else {
         const auto kernargSize = gpuKernel.KernargSegmentByteSize();
@@ -3226,7 +3227,7 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes,
       nontemporalMemcpy(argBuffer, parameters,
                         std::min(gpuKernel.KernargSegmentByteSize(),
                                  signature.paramsSize()));
-      if (pcieKernargs) {
+      if (pcieKernargs && !isGraphCapture) {
         *dev().info().hdpMemFlushCntl = 1u;
       }
     }
@@ -3289,7 +3290,7 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes,
                            (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_RELEASE_FENCE_SCOPE);
       aql_packet->setup = sizes.dimensions() << HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS;
     }
-    if (pcieKernargs) {
+    if (pcieKernargs && !isGraphCapture) {
       if (*dev().info().hdpMemFlushCntl != UINT32_MAX) {
         LogError("Unexpected HDP Register readback value!");
       }
diff --git a/projects/clr/rocclr/utils/flags.hpp b/projects/clr/rocclr/utils/flags.hpp
index f02a9e1224..ebdfc7ee85 100644
--- a/projects/clr/rocclr/utils/flags.hpp
+++ b/projects/clr/rocclr/utils/flags.hpp
@@ -233,7 +233,7 @@ 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.")                           \
-release(bool, HIP_FORCE_DEV_KERNARG, 0,                                       \
+release(bool, HIP_FORCE_DEV_KERNARG, true,                                    \
          "Force device mem for kernel args.")                                 \
 release(bool, DEBUG_CLR_GRAPH_PACKET_CAPTURE, true,                           \
          "Enable/Disable graph packet capturing")                             \