Add CPU/GPU Combine/Reduction/Repartition/Replicate support to realm-execution backend

lib/realm-execution/CMakeLists.txt

@@ -1,3 +1,5 @@
+project(realm-execution LANGUAGES CXX CUDA)
+
 ff_add_library(
   NAME
     realm-execution
@@ -19,4 +21,23 @@ ff_add_library(
     deps::realm
 )
 
+# redops/realm_redop_registry.cc needs to be compiled by nvcc (not g++) so
+# that its GPU reduction kernels (SumReduction<T>::apply_cuda/fold_cuda)
+# actually get generated; see the comment on register_sum_redop for why
+# that matters. Kept as a .cc file (rather than renamed to .cu, as lib/
+# kernels does for its own CUDA sources) so it stays subject to proj's
+# normal .h/.cc file-group layout check — realm-execution isn't in
+# .proj.toml's layout_ignore_paths the way lib/kernels is.
+set_source_files_properties(
+  src/realm-execution/redops/realm_redop_registry.cc
+  PROPERTIES
+  LANGUAGE CUDA
+)
+
+set_target_properties(
+  realm-execution
+  PROPERTIES
+  CUDA_STANDARD 17
+)
+
 add_subdirectory(test)

lib/realm-execution/include/realm-execution/redops/realm_redop_registry.h

@@ -1,15 +1,18 @@
 #ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_REDOPS_REALM_REDOP_REGISTRY_H
 #define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_REDOPS_REALM_REDOP_REGISTRY_H
 
-#include "realm-execution/realm.h"
-#include "realm-execution/redops/redop_id_t.dtg.h"
-
 namespace FlexFlow {
 
 /**
- * \brief Registers all known reduction operators (redops).
+ * \brief Registers all known reduction operators (redops), fetching the
+ * process-global Realm runtime itself.
+ *
+ * \note Deliberately takes no PRealm-typed argument and is implemented in a
+ * .cu file (needed to compile GPU reduction kernels for the redops) —
+ * including realm-execution/realm.h (i.e. PRealm's prealm.h) from a .cu
+ * translation unit fails to compile under nvcc.
  */
-void register_all_redops(Realm::Runtime);
+void register_all_redops();
 
 } // namespace FlexFlow
 

lib/realm-execution/src/realm-execution/pcg_instance.cc

@@ -16,6 +16,7 @@
 #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
 #include "task-spec/dynamic_graph/loss_insertion.h"
 #include "task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h"
+#include "task-spec/dynamic_graph/parallel_op_data_movement.h"
 #include "task-spec/dynamic_graph/pass_expansion.h"
 #include "task-spec/dynamic_graph/shard_expansion.h"
 #include "task-spec/dynamic_graph/training_operation_attrs.dtg.h"
@@ -270,16 +271,16 @@ static Realm::Event spawn_dynamic_node_invocation(
     OptimizerAttrs const &optimizer_attrs,
     ProfilingSettings const &profiling_settings,
     DistributedFfHandle const &device_handle) {
-  Realm::Event precondition = Realm::Event::merge_events(
+  Realm::Event const precondition = Realm::Event::merge_events(
       Realm::Event::merge_events(input_dependencies),
       Realm::Event::merge_events(output_dependencies));
 
-  TensorInstanceBacking tensor_backing =
+  TensorInstanceBacking const tensor_backing =
       subset_tensor_instance_backing_for_invocation(tensor_instance_backing,
                                                     invocation);
 
   auto spawn_task = [&]() {
-    Realm::Processor target_proc = ctx.processor_from_global_device_id(
+    Realm::Processor const target_proc = ctx.processor_from_global_device_id(
         get_only(assert_unwrap(invocation.node_attrs.device_ids)));
     return spawn_op_task(ctx,
                          target_proc,
@@ -299,48 +300,71 @@ static Realm::Event spawn_dynamic_node_invocation(
         ctx, input, output, tensor_instance_backing, precondition);
   };
 
-  auto issue_replicate = [&]() {
+  // N inputs → 1 output (copy semantics). Each of the N copies targets the
+  // same destination instance, so they are chained (each waiting on the
+  // previous one's completion) rather than fired concurrently — issuing them
+  // all against the same precondition would race N unsynchronized writes to
+  // overlapping destination memory.
+  auto issue_gather = [&]() {
+    DynamicValueAttrs const &output = get_only(invocation.outputs).second;
+    Realm::Event last = precondition;
+    for (auto const &[slot, input] : invocation.inputs) {
+      last = issue_p2p_copy(ctx, input, output, tensor_instance_backing, last);
+    }
+    return last;
+  };
+
+  auto issue_broadcast = [&]() {
     DynamicValueAttrs const &input = get_only(invocation.inputs).second;
-    std::vector<DynamicValueAttrs> outputs =
+    std::vector<DynamicValueAttrs> const outputs =
         vector_of(values(invocation.outputs));
     return issue_collective_broadcast(
         ctx, input, outputs, tensor_instance_backing, precondition);
   };
 
-  auto issue_reduction = [&]() {
-    std::vector<DynamicValueAttrs> inputs =
+  auto issue_sum_reduce = [&]() {
+    std::vector<DynamicValueAttrs> const inputs =
         vector_of(values(invocation.inputs));
     DynamicValueAttrs const &output = get_only(invocation.outputs).second;
-    redop_id_t redop_id = get_sum_redop_id_for_data_type(
+    redop_id_t const redop_id = get_sum_redop_id_for_data_type(
         assert_unwrap(output.parallel_tensor_shape).data_type);
     return issue_collective_reduction(
         ctx, inputs, output, tensor_instance_backing, redop_id, precondition);
   };
 
-  TrainingOperationAttrs op_attrs =
+  TrainingOperationAttrs const op_attrs =
       assert_unwrap(invocation.node_attrs.op_attrs);
+
   return op_attrs.visit<Realm::Event>(overload{
-      [&](PCGOperatorAttrs const &pcg_op_attrs) {
+      [&](PCGOperatorAttrs const &pcg_op_attrs) -> Realm::Event {
+        std::optional<ParallelOpMovementKind> const movement_kind =
+            get_parallel_op_movement_kind_for_training_op(
+                op_attrs, assert_unwrap(invocation.node_attrs.task_type));
+        if (movement_kind.has_value()) {
+          switch (movement_kind.value()) {
+            case ParallelOpMovementKind::BROADCAST:
+              return issue_broadcast();
+            case ParallelOpMovementKind::GATHER:
+              return issue_gather();
+            case ParallelOpMovementKind::SUM_REDUCE:
+              return issue_sum_reduce();
+            case ParallelOpMovementKind::RESHUFFLE:
+              return issue_copy();
+          }
+        }
+
         return pcg_op_attrs.visit<Realm::Event>(overload{
-            [&](InputAttrs const &) { return Realm::Event::NO_EVENT; },
-            [&](WeightAttrs const &) { return Realm::Event::NO_EVENT; },
-            [&](ReplicateAttrs const &) {
-              DynamicTaskType task_type =
-                  assert_unwrap(invocation.node_attrs.task_type);
-              switch (task_type) {
-                case DynamicTaskType::FWD:
-                  return issue_replicate();
-                case DynamicTaskType::BWD:
-                  return issue_reduction();
-                default:
-                  PANIC("Unhandled replicate task type ", task_type);
-              }
+            [&](InputAttrs const &) -> Realm::Event {
+              return Realm::Event::NO_EVENT;
+            },
+            [&](WeightAttrs const &) -> Realm::Event {
+              return Realm::Event::NO_EVENT;
             },
-            [&](auto const &) { return spawn_task(); },
+            [&](auto const &) -> Realm::Event { return spawn_task(); },
         });
       },
-      [&](LossAttrs const &) { return spawn_task(); },
-      [&](CopyAttrs const &) { return issue_copy(); },
+      [&](LossAttrs const &) -> Realm::Event { return spawn_task(); },
+      [&](CopyAttrs const &) -> Realm::Event { return issue_copy(); },
   });
 }
 

lib/realm-execution/src/realm-execution/realm_manager.cc

@@ -12,7 +12,7 @@ RealmManager::RealmManager(int *argc, char ***argv)
 
   // Register all tasks and redops at initialization time so we don't need to later
   register_all_tasks().wait();
-  register_all_redops(this->get_runtime());
+  register_all_redops();
 }
 
 RealmManager::~RealmManager() {

lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc

@@ -1,5 +1,10 @@
 #include "realm-execution/redops/realm_redop_registry.h"
-#include "realm-execution/redops/redop_id_t.h"
+#include "realm-execution/redops/redop_id_t.dtg.h"
+#include <cassert>
+#include <realm.h>
+#if defined(__CUDACC__) || defined(__HIPCC__)
+#include <realm/cuda/cuda_redop.h>
+#endif
 
 namespace FlexFlow {
 
@@ -8,16 +13,7 @@ namespace FlexFlow {
 // existing code, despite not otherwise relying or using Legion in any way.
 // https://gitlab.com/StanfordLegion/legion/-/blob/5263aeff477fb94239c50d9306d58c4244e9fc38/runtime/legion/api/redop.inl#L31
 #if !defined(__cpp_lib_atomic_ref) || (__cpp_lib_atomic_ref < 201806L)
-// We only need this crap if we're using a version of c++ < 20
-// Starting with c++20 we can do all this the right way with atomic_ref
 namespace TypePunning {
-// The tenth circle of hell is reserved for members of the C++ committee
-// that decided to deviate from C's support for type punning unions.
-// Add on to it the fact that it took them 9 fucking years to realize
-// that they needed std::atomic_ref and it's plain to see they are all
-// just a bunch of idiots that should never be allowed near a programming
-// language standard ever again. They've clearly never written lock-free
-// code in their lives.
 template <typename T>
 class Pointer {
 public:
@@ -121,6 +117,56 @@ class Alias {
 #define __LEGION_CUDA_HD__
 #endif
 
+#if defined(__CUDACC__) || defined(__HIPCC__)
+// Legion's non-exclusive bool/int64 reduction paths (borrowed below) assume
+// these helpers exist, but they aren't real CUDA/HIP builtins, so we provide
+// them ourselves.
+//
+// __longlong_as_ulonglong/__ulonglong_as_longlong: signed<->unsigned 64-bit
+// reinterpretation (needed because atomicCAS only takes unsigned long long).
+__device__ __forceinline__ unsigned long long int
+    __longlong_as_ulonglong(long long int v) {
+  return static_cast<unsigned long long int>(v);
+}
+__device__ __forceinline__ long long int
+    __ulonglong_as_longlong(unsigned long long int v) {
+  return static_cast<long long int>(v);
+}
+
+// __uint2bool/__bool2uint: read/write a single bool packed into one byte of
+// a 4-byte-aligned word, at byte `offset` within that word (needed because
+// GPU atomicCAS requires 4-byte alignment, but a lone bool doesn't have
+// one).
+__device__ __forceinline__ bool __uint2bool(unsigned int word,
+                                            unsigned int offset) {
+  return static_cast<bool>((word >> (offset * 8)) & 0xFFu);
+}
+__device__ __forceinline__ unsigned int
+    __bool2uint(unsigned int word, bool value, unsigned int offset) {
+  unsigned int const shift = offset * 8;
+  unsigned int const mask = 0xFFu << shift;
+  unsigned int const byte_val = (static_cast<unsigned int>(value) & 0xFFu)
+                                << shift;
+  return (word & ~mask) | byte_val;
+}
+
+// apply_cuda/fold_cuda are identical (just delegate to apply/fold) for every
+// SumReduction<T> specialization, so factored into a macro rather than
+// repeated per type. Required by Realm::Cuda::add_cuda_redop_kernels to
+// build a GPU-resident reduction kernel for this redop.
+#define __LEGION_CUDA_REDOP_METHODS__                                          \
+  template <bool EXCLUSIVE>                                                    \
+  __device__ static void apply_cuda(LHS &lhs, RHS rhs) {                       \
+    apply<EXCLUSIVE>(lhs, rhs);                                                \
+  }                                                                            \
+  template <bool EXCLUSIVE>                                                    \
+  __device__ static void fold_cuda(RHS &rhs1, RHS rhs2) {                      \
+    fold<EXCLUSIVE>(rhs1, rhs2);                                               \
+  }
+#else
+#define __LEGION_CUDA_REDOP_METHODS__
+#endif
+
 template <typename T>
 class SumReduction {
   // Empty definition
@@ -139,6 +185,7 @@ class SumReduction<bool> {
   __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs);
   template <bool EXCLUSIVE>
   __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2);
+  __LEGION_CUDA_REDOP_METHODS__
 };
 
 template <>
@@ -153,6 +200,7 @@ class SumReduction<int32_t> {
   __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs);
   template <bool EXCLUSIVE>
   __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2);
+  __LEGION_CUDA_REDOP_METHODS__
 };
 
 template <>
@@ -167,6 +215,7 @@ class SumReduction<int64_t> {
   __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs);
   template <bool EXCLUSIVE>
   __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2);
+  __LEGION_CUDA_REDOP_METHODS__
 };
 
 template <>
@@ -181,6 +230,7 @@ class SumReduction<float> {
   __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs);
   template <bool EXCLUSIVE>
   __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2);
+  __LEGION_CUDA_REDOP_METHODS__
 };
 
 template <>
@@ -195,6 +245,7 @@ class SumReduction<double> {
   __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs);
   template <bool EXCLUSIVE>
   __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2);
+  __LEGION_CUDA_REDOP_METHODS__
 };
 
 template <>
@@ -523,18 +574,33 @@ __LEGION_CUDA_HD__ inline void SumReduction<double>::fold<false>(RHS &rhs1,
 #endif
 }
 
-void register_all_redops(Realm::Runtime rt) {
+namespace {
+template <typename REDOP>
+void register_sum_redop(::Realm::Runtime &rt, ::Realm::ReductionOpID id) {
+  ::Realm::ReductionOpUntyped *redop =
+      ::Realm::ReductionOpUntyped::create_reduction_op<REDOP>();
+#if defined(__CUDACC__) || defined(__HIPCC__)
+  ::Realm::Cuda::add_cuda_redop_kernels<REDOP>(redop);
+#endif
+  bool ok = rt.register_reduction(id, redop);
+  assert(ok);
+}
+
+} // namespace
+
+void register_all_redops() {
+  ::Realm::Runtime rt = ::Realm::Runtime::get_runtime();
   // Registration is synchronous, so no need to capture events here
-  rt.register_reduction<SumReduction<bool>>(
-      get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_BOOL_REDOP_ID));
-  rt.register_reduction<SumReduction<int32_t>>(
-      get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_INT32_REDOP_ID));
-  rt.register_reduction<SumReduction<int64_t>>(
-      get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_INT64_REDOP_ID));
-  rt.register_reduction<SumReduction<float>>(
-      get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_FLOAT_REDOP_ID));
-  rt.register_reduction<SumReduction<double>>(
-      get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_DOUBLE_REDOP_ID));
+  register_sum_redop<SumReduction<bool>>(
+      rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_BOOL_REDOP_ID));
+  register_sum_redop<SumReduction<int32_t>>(
+      rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT32_REDOP_ID));
+  register_sum_redop<SumReduction<int64_t>>(
+      rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT64_REDOP_ID));
+  register_sum_redop<SumReduction<float>>(
+      rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_FLOAT_REDOP_ID));
+  register_sum_redop<SumReduction<double>>(
+      rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_DOUBLE_REDOP_ID));
 }
 
 } // namespace FlexFlow