Launch_nd: Abstraction to switch between flat and nested loops

docs/sphinx/user_guide/cook_book.rst

@@ -22,4 +22,4 @@ to provide users with complete beyond usage examples beyond what can be found in
 
    cook_book/reduction
    cook_book/multi-reduction
-
+   cook_book/launch-nd

docs/sphinx/user_guide/cook_book/launch-nd.rst

@@ -0,0 +1,116 @@
+.. ##
+.. ## Copyright (c) Lawrence Livermore National Security, LLC and other
+.. ## RAJA Project Developers. See top-level LICENSE and COPYRIGHT
+.. ## files for dates and other details. No copyright assignment is required
+.. ## to contribute to RAJA.
+.. ##
+.. ## SPDX-License-Identifier: (BSD-3-Clause)
+.. ##
+
+.. _cook-book-launch-nd-label:
+
+============================
+Cooking with RAJA::launch_nd
+============================
+
+``RAJA::launch_nd`` runs a logical 2-D or 3-D loop body through selectable
+launch-backed mappings. It is intended for kernels where the source code should
+stay written in logical multi-dimensional indices, but the best GPU mapping is
+not known from the source alone.
+
+The interface supports two mapping policy families:
+
+ * ``RAJA::launch_nd_flattened_policy<ExecPolicy, LayoutTag>`` maps the product
+   of the logical dimensions to a 1-D launch. The ``ExecPolicy`` is a regular
+   forall-style policy such as ``RAJA::cuda_exec<256>`` or
+   ``RAJA::hip_exec<256>``. RAJA derives the launch parameters and performs the
+   linear-to-logical index reconstruction internally.
+ * ``RAJA::launch_nd_grid_policy<LaunchPolicy, LoopPolicies...>`` maps the
+   logical dimensions directly to a true 2-D or 3-D launch. The user supplies
+   the launch policy, one loop policy per segment, and the ``RAJA::LaunchParams``.
+
+This makes it possible to put both mappings behind one abstraction and choose
+the mapping policy from problem size, backend, or measurements while preserving
+one logical loop body.
+
+----------------------------------
+Why Compare Flat and Grid Mappings
+----------------------------------
+
+Many application kernels are logically 2-D or 3-D but have historically run on
+a 1-D iteration space, with division and modulo operations used to recover the
+logical indices. That approach can expose more parallelism when one logical
+dimension is small.
+
+For example, a ``cells x components`` kernel with only a few components may not
+fit a fixed ``16 x 16`` block well. A true 2-D grid can leave many component
+threads inactive in each block. The flattened mapping instead launches over
+``cells * components`` as a 1-D space, which can produce fuller blocks. The
+tradeoff is the extra index reconstruction arithmetic. The true grid mapping
+can still win when the dimensions fit the block shape, when the body is very
+small and index reconstruction dominates, or when direct multi-dimensional
+mapping improves memory access or scheduling.
+
+----------------
+Mapping Policies
+----------------
+
+The example source defines backend-specific policy aliases. CUDA uses direct
+global loop policies for the true grid mapping:
+
+.. literalinclude:: ../../../../examples/launch_nd.cpp
+   :start-after: // _launch_nd_policy_aliases_start
+   :end-before: // _launch_nd_policy_aliases_end
+   :language: C++
+
+The flattened policy uses a forall-style execution policy such as
+``RAJA::cuda_exec<block_size_1d>`` or ``RAJA::hip_exec<block_size_1d>``. The
+grid policy uses ``RAJA::LaunchPolicy`` and direct global loop policies such as
+``RAJA::cuda_global_y_direct`` and ``RAJA::cuda_global_x_direct``.
+
+Both mappings call the same logical body through one ``RAJA::launch_nd`` call:
+
+.. literalinclude:: ../../../../examples/launch_nd.cpp
+   :start-after: // _launch_nd_call_start
+   :end-before: // _launch_nd_call_end
+   :language: C++
+
+----------------------
+Runtime Policy Choice
+----------------------
+
+The mapping can be selected at run time by choosing which policy object is
+passed to the common implementation:
+
+.. literalinclude:: ../../../../examples/launch_nd.cpp
+   :start-after: // _launch_nd_runtime_select_start
+   :end-before: // _launch_nd_runtime_select_end
+   :language: C++
+
+Run the example both ways and compare timing with the profiling tool normally
+used for the target backend::
+
+  ./launch_nd flat
+  ./launch_nd grid
+
+The example uses ``num_cells = 257`` and ``num_comp = 5`` to show a case where
+the ``16 x 16`` grid mapping has a small logical component dimension. Change
+``num_cells``, ``num_comp``, ``block_size_1d``, ``block_x``, and ``block_y`` in
+``RAJA/examples/launch_nd.cpp`` to explore when the flattened or true-grid
+mapping is better for a kernel shape.
+
+--------------------
+Current Capabilities
+--------------------
+
+``RAJA::launch_nd`` currently supports ``RAJA::TypedRangeSegment`` packs created
+with ``RAJA::segments``. The grid mapping supports 2-D and 3-D loops and
+requires one loop policy per segment. The flattened mapping uses
+``RAJA::layout_right`` by default, or ``RAJA::layout_left`` when that layout tag
+is supplied, to control which logical index is unit stride in the flattened
+space.
+
+Use direct ``RAJA::launch`` when the kernel needs explicit shared memory, team
+synchronization, multiple cooperating loops in a single launch body, or other
+hierarchical launch features that do not fit the single logical body accepted by
+``RAJA::launch_nd``.

