Add kernel fusing using RAJA

CMakeLists.txt

@@ -59,8 +59,6 @@ option(ENABLE_SAMRAI_TESTS "Enable SAMRAI Test Programs" On)
 option(ENABLE_PERF_TESTS "Enable Performance Tests." Off)
 set(NUM_PERF_PROCS 8 CACHE INT "Number of processors for performance tests.")
 option(ENABLE_CHECK_ASSERTIONS "Enable assertion checking." On)
-option(ENABLE_CHECK_DEV_ASSERTIONS "Enable SAMRAI developer assertion checking." Off)
-option(ENABLE_CHECK_DIM_ASSERTIONS "Enable assertion checking for dimensions." Off)
 option(ENABLE_BOX_COUNTING "Turns on box telemetry." Off)
 option(ENABLE_DEPRECATED "Build with deprecated features." On)
 option(ENABLE_TIMERS "Enable SAMRAI timers." On)
@@ -72,6 +70,7 @@ set(CUDA_ARCH "sm_70" CACHE STRING "Compute architecture to pass to CUDA builds"
 set(CMAKE_CUDA_FLAGS "" CACHE STRING "")
 set(CMAKE_INSTALL_LIBDIR lib)
 #set(CMAKE_INSTALL_RPATH_USE_LINK_PATH Off CACHE Bool "Rpath uses Link path")
+set(SAMRAI_RAJA_WORKGROUP_THREADS 512 CACHE INT "Number of workgroup threads")
 
 include(GNUInstallDirs)
 

config/SAMRAI_config.h.cmake.in

@@ -340,6 +340,8 @@
 /* Maximum dimension allowed */
 #define SAMRAI_MAXIMUM_DIMENSION @SAMRAI_MAXIMUM_DIMENSION@
 
+#define SAMRAI_RAJA_WORKGROUP_THREADS @SAMRAI_RAJA_WORKGROUP_THREADS@
+
 /* Define to 1 if you have the ANSI C header files. */
 #undef STDC_HEADERS
 
@@ -358,7 +360,9 @@
 /* Configure for compiling on BGL family of machines */
 #undef __BGL_FAMILY__
 
-
+#ifdef HAVE_RAJA
+#define SAMRAI_HAVE_KERNEL_FUSER
+#endif
 
 namespace SAMRAI {
    static const unsigned short MAX_DIM_VAL = SAMRAI_MAXIMUM_DIMENSION;

source/SAMRAI/hier/CMakeLists.txt

@@ -140,6 +140,10 @@ if (ENABLE_MPI)
   set (hier_depends ${hier_depends} mpi)
 endif ()
 
+if (ENABLE_CUDA)
+  set (hier_depends ${hier_depends} cuda)
+endif ()
+
 blt_add_library(
   NAME SAMRAI_hier
   SOURCES ${hier_sources}

source/SAMRAI/hier/ForAll.h

@@ -19,6 +19,7 @@
 #include "SAMRAI/hier/Box.h"
 #include "SAMRAI/hier/Index.h"
 #include "SAMRAI/tbox/ExecutionPolicy.h"
+#include "SAMRAI/tbox/KernelFuser.h"
 
 #include <type_traits>
 #include <tuple>
@@ -145,8 +146,37 @@ struct for_all<1> {
           RAJA::make_tuple(make_range(ifirst, ilast, 0)),
           body);
    }
+
+   template <typename Policy, typename LoopBody,
+             typename std::enable_if<std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+   inline static void eval(tbox::KernelFuser* fuser, const hier::Index& ifirst, const hier::Index& ilast, LoopBody body)
+   {
+      if (fuser == nullptr) {
+         RAJA::kernel<typename tbox::detail::policy_traits<Policy>::Policy1d>(
+             RAJA::make_tuple(make_range(ifirst, ilast, 0)),
+             body);
+      } else {
+         fuser->enqueue(ifirst(0), ilast(0), body);
+      }
+   }
+
+   template <typename Policy, typename LoopBody,
+             typename std::enable_if<!std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+   inline static void eval(tbox::KernelFuser* fuser, const hier::Index& ifirst, const hier::Index& ilast, LoopBody body)
+   {
+      if (fuser == nullptr) {
+         RAJA::kernel<Policy>(
+             RAJA::make_tuple(make_range(ifirst, ilast, 0)),
+             body);
+      } else {
+         fuser->enqueue(ifirst(0), ilast(0), body);
+      }
+   }
 };
 
+
+// 2D and 3D don't use the fuser for anything pending suppor for
+// multidimensional loops in KernelFuser.
 template <>
 struct for_all<2> {
    template <typename Policy, typename LoopBody,
@@ -168,6 +198,28 @@ struct for_all<2> {
                            make_range(ifirst, ilast, 1)),
           body);
    }
