[flang-commits] [flang] 7d166fa - [flang][cuda] Correct the number of blocks when setting the grid to `*` (#121000)

[via flang-commits]

Author: khaki3
Date: 2024-12-23T17:14:38-08:00
New Revision: 7d166fa38470a23f3134a3793b9236b2a5c68fcf

URL: https://github.com/llvm/llvm-project/commit/7d166fa38470a23f3134a3793b9236b2a5c68fcf
DIFF: https://github.com/llvm/llvm-project/commit/7d166fa38470a23f3134a3793b9236b2a5c68fcf.diff

LOG: [flang][cuda] Correct the number of blocks when setting the grid to `*` (#121000)

We set the `gridX` argument of `_FortranACUFLaunchKernel` to `-1` when
`*` is passed to the grid parameter. We store it in one of `dim3`
members. However, `dim3` members are unsigned, so positive-value checks
we use later, such as `gridDim.x > 0`, are invalid. This PR utilizes the
original gird-size arguments to compute the number of blocks.

Added: 
    

Modified: 
    flang/runtime/CUDA/kernel.cpp

Removed: 
    


################################################################################
diff  --git a/flang/runtime/CUDA/kernel.cpp b/flang/runtime/CUDA/kernel.cpp
index 88cdf3cf426229..bdc04ccb17672b 100644
--- a/flang/runtime/CUDA/kernel.cpp
+++ b/flang/runtime/CUDA/kernel.cpp
@@ -48,13 +48,13 @@ void RTDEF(CUFLaunchKernel)(const void *kernel, intptr_t gridX, intptr_t gridY,
       maxBlocks = multiProcCount * maxBlocks;
     }
     if (maxBlocks > 0) {
-      if (gridDim.x > 0) {
+      if (gridX > 0) {
         maxBlocks = maxBlocks / gridDim.x;
       }
-      if (gridDim.y > 0) {
+      if (gridY > 0) {
         maxBlocks = maxBlocks / gridDim.y;
       }
-      if (gridDim.z > 0) {
+      if (gridZ > 0) {
         maxBlocks = maxBlocks / gridDim.z;
       }
       if (maxBlocks < 1) {
@@ -113,13 +113,13 @@ void RTDEF(CUFLaunchClusterKernel)(const void *kernel, intptr_t clusterX,
       maxBlocks = multiProcCount * maxBlocks;
     }
     if (maxBlocks > 0) {
-      if (config.gridDim.x > 0) {
+      if (gridX > 0) {
         maxBlocks = maxBlocks / config.gridDim.x;
       }
-      if (config.gridDim.y > 0) {
+      if (gridY > 0) {
         maxBlocks = maxBlocks / config.gridDim.y;
       }
-      if (config.gridDim.z > 0) {
+      if (gridZ > 0) {
         maxBlocks = maxBlocks / config.gridDim.z;
       }
       if (maxBlocks < 1) {