From 1d785911ed35693afa5566eeac1fa7c21ff9ebd8 Mon Sep 17 00:00:00 2001
From: Jaydeep Patel <jaydeepkumar.patel@amd.com>
Date: Tue, 25 Jul 2023 18:45:03 +0000
Subject: [PATCH] SWDEV-408283 - Sync scratchRegs_, privateMemSize_ and
 workitemPrivateSegmentSize.

Change-Id: I623a7140810ff9867f8816bf4c8621a1fe921744


[ROCm/clr commit: ff1a999f66b721430796a3efcc56c929f7c95b03]
---
 projects/clr/rocclr/device/pal/palkernel.cpp  | 6 ++++++
 projects/clr/rocclr/device/pal/palvirtual.cpp | 7 -------
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/projects/clr/rocclr/device/pal/palkernel.cpp b/projects/clr/rocclr/device/pal/palkernel.cpp
index 84d4040fb1..e71c334718 100644
--- a/projects/clr/rocclr/device/pal/palkernel.cpp
+++ b/projects/clr/rocclr/device/pal/palkernel.cpp
@@ -552,6 +552,12 @@ bool LightningKernel::postLoad() {
   if (workGroupInfo_.size_ == 0) {
     return false;
   }
+  if ((workGroupInfo_.usedStackSize_ & 0x1) == 0x1) {
+    workGroupInfo_.scratchRegs_ =
+        std::max<uint32_t>(device().StackSize(), workGroupInfo_.scratchRegs_ * sizeof(uint32_t)) ;
+    workGroupInfo_.scratchRegs_ = amd::alignUp(workGroupInfo_.scratchRegs_, 16) / sizeof(uint32_t);
+    workGroupInfo_.privateMemSize_ = workGroupInfo_.scratchRegs_ * sizeof(uint32_t);
+  }
 
   // handle the printf metadata if any
   std::vector<std::string> printfStr;
diff --git a/projects/clr/rocclr/device/pal/palvirtual.cpp b/projects/clr/rocclr/device/pal/palvirtual.cpp
index e57de246b2..c7ae5f79e9 100644
--- a/projects/clr/rocclr/device/pal/palvirtual.cpp
+++ b/projects/clr/rocclr/device/pal/palvirtual.cpp
@@ -2678,13 +2678,6 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes, const
       dispatchParam.scratchSize = scratch->size_;
       dispatchParam.scratchOffset = scratch->offset_;
       dispatchParam.workitemPrivateSegmentSize = hsaKernel.spillSegSize();
-      if ((hsaKernel.workGroupInfo()->usedStackSize_ & 0x1) == 0x1) {
-        dispatchParam.workitemPrivateSegmentSize =
-            std::max<uint64_t>(dev().StackSize(), dispatchParam.workitemPrivateSegmentSize);
-        if (dispatchParam.workitemPrivateSegmentSize > 16 * Ki) {
-          dispatchParam.workitemPrivateSegmentSize = 16 * Ki;
-        }
-      }
     }
     dispatchParam.pCpuAqlCode = hsaKernel.cpuAqlCode();
     dispatchParam.hsaQueueVa = hsaQueueMem_->vmAddress();