From 8cea0215fbeaf0661ec281b77782bb7cf0ec1426 Mon Sep 17 00:00:00 2001
From: German Andryeyev <German.Andryeyev@amd.com>
Date: Tue, 20 Apr 2021 11:55:36 -0400
Subject: [PATCH] SWDEV-282397 - Alloc scratch memory on the current queue

Device enqueue has an option to execute scheduler on the current
queue and it's enabled by default. Make sure scratch is allocated
on the current queue for that case. Add max vgpr tracking per
program to adjust scratch size accordingly.

Change-Id: I2a6d796913a4551a1e7f343a2465d589eec60d8a


[ROCm/clr commit: e553b2763aa09c74b8f2719dfdf3565b75ee8794]
---
 projects/clr/rocclr/device/pal/paldevice.cpp  | 13 +++++++++----
 projects/clr/rocclr/device/pal/palprogram.cpp |  4 ++++
 projects/clr/rocclr/device/pal/palprogram.hpp | 10 ++++++++--
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/projects/clr/rocclr/device/pal/paldevice.cpp b/projects/clr/rocclr/device/pal/paldevice.cpp
index 837109188f..c530a765f9 100644
--- a/projects/clr/rocclr/device/pal/paldevice.cpp
+++ b/projects/clr/rocclr/device/pal/paldevice.cpp
@@ -2043,19 +2043,24 @@ bool Device::validateKernel(const amd::Kernel& kernel, const device::VirtualDevi
     }
   }
 
-  if (devKernel->hsa()) {
-    const HSAILKernel* hsaKernel = static_cast<const HSAILKernel*>(devKernel);
-    if (hsaKernel->dynamicParallelism()) {
+  const HSAILKernel* hsaKernel = static_cast<const HSAILKernel*>(devKernel);
+  if (hsaKernel->dynamicParallelism()) {
+    if (settings().useDeviceQueue_) {
       amd::DeviceQueue* defQueue = kernel.program().context().defDeviceQueue(*this);
       if (defQueue != nullptr) {
         vgpu = static_cast<VirtualGPU*>(defQueue->vDev());
         if (!allocScratch(hsaKernel->prog().maxScratchRegs(), vgpu,
-                          devKernel->workGroupInfo()->usedVGPRs_)) {
+                          hsaKernel->prog().maxVgprs())) {
           return false;
         }
       } else {
         return false;
       }
+    } else {
+      if (!allocScratch(hsaKernel->prog().maxScratchRegs(), vgpu,
+                        hsaKernel->prog().maxVgprs())) {
+        return false;
+      }
     }
   }
 
diff --git a/projects/clr/rocclr/device/pal/palprogram.cpp b/projects/clr/rocclr/device/pal/palprogram.cpp
index 605642785c..0e720e2ccc 100644
--- a/projects/clr/rocclr/device/pal/palprogram.cpp
+++ b/projects/clr/rocclr/device/pal/palprogram.cpp
@@ -181,6 +181,7 @@ HSAILProgram::HSAILProgram(Device& device, amd::Program& owner)
       codeSegGpu_(nullptr),
       codeSegment_(nullptr),
       maxScratchRegs_(0),
+      maxVgprs_(0),
       executable_(nullptr),
       loaderContext_(this) {
   assert(device.isOnline());
@@ -194,6 +195,7 @@ HSAILProgram::HSAILProgram(NullDevice& device, amd::Program& owner)
       codeSegGpu_(nullptr),
       codeSegment_(nullptr),
       maxScratchRegs_(0),
+      maxVgprs_(0),
       executable_(nullptr),
       loaderContext_(this) {
   assert(!device.isOnline());
@@ -317,6 +319,7 @@ bool HSAILProgram::setKernels(amd::option::Options* options, void* binary, size_
       // with dynamic parallelism, since runtime doesn't know which child kernel will be called
       maxScratchRegs_ =
           std::max(static_cast<uint>(aKernel->workGroupInfo()->scratchRegs_), maxScratchRegs_);
+      maxVgprs_ = std::max(static_cast<uint>(aKernel->workGroupInfo()->usedVGPRs_), maxVgprs_);
     }
     // Allocate kernel table for device enqueuing
     if (!isNull() && dynamicParallelism && !allocKernelTable()) {
@@ -781,6 +784,7 @@ bool LightningProgram::setKernels(amd::option::Options* options, void* binary, s
     // with dynamic parallelism, since runtime doesn't know which child kernel will be called
     maxScratchRegs_ =
         std::max(static_cast<uint>(kernel->workGroupInfo()->scratchRegs_), maxScratchRegs_);
+    maxVgprs_ = std::max(static_cast<uint>(kernel->workGroupInfo()->usedVGPRs_), maxVgprs_);
   }
   DestroySegmentCpuAccess();
 #endif  // defined(USE_COMGR_LIBRARY)
diff --git a/projects/clr/rocclr/device/pal/palprogram.hpp b/projects/clr/rocclr/device/pal/palprogram.hpp
index f2e26297c3..74eb902f60 100644
--- a/projects/clr/rocclr/device/pal/palprogram.hpp
+++ b/projects/clr/rocclr/device/pal/palprogram.hpp
@@ -177,6 +177,9 @@ class HSAILProgram : public device::Program {
   //! Returns the maximum number of scratch regs used in the program
   uint maxScratchRegs() const { return maxScratchRegs_; }
 
+  //! Returns the maximum number of VGPR(s) used in the program
+  uint maxVgprs() const { return maxVgprs_; }
+
   //! Add internal static sampler
   void addSampler(Sampler* sampler) { staticSamplers_.push_back(sampler); }
 
@@ -237,8 +240,11 @@ class HSAILProgram : public device::Program {
   Memory* kernels_;                    //!< Table with kernel object pointers
   Memory* codeSegGpu_;                 //!< GPU memory with code objects
   Segment* codeSegment_;               //!< Pointer to the code segment for this program
-  uint
-      maxScratchRegs_;  //!< Maximum number of scratch regs used in the program by individual kernel
+  uint maxScratchRegs_;                //!< Maximum number of scratch regs used
+                                       //!< in the program by individual kernel
+  uint maxVgprs_;                      //!< Maximum number of VGPR(s) used
+                                       //!< in the program by individual kernel
+
   std::list<Sampler*> staticSamplers_;  //!< List od internal static samplers
 
   amd::hsa::loader::Loader* loader_;          //!< Loader object