Use sid::make_loop in sid::make_unrolled_loop if unroll factor is 1

GridTools · Sep 25, 2024 · ffd070e · ffd070e
1 parent ef4fed7
commit ffd070e
Showing 1 changed file with 22 additions and 8 deletions.
diff --git a/include/gridtools/sid/loop.hpp b/include/gridtools/sid/loop.hpp
@@ -638,20 +638,25 @@ namespace gridtools {
             return {};
         }
 
-        template <class Key, int UnrollFactor, class NumSteps, class Step = integral_constant<int, 1>>
+        template <class Key,
+            int UnrollFactor,
+            class NumSteps,
+            class Step = integral_constant<int, 1>,
+            std::enable_if_t<(UnrollFactor > 1), int> = 0>
         constexpr GT_FUNCTION auto make_unrolled_loop(NumSteps num_steps, Step step = {}) {
             using u = integral_constant<int, UnrollFactor>;
             return [step,
                        unrolled = make_loop<Key>(num_steps / u(), step * u()),
                        epilogue = make_loop<Key>(num_steps % u(), step),
                        epilogue_start = step * ((num_steps / u()) * u())](auto &&fun) {
-                return [unrolled = unrolled([step, fun=std::forward<decltype(fun)>(fun)](auto &&ptr, auto const strides) {
-                    ::gridtools::host_device::for_each<meta::make_indices_c<UnrollFactor>>([&](auto) {
-                        fun(std::forward<decltype(ptr)>(ptr), strides);
-                        shift(std::forward<decltype(ptr)>(ptr), get_stride<Key>(strides), step);
-                    });
-                    shift(std::forward<decltype(ptr)>(ptr), get_stride<Key>(strides), -step * u());
-                }),
+                return [unrolled =
+                               unrolled([step, fun = std::forward<decltype(fun)>(fun)](auto &&ptr, auto const strides) {
+                                   ::gridtools::host_device::for_each<meta::make_indices_c<UnrollFactor>>([&](auto) {
+                                       fun(std::forward<decltype(ptr)>(ptr), strides);
+                                       shift(std::forward<decltype(ptr)>(ptr), get_stride<Key>(strides), step);
+                                   });
+                                   shift(std::forward<decltype(ptr)>(ptr), get_stride<Key>(strides), -step * u());
+                               }),
                            epilogue = epilogue(std::forward<decltype(fun)>(fun)),
                            epilogue_start](auto &&ptr, auto const &strides) {
                     unrolled(std::forward<decltype(ptr)>(ptr), strides);
@@ -662,6 +667,15 @@ namespace gridtools {
             };
         }
 
+        template <class Key,
+            int UnrollFactor,
+            class NumSteps,
+            class Step = integral_constant<int, 1>,
+            std::enable_if_t<(UnrollFactor == 1), int> = 0>
+        constexpr GT_FUNCTION auto make_unrolled_loop(NumSteps num_steps, Step step = {}) {
+            return make_loop<Key>(num_steps, step);
+        }
+
         /**
          *   A helper that allows to use `SID`s with C++11 range based loop
          *