WIP: Faster Unstructured Backend

include/gridtools/fn/backend/common.hpp

@@ -31,10 +31,30 @@ namespace gridtools::fn::backend {
                 meta::rename<tuple, Dims>());
         }
 
+        template <class Dims, class Sizes, class UnrollFactors>
+        constexpr GT_FUNCTION auto make_unrolled_loops(Sizes const &sizes, UnrollFactors) {
+            return tuple_util::host_device::fold(
+                [&](auto outer, auto dim) {
+                    using unroll_factor = std::remove_reference_t<decltype(host_device::at_key<decltype(dim)>(
+                        std::declval<UnrollFactors const &>()))>;
+                    return [outer = std::move(outer),
+                               inner = sid::make_unrolled_loop<decltype(dim), unroll_factor::value>(
+                                   host_device::at_key<decltype(dim)>(sizes))](
+                               auto &&...args) { return outer(inner(std::forward<decltype(args)>(args)...)); };
+                },
+                host_device::identity(),
+                meta::rename<tuple, Dims>());
+        }
+
         template <class Sizes>
         constexpr GT_FUNCTION auto make_loops(Sizes const &sizes) {
             return make_loops<get_keys<Sizes>>(sizes);
         }
+
+        template <class Sizes, class UnrollFactors>
+        constexpr GT_FUNCTION auto make_unrolled_loops(Sizes const &sizes, UnrollFactors unroll_factors) {
+            return make_unrolled_loops<get_keys<Sizes>>(sizes, unroll_factors);
+        }
     } // namespace common
 
     template <class T>

include/gridtools/sid/loop.hpp

@@ -15,6 +15,7 @@
 #include <utility>
 
 #include "../common/defs.hpp"
+#include "../common/for_each.hpp"
 #include "../common/functional.hpp"
 #include "../common/host_device.hpp"
 #include "../common/integral_constant.hpp"
@@ -637,6 +638,30 @@ namespace gridtools {
             return {};
         }
 
