From 505bce23bea92d6ec5512303a67b5020689a3d57 Mon Sep 17 00:00:00 2001
From: pghafari <pghafari@amd.com>
Date: Mon, 24 Apr 2023 21:38:53 -0400
Subject: [PATCH] SWDEV-369567 - dimension check update

Change-Id: I31deb5070d4852db7dd6d37f9a514965f387698e


[ROCm/clr commit: 31b362bf6ed8e0843f565c7820e2bf9e14fe8e88]
---
 projects/clr/hipamd/src/hip_device.cpp |  5 +++--
 projects/clr/hipamd/src/hip_module.cpp | 12 +++++++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_device.cpp b/projects/clr/hipamd/src/hip_device.cpp
index 55a8aa9a18..d545a5ba47 100644
--- a/projects/clr/hipamd/src/hip_device.cpp
+++ b/projects/clr/hipamd/src/hip_device.cpp
@@ -284,6 +284,7 @@ hipError_t ihipGetDeviceProperties(hipDeviceProp_t* props, hipDevice_t device) {
   auto* deviceHandle = g_devices[device]->devices()[0];
 
   constexpr auto int32_max = static_cast<uint64_t>(std::numeric_limits<int32_t>::max());
+  constexpr auto uint16_max = static_cast<uint64_t>(std::numeric_limits<uint16_t>::max());
   hipDeviceProp_t deviceProps = {0};
 
   const auto& info = deviceHandle->info();
@@ -298,8 +299,8 @@ hipError_t ihipGetDeviceProperties(hipDeviceProp_t* props, hipDevice_t device) {
   deviceProps.maxThreadsDim[1] = info.maxWorkItemSizes_[1];
   deviceProps.maxThreadsDim[2] = info.maxWorkItemSizes_[2];
   deviceProps.maxGridSize[0] = int32_max;
-  deviceProps.maxGridSize[1] = int32_max;
-  deviceProps.maxGridSize[2] = int32_max;
+  deviceProps.maxGridSize[1] = uint16_max;
+  deviceProps.maxGridSize[2] = uint16_max;
   deviceProps.clockRate = info.maxEngineClockFrequency_ * 1000;
   deviceProps.memoryClockRate = info.maxMemoryClockFrequency_ * 1000;
   deviceProps.memoryBusWidth = info.globalMemChannels_;
diff --git a/projects/clr/hipamd/src/hip_module.cpp b/projects/clr/hipamd/src/hip_module.cpp
index 1675baa9f2..b2630af42b 100644
--- a/projects/clr/hipamd/src/hip_module.cpp
+++ b/projects/clr/hipamd/src/hip_module.cpp
@@ -229,7 +229,7 @@ hipError_t ihipLaunchKernel_validate(hipFunction_t f, uint32_t globalWorkSizeX,
     return hipErrorInvalidValue;
   }
   // Make sure dispatch doesn't exceed max workgroup size limit
-  if (blockDimX * blockDimY * blockDimZ > device->info().maxWorkGroupSize_) {
+  if (blockDimX * blockDimY * blockDimZ > info.maxWorkGroupSize_) {
     return hipErrorInvalidValue;
   }
   hip::DeviceFunc* function = hip::DeviceFunc::asFunction(f);
@@ -439,13 +439,15 @@ hipError_t hipModuleLaunchKernel(hipFunction_t f, uint32_t gridDimX, uint32_t gr
 
   STREAM_CAPTURE(hipModuleLaunchKernel, hStream, f, gridDimX, gridDimY, gridDimZ, blockDimX,
                  blockDimY, blockDimZ, sharedMemBytes, kernelParams, extra);
-
+  if (gridDimX > std::numeric_limits<int32_t>::max() ||
+      gridDimY > std::numeric_limits<uint16_t>::max() ||
+      gridDimZ > std::numeric_limits<uint16_t>::max()) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
   size_t globalWorkSizeX = static_cast<size_t>(gridDimX) * blockDimX;
   size_t globalWorkSizeY = static_cast<size_t>(gridDimY) * blockDimY;
   size_t globalWorkSizeZ = static_cast<size_t>(gridDimZ) * blockDimZ;
-  if (globalWorkSizeX > std::numeric_limits<uint32_t>::max() ||
-      globalWorkSizeY > std::numeric_limits<uint32_t>::max() ||
-      globalWorkSizeZ > std::numeric_limits<uint32_t>::max()) {
+  if (globalWorkSizeX > std::numeric_limits<uint32_t>::max()) {
     HIP_RETURN(hipErrorInvalidConfiguration);
   }
   HIP_RETURN(ihipModuleLaunchKernel(