From a808c4b23a54c84b89893833718ffb7e4e98e423 Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Fri, 13 Dec 2024 13:31:52 +0000
Subject: [PATCH] SWDEV-489255 - Update stack size limit in rocvirtual

Change-Id: I2aac9d211f64b3d6c121d8b010d215dcbdeac3aa
---
 rocclr/device/device.cpp          | 7 +------
 rocclr/device/device.hpp          | 6 ++++++
 rocclr/device/rocm/rocvirtual.cpp | 8 +++-----
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/rocclr/device/device.cpp b/rocclr/device/device.cpp
index fb45f77dee..ebfdc49782 100644
--- a/rocclr/device/device.cpp
+++ b/rocclr/device/device.cpp
@@ -938,12 +938,7 @@ bool Device::disableP2P(amd::Device* ptrDev) {
 }
 
 bool Device::UpdateStackSize(uint64_t stackSize) {
-  // Amount of space used by each wave is in units of 256 dwords.
-  // As per COMPUTE_TMPRING_SIZE.WAVE_SIZE 24:12
-  // The field size supports a range of 0->(2M-256) dwords per wave64.
-  // Per lane this works out to 131056 bytes or 128K - 16
-  uint64_t kStackSize = ((128 * Ki) - 16);
-  if (stackSize > kStackSize) {
+  if (stackSize > kMaxStackSize) {
     return false;
   }
   stack_size_ = amd::alignUp(stackSize, 16);
diff --git a/rocclr/device/device.hpp b/rocclr/device/device.hpp
index 28852ca0ca..92b38dcc1f 100644
--- a/rocclr/device/device.hpp
+++ b/rocclr/device/device.hpp
@@ -1653,6 +1653,12 @@ class Device : public RuntimeObject {
   static constexpr size_t kMGInfoSizePerDevice = kMGSyncDataSize + sizeof(MGSyncInfo);
   static constexpr size_t kSGInfoSize = kMGSyncDataSize;
 
+  // Amount of space used by each wave is in units of 256 dwords.
+  // As per COMPUTE_TMPRING_SIZE.WAVE_SIZE 24:12
+  // The field size supports a range of 0->(2M-256) dwords per wave64.
+  // Per lane this works out to 131056 bytes or 128K - 16
+  static constexpr size_t kMaxStackSize = ((128 * Ki) - 16);
+
   typedef std::list<CommandQueue*> CommandQueues;
 
   struct BlitProgram : public amd::HeapObject {
diff --git a/rocclr/device/rocm/rocvirtual.cpp b/rocclr/device/rocm/rocvirtual.cpp
index 3506799ca0..360e02722d 100644
--- a/rocclr/device/rocm/rocvirtual.cpp
+++ b/rocclr/device/rocm/rocvirtual.cpp
@@ -3479,11 +3479,9 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes,
     dispatchPacket.private_segment_size = devKernel->workGroupInfo()->privateMemSize_;
 
     if ((devKernel->workGroupInfo()->usedStackSize_ & 0x1) == 0x1) {
-      dispatchPacket.private_segment_size =
-              std::max<uint64_t>(dev().StackSize(), dispatchPacket.private_segment_size);
-      if (dispatchPacket.private_segment_size > 16 * Ki) {
-        dispatchPacket.private_segment_size = 16 * Ki;
-      }
+      dispatchPacket.private_segment_size = std::min<uint64_t>(
+          std::max<uint64_t>(dev().StackSize(), dispatchPacket.private_segment_size),
+          Device::kMaxStackSize);
     }
 
     // Pass the header accordingly