From be5f097e8ed08d2533468a8d84ac6a28498daeed Mon Sep 17 00:00:00 2001
From: Marko Arandjelovic <marandje@amd.com>
Date: Wed, 5 Jun 2024 14:27:01 +0200
Subject: [PATCH] SWDEV-465204 - Fix hipModuleLaunchKernel data validation

Change-Id: I129f265a5eb79d0a13da4f12e78e06ba307b17ee
---
 hipamd/src/hip_module.cpp | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/hipamd/src/hip_module.cpp b/hipamd/src/hip_module.cpp
index fab1f5f3c2..9d4a9e9512 100644
--- a/hipamd/src/hip_module.cpp
+++ b/hipamd/src/hip_module.cpp
@@ -215,8 +215,12 @@ hipError_t ihipLaunchKernel_validate(hipFunction_t f, uint32_t globalWorkSizeX,
                    "Both, kernelParams and extra Params are provided, only one should be provided");
     return hipErrorInvalidValue;
   }
-  if (globalWorkSizeX == 0 || globalWorkSizeY == 0 || globalWorkSizeZ == 0 || blockDimX == 0 ||
-      blockDimY == 0 || blockDimZ == 0) {
+
+  if (globalWorkSizeX == 0 || globalWorkSizeY == 0 || globalWorkSizeZ == 0) {
+    return hipErrorInvalidValue;
+  }
+
+  if (blockDimX == 0 || blockDimY == 0 || blockDimZ == 0) {
     return hipErrorInvalidConfiguration;
   }
 
@@ -372,7 +376,7 @@ hipError_t ihipModuleLaunchKernel(hipFunction_t f, uint32_t globalWorkSizeX,
   hip::DeviceFunc* function = hip::DeviceFunc::asFunction(f);
   amd::Kernel* kernel = function->kernel();
 
-  amd::ScopedLock lock (DEBUG_HIP_KERNARG_COPY_OPT ? &function->dflock_ : nullptr); 
+  amd::ScopedLock lock (DEBUG_HIP_KERNARG_COPY_OPT ? &function->dflock_ : nullptr);
 
   hipError_t status = ihipLaunchKernel_validate(
       f, globalWorkSizeX, globalWorkSizeY, globalWorkSizeZ, blockDimX, blockDimY, blockDimZ,
@@ -449,9 +453,10 @@ 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<int32_t>::max()/1024 ||
-      gridDimZ > std::numeric_limits<int32_t>::max()/1024) {
+
+  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()) + 1;
+  if (gridDimX > int32_max || gridDimY > uint16_max || gridDimZ > uint16_max) {
     HIP_RETURN(hipErrorInvalidValue);
   }
   size_t globalWorkSizeX = static_cast<size_t>(gridDimX) * blockDimX;