examples/CMakeLists.txt

@@ -61,6 +61,10 @@ raja_add_executable(
   NAME launch_flatten
   SOURCES launch_flatten.cpp)
 
+raja_add_executable(
+  NAME launch_nd
+  SOURCES launch_nd.cpp)
+
 raja_add_executable(
   NAME launch_reductions
   SOURCES launch_reductions.cpp)

examples/launch_nd.cpp

@@ -0,0 +1,175 @@
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// Copyright (c) Lawrence Livermore National Security, LLC and other
+// RAJA Project Developers. See top-level LICENSE and COPYRIGHT
+// files for dates and other details. No copyright assignment is required
+// to contribute to RAJA.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include "RAJA/RAJA.hpp"
+
+/*
+ *  RAJA::launch_nd mapping policies
+ *
+ *  This example shows how to select a launch_nd mapping at run time while
+ *  keeping one logical 2-D loop body. Run with:
+ *
+ *    ./launch_nd flat
+ *    ./launch_nd grid
+ *
+ *  The flattened policy takes a regular RAJA forall execution policy, such as
+ *  cuda_exec<256> or hip_exec<256>, and RAJA maps it through launch internally.
+ *  The grid policy takes an explicit launch policy, loop policies, and
+ *  LaunchParams for a true 2-D launch mapping.
+ */
+
+namespace
+{
+
+constexpr int block_size_1d = 256;
+constexpr int block_x       = 16;
+constexpr int block_y       = 16;
+constexpr int num_cells     = 257;
+constexpr int num_comp      = 5;
+constexpr int total         = num_cells * num_comp;
+
+int ceil_div(int value, int divisor) { return (value + divisor - 1) / divisor; }
+
+// _launch_nd_mapping_enum_start
+enum class Mapping
+{
+  Flat,
+  Grid
+};
+// _launch_nd_mapping_enum_end
+
+// _launch_nd_policy_aliases_start
+#if defined(RAJA_ENABLE_CUDA)
+using launch_policy = RAJA::LaunchPolicy<RAJA::cuda_launch_t<false>>;
+using flat_exec     = RAJA::cuda_exec<block_size_1d>;
+using cell_loop     = RAJA::LoopPolicy<RAJA::cuda_global_y_direct>;
+using comp_loop     = RAJA::LoopPolicy<RAJA::cuda_global_x_direct>;
+using resource_type = RAJA::resources::Cuda;
+
+#elif defined(RAJA_ENABLE_HIP)
+using launch_policy = RAJA::LaunchPolicy<RAJA::hip_launch_t<false>>;
+using flat_exec     = RAJA::hip_exec<block_size_1d>;
+using cell_loop     = RAJA::LoopPolicy<RAJA::hip_global_y_direct>;
+using comp_loop     = RAJA::LoopPolicy<RAJA::hip_global_x_direct>;
+using resource_type = RAJA::resources::Hip;
+
+#else
+using launch_policy = RAJA::LaunchPolicy<RAJA::seq_launch_t>;
+using flat_exec     = RAJA::seq_exec;
+using cell_loop     = RAJA::LoopPolicy<RAJA::seq_exec>;
+using comp_loop     = RAJA::LoopPolicy<RAJA::seq_exec>;
+using resource_type = RAJA::resources::Host;
+#endif
+// _launch_nd_policy_aliases_end
+
+Mapping get_mapping(int argc, char** argv)
+{
+  if (argc < 2)
+  {
+    return Mapping::Flat;
+  }
+
+  const std::string arg = argv[1];
+  if (arg == "flat" || arg == "flattened")
+  {
+    return Mapping::Flat;
+  }
+  if (arg == "grid")
+  {
+    return Mapping::Grid;
+  }
+
+  std::cout << "Unknown mapping '" << arg << "'. Use 'flat' or 'grid'.\n";
+  return Mapping::Flat;
+}
+
+template<typename LaunchNdPolicy>
+int run_mapping(RAJA::resources::Resource res,
+                LaunchNdPolicy policy,
+                char const* mapping_name)
+{
+  int* values_ptr = res.allocate<int>(total);
+  res.memset(values_ptr, 0, sizeof(int) * total);
+
+  auto cells = RAJA::TypedRangeSegment<int>(0, num_cells);
+  auto comps = RAJA::TypedRangeSegment<int>(0, num_comp);
+
+  // _launch_nd_call_start
+  RAJA::launch_nd(res, policy, RAJA::segments(cells, comps),
+                  [=] RAJA_HOST_DEVICE(int cell, int comp) {
+                    const int idx   = comp + num_comp * cell;
+                    values_ptr[idx] = 1000 * cell + comp;
+                  });
+  // _launch_nd_call_end
+
+  std::vector<int> values(total);
+  res.wait();
+  res.memcpy(values.data(), values_ptr, sizeof(int) * total);
+  res.wait();
+
+  int errors = 0;
+  for (int cell = 0; cell < num_cells; ++cell)
+  {
+    for (int comp = 0; comp < num_comp; ++comp)
+    {
+      const int idx      = comp + num_comp * cell;
+      const int expected = 1000 * cell + comp;
+      if (values[idx] != expected)
+      {
+        ++errors;
+      }
+    }
+  }
+
+  std::cout << "  mapping: " << mapping_name << '\n';
+  std::cout << "  logical size: " << num_cells << " x " << num_comp << '\n';
+  std::cout << "  result -- " << (errors == 0 ? "PASS" : "FAIL") << '\n';
+
+  res.deallocate(values_ptr);
+  return errors;
+}
+
+}  // namespace
+
+int main(int argc, char** argv)
+{
+  std::cout << "\nRAJA launch_nd mapping example...\n";
+
+  RAJA::resources::Resource res(resource_type {});
+  const Mapping mapping = get_mapping(argc, argv);
+
+  int errors = 0;
+  // _launch_nd_runtime_select_start
+  if (mapping == Mapping::Flat)
+  {
+    auto policy = RAJA::launch_nd_flattened_policy<flat_exec> {};
+    errors      = run_mapping(res, policy, "flat");
+    std::cout << "  flattened launch threads per block: " << block_size_1d
+              << '\n';
+  }
+  else
+  {
+    auto policy =
+        RAJA::launch_nd_grid_policy<launch_policy, cell_loop, comp_loop>(
+            RAJA::LaunchParams(RAJA::Teams(ceil_div(num_comp, block_x),
+                                           ceil_div(num_cells, block_y)),
+                               RAJA::Threads(block_x, block_y)));
+    errors = run_mapping(res, policy, "grid");
+    std::cout << "  grid launch block shape: " << block_x << " x " << block_y
+              << '\n';
+  }
+  // _launch_nd_runtime_select_end
+
+  std::cout << "\nDONE!...\n";
+  return errors == 0 ? 0 : 1;
+}

include/RAJA/RAJA.hpp

@@ -124,6 +124,7 @@
 #include "RAJA/util/StaticLayout.hpp"
 #include "RAJA/util/IndexLayout.hpp"
 #include "RAJA/util/View.hpp"
+#include "RAJA/pattern/launch_nd.hpp"
 
 
 //