Realm backend

.flake/pkgs/legion.nix

@@ -18,13 +18,13 @@ in
 
 stdenv.mkDerivation rec {
   pname = "legion_flexflow";
-  version = "2024-03-13";
+  version = "2025-01-21";
 
   src = fetchFromGitLab {
     owner = "StanfordLegion";
     repo = "legion";
-    rev = "24e8c452341dea41427e0ce61e154d61715e6835";
-    sha256 = "sha256-NjCSjphOIew/V24i74I6DModSGcWKLeiSIjts3cFtx4=";
+    rev = "0c5a181e59c07e3af1091a2007378ff9355047fa";
+    sha256 = "sha256-oapo7klN17gmRsmaSsrpup4YJ0dtHxiKFtwz8jyPqzU=";
     fetchSubmodules = true;
   };
 
@@ -33,7 +33,7 @@ stdenv.mkDerivation rec {
   ];
 
   cmakeFlags = [
-    "-DLegion_USE_Python=1"
+    "-DLegion_USE_Python=0"
     "-DLegion_BUILD_BINDINGS=1"
     "-DLegion_USE_CUDA=1"
     "-DLegion_CUDA_ARCH=${lib.concatStringsSep "," cudaCapabilities}"

.proj.toml

@@ -70,6 +70,13 @@ has-cpu-only-benchmarks = false
 has-cuda-tests = true
 has-cuda-benchmarks = false
 
+[targets.realm-backend]
+type = "lib"
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = true
+has-cuda-benchmarks = false
+
 [targets.models]
 type = "lib"
 has-cpu-only-tests = true

lib/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(runtime)
 add_subdirectory(op-attrs)
 add_subdirectory(kernels)
 add_subdirectory(local-execution)
+add_subdirectory(realm-backend)
 add_subdirectory(task-spec)
 add_subdirectory(utils)
 add_subdirectory(ffi)

lib/local-execution/include/local-execution/local_args_backing.h

@@ -24,24 +24,14 @@ std::optional<DeviceSpecificDeviceStates>
 std::unordered_map<slot_id_t, ConcreteArgSpec>
     construct_arg_slots_backing(TaskBinding const &, RuntimeArgConfig const &);
 
-std::optional<DeviceSpecificDeviceStates>
-    create_per_device_op_state(LocalTaskRegistry const &,
-                               LocalTensorBacking const &,
-                               RuntimeArgConfig const &,
-                               Allocator &,
-                               TrainingLayerPlusContext const &);
-
 TaskArgumentAccessor get_task_arg_accessor(LocalTensorBacking const &,
                                            RuntimeArgConfig const &,
                                            TaskInvocation const &,
                                            Allocator &);
 
 LocalArgsBacking make_local_args_backing_for_computation_graph(
-    LocalTaskRegistry const &,
-    TrainingComputationGraph const &,
     RuntimeArgConfig const &,
-    LocalTensorBacking const &,
-    Allocator &);
+    std::unordered_map<layer_guid_t, std::optional<DeviceSpecificDeviceStates>> const &);
 
 } // namespace FlexFlow
 

lib/local-execution/include/local-execution/local_training_backing.h

@@ -6,6 +6,7 @@
 #include "pcg/optimizer_attrs.dtg.h"
 #include "task-spec/training_computation_graph.dtg.h"
 #include "task-spec/training_tensor_guid_t.dtg.h"
