Fix launch when no work while keeping old behavior

LLNL · Feb 11, 2025 · ad29032 · ad29032
1 parent 01c41e0
commit ad29032
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 18 deletions.
diff --git a/include/RAJA/policy/cuda/kernel/CudaKernel.hpp b/include/RAJA/policy/cuda/kernel/CudaKernel.hpp
@@ -549,9 +549,13 @@ struct StatementExecutor<
 
 
     // Only launch kernel if we have something to iterate over
-    int num_blocks  = launch_dims.num_blocks();
-    int num_threads = launch_dims.num_threads();
-    if (num_blocks > 0 || num_threads > 0)
+    bool active_threads = launch_dims.threads_are_active();
+    bool active_blocks  = launch_dims.blocks_are_active();
+    int  num_blocks     = launch_dims.num_blocks();
+    int  num_threads    = launch_dims.num_threads();
+    if ((active_threads || active_blocks) &&
+        (!active_blocks || num_blocks > 0) &&
+        (!active_threads || num_threads > 0))
     {
 
       //

diff --git a/include/RAJA/policy/cuda/kernel/internal.hpp b/include/RAJA/policy/cuda/kernel/internal.hpp
@@ -98,20 +98,40 @@ struct LaunchDims
     return result;
   }
 
+  RAJA_INLINE
+  int blocks_are_active() const
+  {
+    return active.blocks.x || active.blocks.y || active.blocks.z;
+  }
+
+  RAJA_INLINE
+  int threads_are_active() const
+  {
+    return active.threads.x || active.threads.y || active.threads.z;
+  }
+
   RAJA_INLINE
   int num_blocks() const
   {
-    return (active.blocks.x ? dims.blocks.x : 1) *
-           (active.blocks.y ? dims.blocks.y : 1) *
-           (active.blocks.z ? dims.blocks.z : 1);
+    if (blocks_are_active()) {
+      return (active.blocks.x ? dims.blocks.x : 1) *
+             (active.blocks.y ? dims.blocks.y : 1) *
+             (active.blocks.z ? dims.blocks.z : 1);
+    } else {
+      return 0;
+    }
   }
 
   RAJA_INLINE
   int num_threads() const
   {
-    return (active.threads.x ? dims.threads.x : 1) *
-           (active.threads.y ? dims.threads.y : 1) *
-           (active.threads.z ? dims.threads.z : 1);
+    if (threads_are_active()) {
+      return (active.threads.x ? dims.threads.x : 1) *
+             (active.threads.y ? dims.threads.y : 1) *
+             (active.threads.z ? dims.threads.z : 1);
+    } else {
+      return 0;
+    }
   }
 
   RAJA_INLINE

diff --git a/include/RAJA/policy/hip/kernel/HipKernel.hpp b/include/RAJA/policy/hip/kernel/HipKernel.hpp
@@ -509,9 +509,13 @@ struct StatementExecutor<
 
 
     // Only launch kernel if we have something to iterate over
-    int num_blocks  = launch_dims.num_blocks();
-    int num_threads = launch_dims.num_threads();
-    if (num_blocks > 0 || num_threads > 0)
+    bool active_threads = launch_dims.threads_are_active();
+    bool active_blocks  = launch_dims.blocks_are_active();
+    int  num_blocks     = launch_dims.num_blocks();
+    int  num_threads    = launch_dims.num_threads();
+    if ((active_threads || active_blocks) &&
+        (!active_blocks || num_blocks > 0) &&
+        (!active_threads || num_threads > 0))
     {
 
       //

diff --git a/include/RAJA/policy/hip/kernel/internal.hpp b/include/RAJA/policy/hip/kernel/internal.hpp
@@ -98,20 +98,40 @@ struct LaunchDims
     return result;
   }
 
+  RAJA_INLINE
+  int blocks_are_active() const
+  {
+    return active.blocks.x || active.blocks.y || active.blocks.z;
+  }
+
+  RAJA_INLINE
+  int threads_are_active() const
+  {
+    return active.threads.x || active.threads.y || active.threads.z;
+  }
+
   RAJA_INLINE
   int num_blocks() const
   {
-    return (active.blocks.x ? dims.blocks.x : 1) *
-           (active.blocks.y ? dims.blocks.y : 1) *
-           (active.blocks.z ? dims.blocks.z : 1);
+    if (blocks_are_active()) {
+      return (active.blocks.x ? dims.blocks.x : 1) *
+             (active.blocks.y ? dims.blocks.y : 1) *
+             (active.blocks.z ? dims.blocks.z : 1);
+    } else {
+      return 0;
+    }
   }
 
   RAJA_INLINE
   int num_threads() const
   {
-    return (active.threads.x ? dims.threads.x : 1) *
-           (active.threads.y ? dims.threads.y : 1) *
-           (active.threads.z ? dims.threads.z : 1);
+    if (threads_are_active()) {
+      return (active.threads.x ? dims.threads.x : 1) *
+             (active.threads.y ? dims.threads.y : 1) *
+             (active.threads.z ? dims.threads.z : 1);
+    } else {
+      return 0;
+    }
   }
 
   RAJA_INLINE