+
+   template <typename Policy, typename LoopBody,
+             typename std::enable_if<std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+   inline static void eval(tbox::KernelFuser* fuser, const hier::Index& ifirst, const hier::Index& ilast, LoopBody body)
+   {
+      NULL_USE(fuser);
+      RAJA::kernel<typename tbox::detail::policy_traits<Policy>::Policy2d>(
+          RAJA::make_tuple(make_range(ifirst, ilast, 0),
+                           make_range(ifirst, ilast, 1)),
+          body);
+   }
+
+   template <typename Policy, typename LoopBody,
+             typename std::enable_if<!std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+   inline static void eval(tbox::KernelFuser* fuser, const hier::Index& ifirst, const hier::Index& ilast, LoopBody body)
+   {
+      NULL_USE(fuser);
+      RAJA::kernel<Policy>(
+          RAJA::make_tuple(make_range(ifirst, ilast, 0),
+                           make_range(ifirst, ilast, 1)),
+          body);
+   }
 };
 
 template <>
@@ -193,6 +245,30 @@ struct for_all<3> {
                            make_range(ifirst, ilast, 2)),
           body);
    }
+
+   template <typename Policy, typename LoopBody,
+             typename std::enable_if<std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+   inline static void eval(tbox::KernelFuser* fuser, const hier::Index& ifirst, const hier::Index& ilast, LoopBody body)
+   {
+      NULL_USE(fuser);
+      RAJA::kernel<typename tbox::detail::policy_traits<Policy>::Policy3d>(
+          RAJA::make_tuple(make_range(ifirst, ilast, 0),
+                           make_range(ifirst, ilast, 1),
+                           make_range(ifirst, ilast, 2)),
+          body);
+   }
+
+   template <typename Policy, typename LoopBody,
+             typename std::enable_if<!std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+   inline static void eval(tbox::KernelFuser* fuser, const hier::Index& ifirst, const hier::Index& ilast, LoopBody body)
+   {
+      NULL_USE(fuser);
+      RAJA::kernel<Policy>(
+          RAJA::make_tuple(make_range(ifirst, ilast, 0),
+                           make_range(ifirst, ilast, 1),
+                           make_range(ifirst, ilast, 2)),
+          body);
+   }
 };
 
 }  // namespace detail
@@ -205,20 +281,52 @@ inline void for_all(int begin, int end, LoopBody body)
    RAJA::forall<typename tbox::detail::policy_traits<Policy>::Policy>(RAJA::RangeSegment(begin, end), body);
 }
 
+template <typename Policy, typename LoopBody,
+          typename std::enable_if<std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+inline void for_all(tbox::KernelFuser* fuser, int begin, int end, LoopBody body)
+{
+   if (fuser == nullptr) {
+      RAJA::forall<typename tbox::detail::policy_traits<Policy>::Policy>(RAJA::RangeSegment(begin, end), body);
+   } else {
+      fuser->enqueue(begin, end, body);
+   }
+}
+
 template <typename Policy, typename LoopBody,
           typename std::enable_if<!std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
 inline void for_all(int begin, int end, LoopBody body)
 {
    RAJA::forall<Policy>(RAJA::RangeSegment(begin, end), body);
 }
 
+template <typename Policy, typename LoopBody,
+          typename std::enable_if<!std::is_base_of<tbox::policy::base, Policy>::value, int>::type = 0>
+inline void for_all(tbox::KernelFuser* fuser, int begin, int end, LoopBody body)
+{
+   if (fuser == nullptr) {
+      RAJA::forall<Policy>(RAJA::RangeSegment(begin, end), body);
+   } else {
+      fuser->enqueue(begin, end, body);
+   }
+}
+
 // does NOT include end
 template <typename LoopBody>
 inline void parallel_for_all(int begin, int end, LoopBody body)
 {
    for_all<tbox::policy::parallel>(begin, end, body);
 }
 
+template <typename LoopBody>
+inline void parallel_for_all(tbox::KernelFuser* fuser, int begin, int end, LoopBody body)
+{
+   if (fuser == nullptr) {
+      for_all<tbox::policy::parallel>(begin, end, body);
+   } else {
+      for_all<tbox::policy::parallel>(fuser, begin, end, body);
+   }
+}
+
 template <typename LoopBody>
 inline void host_parallel_for_all(int begin, int end, LoopBody body)
 {
@@ -231,12 +339,25 @@ inline void for_all(const hier::Box& box, const int dim, LoopBody body)
    for_all<Policy>(box.lower()(dim), box.upper()(dim) + 1, body);
 }
 
+
+template <typename Policy, typename LoopBody>
+inline void for_all(tbox::KernelFuser* fuser, const hier::Box& box, const int dim, LoopBody body)
+{
+   for_all<Policy>(fuser, box.lower()(dim), box.upper()(dim) + 1, body);
+}
+
 template <typename LoopBody>
 inline void parallel_for_all(const hier::Box& box, const int dim, LoopBody body)
 {
    for_all<tbox::policy::parallel>(box.lower()(dim), box.upper()(dim) + 1, body);
 }
 
