Add __launch_bounds__

GridTools · Jun 18, 2024 · 6e77d53 · 6e77d53
1 parent f70b885
commit 6e77d53
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/include/gridtools/fn/backend/gpu.hpp b/include/gridtools/fn/backend/gpu.hpp
@@ -117,8 +117,20 @@ namespace gridtools::fn::backend {
             return std::make_tuple(std::move(global_thread_indices), std::move(block_sizes));
         }
 
-        template <class ThreadBlockSizes, class LoopBlockSizes, class Sizes, class PtrHolder, class Strides, class Fun>
-        __global__ void kernel(Sizes sizes, PtrHolder ptr_holder, Strides strides, Fun fun) {
+        template <class Int, Int... i>
+        constexpr int iseq_product(std::integer_sequence<Int, i...>) {
+            return (1 * ... * i);
+        }
+
+        template <class ThreadBlockSizes,
+            class LoopBlockSizes,
+            class Sizes,
+            class PtrHolder,
+            class Strides,
+            class Fun,
+            int NumThreads = iseq_product(meta::list_to_iseq<block_sizes_for_sizes<ThreadBlockSizes, Sizes>>())>
+        __global__ void __launch_bounds__(NumThreads)
+            kernel(Sizes sizes, PtrHolder ptr_holder, Strides strides, Fun fun) {
             auto const [thread_idx, block_size] = global_thread_index<ThreadBlockSizes, LoopBlockSizes>(sizes);
             if (!tuple_util::all_of(std::less(), thread_idx, sizes))
                 return;

diff --git a/tests/include/fn_select.hpp b/tests/include/fn_select.hpp
@@ -47,7 +47,7 @@ namespace {
                                  gridtools::integral_constant>,
             gridtools::meta::list<gridtools::integral_constant<int, sizes>...>>;
 
-    using fn_backend_t = gridtools::fn::backend::gpu<block_sizes_t<32, 8, 1>, block_sizes_t<2, 2, 2>>;
+    using fn_backend_t = gridtools::fn::backend::gpu<block_sizes_t<32, 8, 1>, block_sizes_t<1, 1, 1>>;
 } // namespace
 #endif