+#include "utils/containers/generate_map.h"
 #include "utils/units/milliseconds_t.h"
 
 namespace FlexFlow {
@@ -18,6 +19,13 @@ LocalTrainingBacking make_local_training_backing_for_computation_graph(
     RuntimeArgConfig const &runtime_arg_config,
     OptimizerAttrs const &optimizer_attrs);
 
+std::optional<DeviceSpecificDeviceStates>
+    create_per_device_op_state(LocalTaskRegistry const &,
+                               LocalTensorBacking const &,
+                               RuntimeArgConfig const &,
+                               Allocator &,
+                               TrainingLayerPlusContext const &);
+
 std::optional<milliseconds_t> execute_forward(LocalTaskRegistry const &,
                                               LocalTensorBacking const &,
                                               LocalArgsBacking const &,

lib/local-execution/src/local-execution/local_args_backing.cc

@@ -35,38 +35,6 @@ std::unordered_map<slot_id_t, ConcreteArgSpec>
   ;
 }
 
-std::optional<DeviceSpecificDeviceStates>
-    create_per_device_op_state(LocalTaskRegistry const &local_task_registry,
-                               LocalTensorBacking const &tensor_backing,
-                               RuntimeArgConfig const &runtime_arg_config,
-                               Allocator &allocator,
-                               TrainingLayerPlusContext const &training_layer) {
-  std::optional maybe_registered_task = try_get_registered_task(
-      local_task_registry, training_layer.layer_guid, OpTaskType::INIT);
-
-  ASSERT(maybe_registered_task.has_value());
-
-  registered_task_t registered_task = maybe_registered_task.value();
-  if (registered_task.is_noop_task()) {
-    return std::nullopt;
-  }
-
-  TaskInvocation invocation = lower_to_task_invocation(
-      /*op_task_invocation=*/get_init_op_task_invocation(
-          training_layer.layer_attrs.op_attrs),
-      /*training_layer=*/training_layer,
-      /*device_specific_device_states=*/std::nullopt);
-
-  TaskArgumentAccessor accessor = get_task_arg_accessor(
-      tensor_backing, runtime_arg_config, invocation, allocator);
-  TaskSignatureAndImpl task_sig_impl =
-      local_task_registry.task_mapping.at(invocation.task_id);
-  auto fn =
-      task_sig_impl.impl_function.get<InitOpTaskImplFunction>().function_ptr;
-  DeviceSpecificDeviceStates device_state = fn(accessor);
-  return device_state;
-}
-
 TaskArgumentAccessor
     get_task_arg_accessor(LocalTensorBacking const &local_tensor_backing,
                           RuntimeArgConfig const &runtime_arg_config,
@@ -82,24 +50,9 @@ TaskArgumentAccessor
 }
 
 LocalArgsBacking make_local_args_backing_for_computation_graph(
-    LocalTaskRegistry const &task_registry,
-    TrainingComputationGraph const &training_computation_graph,
     RuntimeArgConfig const &runtime_arg_config,
-    LocalTensorBacking const &local_tensor_backing,
-    Allocator &allocator) {
-  std::unordered_map<layer_guid_t, std::optional<DeviceSpecificDeviceStates>>
-      per_device_op_states = generate_map(
-          topological_ordering(training_computation_graph.computation_graph),
-          [&](layer_guid_t const &layer_guid) {
-            return create_per_device_op_state(
-                task_registry,
-                local_tensor_backing,
-                runtime_arg_config,
-                allocator,
-                get_training_layer_plus_context(training_computation_graph,
-                                                layer_guid));
-          });
-
+    std::unordered_map<layer_guid_t, std::optional<DeviceSpecificDeviceStates>> const &
+        per_device_op_states) {
   return LocalArgsBacking{
       runtime_arg_config,
       per_device_op_states,

lib/local-execution/src/local-execution/local_training_backing.cc

@@ -39,12 +39,22 @@ LocalTrainingBacking make_local_training_backing_for_computation_graph(
       preallocated,
       allocator);
 
+  std::unordered_map<layer_guid_t, std::optional<DeviceSpecificDeviceStates>>
+      per_device_op_states = generate_map(
+          topological_ordering(training_computation_graph.computation_graph),
+          [&](layer_guid_t const &layer_guid) {
+            return create_per_device_op_state(
+                local_task_registry,
+                local_tensor_backing,
+                runtime_arg_config,
+                allocator,
+                get_training_layer_plus_context(training_computation_graph,
+                                                layer_guid));
+          });
+
   LocalArgsBacking local_args_backing =
-      make_local_args_backing_for_computation_graph(local_task_registry,
-                                                    training_computation_graph,
-                                                    runtime_arg_config,
-                                                    local_tensor_backing,
-                                                    allocator);
+      make_local_args_backing_for_computation_graph(runtime_arg_config,
+                                                    per_device_op_states);
 
   return LocalTrainingBacking{
       /*computation_graph=*/training_computation_graph,
@@ -54,6 +64,38 @@ LocalTrainingBacking make_local_training_backing_for_computation_graph(
   };
 }
 
+std::optional<DeviceSpecificDeviceStates>
+    create_per_device_op_state(LocalTaskRegistry const &local_task_registry,
+                               LocalTensorBacking const &tensor_backing,
+                               RuntimeArgConfig const &runtime_arg_config,
+                               Allocator &allocator,
+                               TrainingLayerPlusContext const &training_layer) {
+  std::optional maybe_registered_task = try_get_registered_task(
+      local_task_registry, training_layer.layer_guid, OpTaskType::INIT);
+
+  ASSERT(maybe_registered_task.has_value());
+
+  registered_task_t registered_task = maybe_registered_task.value();
+  if (registered_task.is_noop_task()) {
+    return std::nullopt;
+  }
+
+  TaskInvocation invocation = lower_to_task_invocation(
+      /*op_task_invocation=*/get_init_op_task_invocation(
+          training_layer.layer_attrs.op_attrs),
+      /*training_layer=*/training_layer,
+      /*device_specific_device_states=*/std::nullopt);
+
+  TaskArgumentAccessor accessor = get_task_arg_accessor(
+      tensor_backing, runtime_arg_config, invocation, allocator);
+  TaskSignatureAndImpl task_sig_impl =
+      local_task_registry.task_mapping.at(invocation.task_id);
+  auto fn =
+      task_sig_impl.impl_function.get<InitOpTaskImplFunction>().function_ptr;
+  DeviceSpecificDeviceStates device_state = fn(accessor);
+  return device_state;
+}
+
 std::optional<milliseconds_t>
     execute_forward(LocalTaskRegistry const &local_task_registry,
                     LocalTensorBacking const &local_tensor_backing,

lib/realm-backend/CMakeLists.txt

@@ -0,0 +1,21 @@
+ff_add_library(
+  NAME
+    realm-backend
+  SRC_PATTERNS
+    src/*.cc
+  PUBLIC_INCLUDE
+    include/
+  PRIVATE_INCLUDE
+    src/
+  DEPS
+    op-attrs
+    utils
+    kernels
+    compiler
+    local-execution
+    pcg
+    spdlog
+    legion
+)
+
+add_subdirectory(test)

lib/realm-backend/include/realm-backend/driver.h

@@ -0,0 +1,13 @@
+#ifndef _FLEXFLOW_REALM_BACKEND_DRIVER_H
+#define _FLEXFLOW_REALM_BACKEND_DRIVER_H
+
+#include "realm.h"
+#include "realm/cmdline.h"
+#include "task-spec/op_task_invocation.h"
+
+Realm::Processor::TaskFuncID get_realm_task_id(FlexFlow::task_id_t task_id);
+
+void top_level_task(const void *args, size_t arglen, const void *userdata,
+               size_t userlen, Realm::Processor p);
+
+#endif

lib/realm-backend/include/realm-backend/model_training_instance.h

@@ -0,0 +1,31 @@
+#ifndef _FLEXFLOW_LOCAL_EXECUTION_MODEL_TRAINING_INSTANCE_H
+#define _FLEXFLOW_LOCAL_EXECUTION_MODEL_TRAINING_INSTANCE_H
+
+#include "realm-backend/realm_training_backing.h"
+#include "op-attrs/ops/loss_functions/loss_attrs.dtg.h"
+#include "pcg/tensor_guid_t.dtg.h"
+#include "task-spec/loss_tensor_guid_t.dtg.h"
+
+namespace FlexFlow {
+
+struct ModelTrainingInstance {
+  ModelTrainingInstance(RealmRuntimeState &,
+                        LocalTrainingBacking const &,
+                        LossAttrs const &,
+                        OptimizerAttrs const &);
+
+  RealmRuntimeState &runtime_state;
+  LocalTrainingBacking training_backing;
+  LossAttrs loss_attrs;
+  OptimizerAttrs optimizer_attrs;
+
+public:
+  std::unordered_map<layer_guid_t, std::optional<milliseconds_t>> forward();
+  std::unordered_map<layer_guid_t, std::optional<milliseconds_t>> backward();
+  void update();
+  GenericTensorAccessorR get_loss_tensor_accessor() const;
+};
+
+} // namespace FlexFlow
+
+#endif

lib/realm-backend/include/realm-backend/realm_allocator.h

@@ -0,0 +1,34 @@
+#ifndef _FLEXFLOW_REALM_BACKEND_REALM_ALLOCATOR_H
+#define _FLEXFLOW_REALM_BACKEND_REALM_ALLOCATOR_H
+
+#include "realm-backend/driver.h"
+#include "realm.h"
+#include "kernels/allocation.h"
+#include <realm/event.h>
+
+namespace FlexFlow {
+
+struct RealmAllocatorImpl : public IAllocator {
+  RealmAllocatorImpl() = delete;
+  RealmAllocatorImpl(RealmAllocatorImpl const &) = delete;
+  RealmAllocatorImpl(RealmAllocatorImpl &&) = delete;
+  RealmAllocatorImpl(Realm::Processor);
+  ~RealmAllocatorImpl() = default;
+
+  void *allocate(size_t) override;
+  void deallocate(void *) override;
+
+  DeviceType get_allocation_device_type() const override;
+
+private:
+  std::unordered_map<void *, Realm::RegionInstance> ptrs;
+  Realm::Processor proc;
+  Realm::Memory mem;
+  std::vector<size_t> field_sizes = {sizeof(char)};
+};
+
+Allocator create_realm_memory_allocator(Realm::Processor);
+
+} // namespace FlexFlow
+
+#endif

lib/realm-backend/include/realm-backend/realm_training_backing.h

@@ -0,0 +1,64 @@
+#ifndef _FLEXFLOW_REALM_BACKEND_REALM_TRAINING_BACKING_H
+#define _FLEXFLOW_REALM_BACKEND_REALM_TRAINING_BACKING_H
+
+#include "local-execution/local_training_backing.dtg.h"
+#include "op-attrs/ops/loss_functions/loss_attrs.dtg.h"
+#include "pcg/optimizer_attrs.dtg.h"
+#include "task-spec/training_computation_graph.dtg.h"
+#include "task-spec/training_tensor_guid_t.dtg.h"
+#include "utils/containers/generate_map.h"
+#include "utils/units/milliseconds_t.h"
+#include "realm-backend/driver.h"
+#include "realm-backend/realm_allocator.h"
+#include "realm-backend/task_wrapper.h"
+
+namespace FlexFlow {
+
+struct RealmRuntimeState {
+  Realm::Processor master_proc;
+  Realm::Event master_event;
+  Realm::Memory master_mem;
+  std::vector<Realm::Processor> worker_procs;
+  std::vector<Realm::Event> worker_events;
+  std::vector<Allocator> allocators;
+};
+
+LocalTrainingBacking make_realm_training_backing_for_computation_graph(
+    RealmRuntimeState &runtime_state,
+    std::unordered_map<training_tensor_guid_t, GenericTensorAccessorW> const
+        &preallocated_tensors,
+    TrainingComputationGraph const &training_computation_graph,
+    RuntimeArgConfig const &runtime_arg_config,
+    OptimizerAttrs const &optimizer_attrs);
+
+void register_tasks_for_realm(LocalTaskRegistry const &, RealmRuntimeState &);
+
+std::optional<DeviceSpecificDeviceStates>
+    create_per_device_op_state(LocalTaskRegistry const &,
+                               LocalTensorBacking const &,
+                               RuntimeArgConfig const &,
+                               RealmRuntimeState &,
+                               TrainingLayerPlusContext const &);
+
+Future<std::optional<milliseconds_t>> execute_forward(LocalTaskRegistry const &,
+                                              LocalTensorBacking const &,
+                                              LocalArgsBacking const &,
+                                              TrainingLayerPlusContext const &,
+                                              RealmRuntimeState &);
+
+Future<std::optional<milliseconds_t>> execute_backward(LocalTaskRegistry const &,
+                                               LocalTensorBacking const &,
+                                               LocalArgsBacking const &,
+                                               TrainingLayerPlusContext const &,
+                                               RealmRuntimeState &);
+
+Future<void> compute_loss(LocalTrainingBacking const &, LossAttrs const &, RealmRuntimeState &);
+
+Future<void> execute_update(LocalTrainingBacking const &,
+                    layer_guid_t const &,
+                    OptimizerAttrs const &,
+                    RealmRuntimeState &);
+
+} // namespace FlexFlow
+
+#endif