+template <typename LoopBody>
+inline void parallel_for_all(tbox::KernelFuser* fuser, const hier::Box& box, const int dim, LoopBody body)
+{
+   for_all<tbox::policy::parallel>(fuser, box.lower()(dim), box.upper()(dim) + 1, body);
+}
+
 template <typename LoopBody>
 inline void host_parallel_for_all(const hier::Box& box, const int dim, LoopBody body)
 {
@@ -250,12 +371,29 @@ inline void for_all(const hier::Box& box, LoopBody body)
    detail::for_all<arg_count>::template eval<Policy>(box.lower(), box.upper(), body);
 }
 
+template <typename Policy, typename LoopBody>
+inline void for_all(tbox::KernelFuser* fuser, const hier::Box& box, LoopBody body)
+{
+   if (fuser == nullptr) {
+      for_all<Policy,LoopBody>(box, body);
+   } else {
+      constexpr int arg_count = detail::function_traits<LoopBody>::argument_count;
+      detail::for_all<arg_count>::template eval<Policy>(fuser, box.lower(), box.upper(), body);
+   }
+}
+
 template <typename LoopBody>
 inline void parallel_for_all(const hier::Box& box, LoopBody body)
 {
    for_all<tbox::policy::parallel>(box, body);
 }
 
+template <typename LoopBody>
+inline void parallel_for_all(tbox::KernelFuser* fuser, const hier::Box& box, LoopBody body)
+{
+   for_all<tbox::policy::parallel>(fuser, box, body);
+}
+
 template <typename LoopBody>
 inline void host_parallel_for_all(const hier::Box& box, LoopBody body)
 {

source/SAMRAI/hier/PatchData.cpp

@@ -31,6 +31,30 @@ PatchData::~PatchData()
 {
 }
 
+void
+PatchData::copyFuseable(
+   const PatchData& src,
+   const BoxOverlap& overlap)
+{
+   copy(src, overlap);
+}
+
+void
+PatchData::packStreamFuseable(
+   tbox::MessageStream& stream,
+   const BoxOverlap& overlap) const
+{
+   packStream(stream, overlap);
+}
+
+void
+PatchData::unpackStreamFuseable(
+   tbox::MessageStream& stream,
+   const BoxOverlap& overlap)
+{
+   unpackStream(stream, overlap);
+}
+
 /*
  *************************************************************************
  *

source/SAMRAI/hier/PatchData.h

@@ -20,6 +20,15 @@
 #include "SAMRAI/tbox/Utilities.h"
 
 namespace SAMRAI {
+
+/*
+ * Forward declaration of KernelFuser class - required here because it sucks in
+ * RAJA and requires CUDA.
+ */
+//namespace tbox {
+//class KernelFuser;
+//}
+
 namespace hier {
 
 /**
@@ -160,6 +169,11 @@ class PatchData
       const PatchData& src,
       const BoxOverlap& overlap) = 0;
 
+   virtual void
+   copyFuseable(
+      const PatchData& src,
+      const BoxOverlap& overlap);
+
    /**
     * Copy data from the source into the destination using the designated
     * overlap descriptor.  The overlap description will have been computed
@@ -206,6 +220,18 @@ class PatchData
       tbox::MessageStream& stream,
       const BoxOverlap& overlap) const = 0;
 
+   /**
+    * Pack data lying on the specified index set into the output stream using
+    * the given KernelFuser. The default implementation of this method will
+    * call packStream without the fuser argument. See the abstract stream
+    * virtual base class for more information about the packing operators
+    * defined for streams.
+    */
+   virtual void
+   packStreamFuseable(
+      tbox::MessageStream& stream,
+      const BoxOverlap& overlap) const;
+
    /**
     * Unpack data from the message stream into the specified index set.
     * See the abstract stream virtual base class for more information about
@@ -216,6 +242,18 @@ class PatchData
       tbox::MessageStream& stream,
       const BoxOverlap& overlap) = 0;
 
+   /**
+    * Unpack data from the message stream into the specified index set using
+    * the given KernelFuser. The default implementation of this method will
+    * call unpackStream without the fuser argument. See the abstract stream
+    * virtual base class for more information about the packing operators
+    * defined for streams.
+    */
+   virtual void
+   unpackStreamFuseable(
+      tbox::MessageStream& stream,
+      const BoxOverlap& overlap);
+
    /**
     * Checks that class version and restart file version are equal.  If so,
     * reads in the data members common to all patch data types from restart

source/SAMRAI/hier/PatchLevel.cpp

@@ -9,7 +9,9 @@
  ************************************************************************/
 #include "SAMRAI/hier/PatchLevel.h"
 
+#include "SAMRAI/tbox/Collectives.h"
 #include "SAMRAI/tbox/MathUtilities.h"
+#include "SAMRAI/tbox/StagedKernelFusers.h"
 #include "SAMRAI/tbox/TimerManager.h"
 #include "SAMRAI/hier/BaseGridGeometry.h"
 #include "SAMRAI/hier/BoxContainer.h"