+        template <class Key, int UnrollFactor, class NumSteps, class Step = integral_constant<int, 1>>
+        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());
+                }),
+                           epilogue = epilogue(std::forward<decltype(fun)>(fun)),
+                           epilogue_start](auto &&ptr, auto const &strides) {
+                    unrolled(std::forward<decltype(ptr)>(ptr), strides);
+                    shift(std::forward<decltype(ptr)>(ptr), get_stride<Key>(strides), epilogue_start);
+                    epilogue(std::forward<decltype(ptr)>(ptr), strides);
+                    shift(std::forward<decltype(ptr)>(ptr), get_stride<Key>(strides), -epilogue_start);
+                };
+            };
+        }
+
         /**
          *   A helper that allows to use `SID`s with C++11 range based loop
          *

jenkins/envs/daint_nvcc_cray.sh

@@ -7,6 +7,7 @@ module switch cudatoolkit/11.2.0_3.39-2.1__gf93aa1c
 module switch cce/10.0.2
 
 export GTCMAKE_GT_CLANG_CUDA_MODE=NVCC-CUDA
+export GTCMAKE_CMAKE_CUDA_FLAGS_RELEASE='-O3 -DNDEBUG --use_fast_math'
 
 export CTEST_PARALLEL_LEVEL=1
 

jenkins/envs/daint_nvcc_gcc.sh

@@ -17,5 +17,6 @@ export CUDAHOSTCXX="$CXX"
 
 export GTCMAKE_CMAKE_CXX_FLAGS='-march=haswell'
 export GTCMAKE_CMAKE_CXX_FLAGS_RELEASE='-Ofast -DNDEBUG'
+export GTCMAKE_CMAKE_CUDA_FLAGS_RELEASE='-O3 -DNDEBUG --use_fast_math'
 
 export CTEST_PARALLEL_LEVEL=1

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>>;
+    using fn_backend_t = gridtools::fn::backend::gpu<block_sizes_t<32, 8, 1>, block_sizes_t<1, 1, 1>>;
 } // namespace
 #endif
 
@@ -68,17 +68,19 @@ namespace gridtools::fn::backend {
             return "naive";
         }
     } // namespace naive_impl_
+    using naive_impl_::naive_with_threadpool;
 
     namespace gpu_impl_ {
-        template <class>
+        template <class, class>
         struct gpu;
-        template <class BlockSizes>
-        storage::gpu backend_storage_traits(gpu<BlockSizes>);
-        template <class BlockSizes>
-        timer_cuda backend_timer_impl(gpu<BlockSizes>);
-        template <class BlockSizes>
-        inline char const *backend_name(gpu<BlockSizes> const &) {
+        template <class ThreadBlockSizes, class LoopBlockSizes>
+        storage::gpu backend_storage_traits(gpu<ThreadBlockSizes, LoopBlockSizes>);
+        template <class ThreadBlockSizes, class LoopBlockSizes>
+        timer_cuda backend_timer_impl(gpu<ThreadBlockSizes, LoopBlockSizes>);
+        template <class ThreadBlockSizes, class LoopBlockSizes>
+        inline char const *backend_name(gpu<ThreadBlockSizes, LoopBlockSizes> const &) {
             return "gpu";
         }
     } // namespace gpu_impl_
+    using gpu_impl_::gpu;
 } // namespace gridtools::fn::backend

tests/regression/fn/fn_unstructured_nabla.cpp

@@ -184,6 +184,12 @@ namespace {
     static constexpr int vertex_field_id = 0;
     static constexpr int edge_field_id = 1;
 
+    using k_blocked_backend_t = meta::if_<meta::is_instantiation_of<backend::gpu, fn_backend_t>,
+        backend::gpu<meta::list<meta::list<integral_constant<int, 0>, integral_constant<int, 32>>,
+                         meta::list<integral_constant<int, 1>, integral_constant<int, 8>>>,
+            meta::list<meta::list<integral_constant<int, 1>, integral_constant<int, 5>>>>,
+        fn_backend_t>;
+
     constexpr inline auto make_comp = [](auto backend, auto const &mesh, auto &nabla) {
         using mesh_t = std::remove_reference_t<decltype(mesh)>;
         using float_t = typename mesh_t::float_t;
@@ -253,19 +259,21 @@ namespace {
         TypeParam::benchmark("fn_unstructured_nabla_field_of_tuples", comp);
     }
 
-    GT_REGRESSION_TEST(fn_unstructured_nabla_fused_field_of_tuples, test_environment<>, fn_backend_t) {
+    GT_REGRESSION_TEST(fn_unstructured_nabla_fused_field_of_tuples, test_environment<>, k_blocked_backend_t) {
         using float_t = typename TypeParam::float_t;
 
         auto mesh = TypeParam::fn_unstructured_mesh();
         auto nabla = mesh.template make_storage<tuple<float_t, float_t>>(mesh.nvertices(), mesh.nlevels());
-        auto comp = make_comp_fused(fn_backend_t(), mesh, nabla);
+        auto comp = make_comp_fused(k_blocked_backend_t(), mesh, nabla);
         comp();
         auto expected = make_expected(mesh);
         TypeParam::verify(expected, nabla);
         TypeParam::benchmark("fn_unstructured_nabla_fused_field_of_tuples", comp);
     }
 
     GT_REGRESSION_TEST(fn_unstructured_nabla_tuple_of_fields, test_environment<>, fn_backend_t) {
+        using float_t = typename TypeParam::float_t;
+
         auto mesh = TypeParam::fn_unstructured_mesh();
         auto nabla0 = mesh.template make_storage<float_t, vertex_field_id>(mesh.nvertices(), mesh.nlevels());
         auto nabla1 = mesh.template make_storage<float_t, vertex_field_id>(mesh.nvertices(), mesh.nlevels());
@@ -280,14 +288,16 @@ namespace {
         TypeParam::benchmark("fn_unstructured_nabla_tuple_of_fields", comp);
     }
 
-    GT_REGRESSION_TEST(fn_unstructured_nabla_fused_tuple_of_fields, test_environment<>, fn_backend_t) {
+    GT_REGRESSION_TEST(fn_unstructured_nabla_fused_tuple_of_fields, test_environment<>, k_blocked_backend_t) {
+        using float_t = typename TypeParam::float_t;
+
         auto mesh = TypeParam::fn_unstructured_mesh();
         auto nabla0 = mesh.template make_storage<float_t, vertex_field_id>(mesh.nvertices(), mesh.nlevels());
         auto nabla1 = mesh.template make_storage<float_t, vertex_field_id>(mesh.nvertices(), mesh.nlevels());
         auto nabla =
             sid::composite::keys<integral_constant<int, 0>, integral_constant<int, 1>>::make_values(nabla0, nabla1);
 
-        auto comp = make_comp_fused(fn_backend_t(), mesh, nabla);
+        auto comp = make_comp_fused(k_blocked_backend_t(), mesh, nabla);
         comp();
         auto expected = make_expected(mesh);
         TypeParam::verify([&](int vertex, int k) { return get<0>(expected(vertex, k)); }, nabla0);

tests/unit_tests/fn/test_fn_backend_gpu.cu

@@ -63,10 +63,16 @@ namespace gridtools::fn::backend {
 
             column_stage<int_t<1>, sum_scan, 0, 1> cs;
 
-            using block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>, meta::list<int_t<2>, int_t<2>>>;
+            using thread_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>, meta::list<int_t<2>, int_t<2>>>;
+            using loop_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<1>>, meta::list<int_t<2>, int_t<2>>>;
 
-            apply_column_stage(
-                gpu<block_sizes_t>(), sizes, cs, make_iterator_mock(), composite, int_t<1>(), tuple(42, 1));
+            apply_column_stage(gpu<thread_block_sizes_t, loop_block_sizes_t>(),
+                sizes,
+                cs,
+                make_iterator_mock(),
+                composite,
+                int_t<1>(),
+                tuple(42, 1));
 
             cudaMemcpy(outh, out.get(), 5 * 7 * 3 * sizeof(int), cudaMemcpyDeviceToHost);
             for (int i = 0; i < 5; ++i)
@@ -100,10 +106,16 @@ namespace gridtools::fn::backend {
 
             column_stage<int_t<0>, sum_scan, 0, 1> cs;
 
-            using block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>, meta::list<int_t<2>, int_t<2>>>;
+            using thread_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>, meta::list<int_t<2>, int_t<2>>>;
+            using loop_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<2>>, meta::list<int_t<2>, int_t<2>>>;
 
-            apply_column_stage(
-                gpu<block_sizes_t>(), sizes, cs, make_iterator_mock(), composite, int_t<0>(), tuple(42, 1));
+            apply_column_stage(gpu<thread_block_sizes_t, loop_block_sizes_t>(),
+                sizes,
+                cs,
+                make_iterator_mock(),
+                composite,
+                int_t<0>(),
+                tuple(42, 1));
 
             cudaMemcpy(outh, out.get(), 5 * sizeof(int), cudaMemcpyDeviceToHost);
             int res = 42;
@@ -139,13 +151,22 @@ namespace gridtools::fn::backend {
 
             column_stage<int_t<1>, sum_scan, 0, 1> cs;
 
-            using block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>,
+            using thread_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>,
                 meta::list<int_t<2>, int_t<2>>,
                 meta::list<int_t<3>, int_t<2>>,
                 meta::list<int_t<4>, int_t<1>>>;
+            using loop_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<2>>,
+                meta::list<int_t<2>, int_t<2>>,
+                meta::list<int_t<3>, int_t<2>>,
+                meta::list<int_t<4>, int_t<2>>>;
 
-            apply_column_stage(
-                gpu<block_sizes_t>(), sizes, cs, make_iterator_mock(), composite, int_t<1>(), tuple(42, 1));
+            apply_column_stage(gpu<thread_block_sizes_t, loop_block_sizes_t>(),
+                sizes,
+                cs,
+                make_iterator_mock(),
+                composite,
+                int_t<1>(),
+                tuple(42, 1));
 
             cudaMemcpy(outh, out.get(), 5 * 7 * 3 * 2 * 3 * sizeof(int), cudaMemcpyDeviceToHost);
             for (int i = 0; i < 5; ++i)
@@ -197,8 +218,8 @@ namespace gridtools::fn::backend {
         };
 
         TEST(backend_gpu, global_tmp) {
-            using block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>, meta::list<int_t<2>, int_t<2>>>;
-            auto alloc = tmp_allocator(gpu<block_sizes_t>());
+            using thread_block_sizes_t = meta::list<meta::list<int_t<0>, int_t<4>>, meta::list<int_t<2>, int_t<2>>>;
+            auto alloc = tmp_allocator(gpu<thread_block_sizes_t>());
             auto sizes = hymap::keys<int_t<0>, int_t<1>, int_t<2>>::values<int_t<5>, int_t<7>, int_t<3>>();
             auto tmp = allocate_global_tmp(alloc, sizes, data_type<int>());
             static_assert(sid::is_sid<decltype(tmp)>());

tests/unit_tests/sid/test_sid_loop.cpp

@@ -58,6 +58,35 @@ namespace gridtools {
                     EXPECT_EQ(88, data[i][j]) << " i:" << i << ", j:" << j;
         }
 
+        TEST(make_unrolled_loop, smoke) {
+            double data[10][10] = {};
+            auto strides = tuple(10_c, 1_c);
+
+            double *ptr = &data[0][0];
+            sid::make_unrolled_loop<i_t, 2>(5, 1)(assignment_f{42})(ptr, strides);
+            for (int i = 0; i < 5; ++i)
+                EXPECT_EQ(42, data[i][0]) << " i:" << i;
+
+            ptr = &data[2][3];
+            sid::make_unrolled_loop<j_t, 3>(4_c, -1_c)(assignment_f{5})(ptr, strides);
+            for (int i = 0; i < 4; ++i)
+                EXPECT_EQ(5, data[2][i]) << " i:" << i;
+
+            ptr = &data[0][0];
+            sid::make_unrolled_loop<i_t, 2>(10_c)(sid::make_unrolled_loop<j_t, 7>(10_c)(assignment_f{88}))(
+                ptr, strides);
+            for (int i = 0; i < 10; ++i)
+                for (int j = 0; j < 10; ++j)
+                    EXPECT_EQ(88, data[i][j]) << " i:" << i << ", j:" << j;
+
+            // pass ptr as r-value
+            sid::make_unrolled_loop<i_t, 3>(10_c)(sid::make_unrolled_loop<j_t, 1>(10_c)(assignment_f{88}))(
+                &data[0][0], strides);
+            for (int i = 0; i < 10; ++i)
+                for (int j = 0; j < 10; ++j)
+                    EXPECT_EQ(88, data[i][j]) << " i:" << i << ", j:" << j;
+        }
+
         TEST(nest_loops, smoke) {
             double data[10][10] = {};
             double *ptr = &data[0][0];