From 4f5c4916fc6bf770cc1d90f546439c44e7895239 Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 13 Aug 2026 14:59:35 -0700 Subject: [PATCH 1/7] Add Combine/Reduction/Repartition support to realm-execution backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generalizes the previously Replicate-only parallel-op handling in the dynamic-graph pipeline (pass_expansion, copy_insertion, shard_expansion) and Realm execution dispatch (pcg_instance) into a single, shared ParallelOpMovementKind abstraction (BROADCAST/GATHER/SUM_REDUCE/RESHUFFLE: 1->N copy, N->1 copy, N->1 sum, N->N reshuffle) that all four parallel ops now go through uniformly: Replicate FWD=BROADCAST BWD=SUM_REDUCE | Combine FWD=GATHER BWD=BROADCAST Reduction FWD=SUM_REDUCE BWD=BROADCAST | Repartition FWD/BWD=RESHUFFLE pass_expansion.cc: BWD expansion for all parallel ops is now the generic is_parallel_training_op case (was Replicate-only) — gradient flows in reverse with no forward activations needed, so BWD inputs are just the grad of FWD outputs and vice versa. copy_insertion.cc: resolves each value's ParallelTensorMapping either directly from the op's own node mapping (only possible on the side whose per-device coordinates are unique — the bidict requirement) or from an adjacent already-resolved value (source, for the non-unique side of BROADCAST; sink, for the non-unique side of GATHER/SUM_REDUCE). RESHUFFLE has no fixed fan-in/fan-out direction (Repartition can scatter, gather, or purely shuffle depending on its degree change), so which side is unique is checked dynamically per-invocation rather than assumed from the movement kind; it also accepts LOSS as a valid sink for the adjacent-value case, since a parallel op's FWD output can feed directly into the loss as the model's terminal output. shard_expansion.cc: RESHUFFLE splits into N separate per-device invocations (like a normal op), while BROADCAST/GATHER/SUM_REDUCE stay as a single invocation with the many-valued side represented as repeated same-slot-name entries distinguished by task_shard — matching how Replicate was already represented, per DynamicTensorSlot.task_shard's docstring. RESHUFFLE pairs each device's input/output coordinates by looking up its own shard binding directly by slot name (coordinate values can otherwise collide between the input and output coordinate spaces). Also removes shard_invocation_for_binding and restrict_tensor_mapping_keys_to_coord (dead code — never called; apply_dynamic_node_invocation_sharding_info is the real path) plus the #includes that became unused as a result. pcg_instance.cc: spawn_dynamic_node_invocation dispatches to issue_broadcast/issue_gather/issue_sum_reduce/issue_copy based on ParallelOpMovementKind rather than switching on each op's attrs type. issue_gather chains its N copies sequentially rather than firing them concurrently, since they all target the same destination instance. Combine/Reduction/Repartition never spawn a task (dispatch is pure Realm copies/reductions), so their now-dead task_id_t.cc entries return nullopt and their realm_task_registry.cc registrations are removed; a few unused lambda-parameter names in task_id_t.cc were tidied at the same time. --- .../src/realm-execution/pcg_instance.cc | 76 ++- ...op_registry.cc => realm_redop_registry.cu} | 0 .../tasks/realm_task_registry.cc | 9 - .../src/realm-execution/tasks/task_id_t.cc | 36 +- .../dynamic_graph/parallel_op_data_movement.h | 36 ++ .../task-spec/dynamic_graph/copy_insertion.cc | 245 ++++++-- .../parallel_op_data_movement.cc | 81 +++ .../task-spec/dynamic_graph/pass_expansion.cc | 6 +- .../dynamic_graph/shard_expansion.cc | 522 ++++++++---------- 9 files changed, 616 insertions(+), 395 deletions(-) rename lib/realm-execution/src/realm-execution/redops/{realm_redop_registry.cc => realm_redop_registry.cu} (100%) create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/parallel_op_data_movement.h create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/parallel_op_data_movement.cc diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc index 1d102f8cca..3e24e3417a 100644 --- a/lib/realm-execution/src/realm-execution/pcg_instance.cc +++ b/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 outputs = + std::vector const outputs = vector_of(values(invocation.outputs)); return issue_collective_broadcast( ctx, input, outputs, tensor_instance_backing, precondition); }; - auto issue_reduction = [&]() { - std::vector inputs = + auto issue_sum_reduce = [&]() { + std::vector 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(overload{ - [&](PCGOperatorAttrs const &pcg_op_attrs) { + [&](PCGOperatorAttrs const &pcg_op_attrs) -> Realm::Event { + std::optional 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(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(); }, }); } diff --git a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu similarity index 100% rename from lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc rename to lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu diff --git a/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc b/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc index 406a2912a3..9d3802910a 100644 --- a/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc +++ b/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc @@ -36,7 +36,6 @@ Realm::Event register_all_tasks() { std::vector init_task_ids = { // Init tasks task_id_t::BATCHNORM_INIT_TASK_ID, - task_id_t::COMBINE_INIT_TASK_ID, task_id_t::CONV2D_INIT_TASK_ID, task_id_t::DROPOUT_INIT_TASK_ID, task_id_t::ELEMENTBINARY_INIT_TASK_ID, @@ -47,8 +46,6 @@ Realm::Event register_all_tasks() { task_id_t::ATTENTION_INIT_TASK_ID, task_id_t::POOL2D_INIT_TASK_ID, task_id_t::REDUCE_INIT_TASK_ID, - task_id_t::REDUCTION_INIT_TASK_ID, - task_id_t::REPARTITION_INIT_TASK_ID, task_id_t::SOFTMAX_INIT_TASK_ID, }; @@ -68,7 +65,6 @@ Realm::Event register_all_tasks() { task_id_t::BATCHNORM_FWD_TASK_ID, task_id_t::BROADCAST_FWD_TASK_ID, task_id_t::CAST_FWD_TASK_ID, - task_id_t::COMBINE_FWD_TASK_ID, task_id_t::CONCAT_FWD_TASK_ID, task_id_t::CONV2D_FWD_TASK_ID, task_id_t::DROPOUT_FWD_TASK_ID, @@ -82,8 +78,6 @@ Realm::Event register_all_tasks() { task_id_t::ATTENTION_FWD_TASK_ID, task_id_t::POOL2D_FWD_TASK_ID, task_id_t::REDUCE_FWD_TASK_ID, - task_id_t::REDUCTION_FWD_TASK_ID, - task_id_t::REPARTITION_FWD_TASK_ID, task_id_t::RESHAPE_FWD_TASK_ID, task_id_t::REVERSE_FWD_TASK_ID, task_id_t::SOFTMAX_FWD_TASK_ID, @@ -95,7 +89,6 @@ Realm::Event register_all_tasks() { task_id_t::BATCHNORM_BWD_TASK_ID, task_id_t::BROADCAST_BWD_TASK_ID, task_id_t::CAST_BWD_TASK_ID, - task_id_t::COMBINE_BWD_TASK_ID, task_id_t::CONCAT_BWD_TASK_ID, task_id_t::CONV2D_BWD_TASK_ID, task_id_t::DROPOUT_BWD_TASK_ID, @@ -109,8 +102,6 @@ Realm::Event register_all_tasks() { task_id_t::ATTENTION_BWD_TASK_ID, task_id_t::POOL2D_BWD_TASK_ID, task_id_t::REDUCE_BWD_TASK_ID, - task_id_t::REDUCTION_BWD_TASK_ID, - task_id_t::REPARTITION_BWD_TASK_ID, task_id_t::RESHAPE_BWD_TASK_ID, task_id_t::REVERSE_BWD_TASK_ID, task_id_t::SOFTMAX_BWD_TASK_ID, diff --git a/lib/realm-execution/src/realm-execution/tasks/task_id_t.cc b/lib/realm-execution/src/realm-execution/tasks/task_id_t.cc index bce32a5d81..25daa6e0e3 100644 --- a/lib/realm-execution/src/realm-execution/tasks/task_id_t.cc +++ b/lib/realm-execution/src/realm-execution/tasks/task_id_t.cc @@ -35,7 +35,7 @@ std::optional [](BatchNormAttrs const &) { return task_id_t::BATCHNORM_INIT_TASK_ID; }, [](BroadcastAttrs const &) { return std::nullopt; }, [](CastAttrs const &) { return std::nullopt; }, - [](CombineAttrs const &attrs) { return task_id_t::COMBINE_INIT_TASK_ID; }, + [](CombineAttrs const &) { return std::nullopt; }, [](ConcatAttrs const &) { return std::nullopt; }, [](Conv2DAttrs const &) { return task_id_t::CONV2D_INIT_TASK_ID; }, [](DropoutAttrs const &) { return task_id_t::DROPOUT_INIT_TASK_ID; }, @@ -57,13 +57,9 @@ std::optional [](NoopAttrs const &) { return std::nullopt; }, [](Pool2DAttrs const &) { return task_id_t::POOL2D_INIT_TASK_ID; }, [](ReduceAttrs const &) { return task_id_t::REDUCE_INIT_TASK_ID; }, - [](ReductionAttrs const &attrs) { - return task_id_t::REDUCTION_INIT_TASK_ID; - }, - [](RepartitionAttrs const &attrs) { - return task_id_t::REPARTITION_INIT_TASK_ID; - }, - [](ReplicateAttrs const &attrs) { return std::nullopt; }, + [](ReductionAttrs const &) { return std::nullopt; }, + [](RepartitionAttrs const &) { return std::nullopt; }, + [](ReplicateAttrs const &) { return std::nullopt; }, [](ReshapeAttrs const &) { return std::nullopt; }, [](ReverseAttrs const &) { return std::nullopt; }, [](SoftmaxAttrs const &) { return task_id_t::SOFTMAX_INIT_TASK_ID; }, @@ -81,7 +77,7 @@ std::optional [](BatchNormAttrs const &) { return task_id_t::BATCHNORM_FWD_TASK_ID; }, [](BroadcastAttrs const &) { return task_id_t::BROADCAST_FWD_TASK_ID; }, [](CastAttrs const &) { return task_id_t::CAST_FWD_TASK_ID; }, - [](CombineAttrs const &attrs) { return task_id_t::COMBINE_FWD_TASK_ID; }, + [](CombineAttrs const &) { return std::nullopt; }, [](ConcatAttrs const &) { return task_id_t::CONCAT_FWD_TASK_ID; }, [](Conv2DAttrs const &) { return task_id_t::CONV2D_FWD_TASK_ID; }, [](DropoutAttrs const &) { return task_id_t::DROPOUT_FWD_TASK_ID; }, @@ -103,13 +99,9 @@ std::optional [](NoopAttrs const &) { return std::nullopt; }, [](Pool2DAttrs const &) { return task_id_t::POOL2D_FWD_TASK_ID; }, [](ReduceAttrs const &) { return task_id_t::REDUCE_FWD_TASK_ID; }, - [](ReductionAttrs const &attrs) { - return task_id_t::REDUCTION_FWD_TASK_ID; - }, - [](RepartitionAttrs const &attrs) { - return task_id_t::REPARTITION_FWD_TASK_ID; - }, - [](ReplicateAttrs const &attrs) { return std::nullopt; }, + [](ReductionAttrs const &) { return std::nullopt; }, + [](RepartitionAttrs const &) { return std::nullopt; }, + [](ReplicateAttrs const &) { return std::nullopt; }, [](ReshapeAttrs const &) { return task_id_t::RESHAPE_FWD_TASK_ID; }, [](ReverseAttrs const &) { return task_id_t::REVERSE_FWD_TASK_ID; }, [](SoftmaxAttrs const &) { return task_id_t::SOFTMAX_FWD_TASK_ID; }, @@ -127,7 +119,7 @@ std::optional [](BatchNormAttrs const &) { return task_id_t::BATCHNORM_BWD_TASK_ID; }, [](BroadcastAttrs const &) { return task_id_t::BROADCAST_BWD_TASK_ID; }, [](CastAttrs const &) { return task_id_t::CAST_BWD_TASK_ID; }, - [](CombineAttrs const &attrs) { return task_id_t::COMBINE_BWD_TASK_ID; }, + [](CombineAttrs const &) { return std::nullopt; }, [](ConcatAttrs const &) { return task_id_t::CONCAT_BWD_TASK_ID; }, [](Conv2DAttrs const &) { return task_id_t::CONV2D_BWD_TASK_ID; }, [](DropoutAttrs const &) { return task_id_t::DROPOUT_BWD_TASK_ID; }, @@ -149,13 +141,9 @@ std::optional [](NoopAttrs const &) { return std::nullopt; }, [](Pool2DAttrs const &) { return task_id_t::POOL2D_BWD_TASK_ID; }, [](ReduceAttrs const &) { return task_id_t::REDUCE_BWD_TASK_ID; }, - [](ReductionAttrs const &attrs) { - return task_id_t::REDUCTION_BWD_TASK_ID; - }, - [](RepartitionAttrs const &attrs) { - return task_id_t::REPARTITION_BWD_TASK_ID; - }, - [](ReplicateAttrs const &attrs) { return std::nullopt; }, + [](ReductionAttrs const &) { return std::nullopt; }, + [](RepartitionAttrs const &) { return std::nullopt; }, + [](ReplicateAttrs const &) { return std::nullopt; }, [](ReshapeAttrs const &) { return task_id_t::RESHAPE_BWD_TASK_ID; }, [](ReverseAttrs const &) { return task_id_t::REVERSE_BWD_TASK_ID; }, [](SoftmaxAttrs const &) { return task_id_t::SOFTMAX_BWD_TASK_ID; }, diff --git a/lib/task-spec/include/task-spec/dynamic_graph/parallel_op_data_movement.h b/lib/task-spec/include/task-spec/dynamic_graph/parallel_op_data_movement.h new file mode 100644 index 0000000000..ffe8bb6002 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/parallel_op_data_movement.h @@ -0,0 +1,36 @@ +#ifndef _FLEXFLOW_PARALLEL_OP_DATA_MOVEMENT_H +#define _FLEXFLOW_PARALLEL_OP_DATA_MOVEMENT_H +#include "op-attrs/operator_type.dtg.h" +#include "task-spec/dynamic_graph/dynamic_task_type.dtg.h" +#include "task-spec/dynamic_graph/training_operation_attrs.dtg.h" +#include + +namespace FlexFlow { + +enum class ParallelOpMovementKind { + BROADCAST, // 1 → N (copy) + GATHER, // N → 1 (copy) + SUM_REDUCE, // N → 1 (sum) + RESHUFFLE, // N → N (copy) +}; + +// Returns the data movement kind for a parallel op in a given pass direction. +// This is the single source of truth that all three pipeline stages +// (pass_expansion, shard_expansion, pcg_instance) derive their behavior from. +ParallelOpMovementKind get_parallel_op_movement_kind(OperatorType op_type, + DynamicTaskType task_type); + +// Returns true if this op type is a parallel (data-movement-only) op. +bool is_parallel_op(OperatorType op_type); + +// Returns the movement kind if op_attrs is a parallel op, nullopt otherwise. +// Avoids requiring pcg_operator_attrs.h at call sites. +std::optional + get_parallel_op_movement_kind_for_training_op( + TrainingOperationAttrs const &op_attrs, DynamicTaskType task_type); + +// Returns true if the TrainingOperationAttrs represents a parallel op. +bool is_parallel_training_op(TrainingOperationAttrs const &op_attrs); +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index ef5695792f..b9cffc2370 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc @@ -3,6 +3,7 @@ #include "op-attrs/tensor_slot_name.dtg.h" #include "pcg/machine_space_coordinate.dtg.h" #include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h" +#include "pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.h" #include "task-spec/dynamic_graph/copy_insertion.h" #include "task-spec/dynamic_graph/dynamic_node_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h" @@ -14,6 +15,7 @@ #include "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.h" +#include "task-spec/dynamic_graph/parallel_op_data_movement.h" #include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" #include "utils/bidict/algorithms/bidict_from_unstructured_relation.h" @@ -77,6 +79,26 @@ void require_value_is_copy_inserted(DynamicValueAttrs const &v) { ASSERT(v.mapping.has_value()); } +// Returns true if the given slot name maps to duplicate coords across shard +// bindings (i.e. cannot form a valid bidict). +static bool slot_name_has_dup_coords(DynamicNodeInvocation const &invocation, + TensorSlotName const &slot_name) { + DynamicNodeMapping const &nm = + assert_unwrap(invocation.node_attrs.mapping); + bidict const + shard_bindings = dynamic_node_mapping_get_shard_bindings(nm); + std::set seen; + for (global_device_id_t const &dev : shard_bindings.left_values()) { + ParallelTensorSpaceCoordinate const c = + ptensor_space_coord_for_slot_name(shard_bindings.at_l(dev), slot_name); + if (contains(seen, c)) { + return true; + } + seen.insert(c); + } + return false; +} + void require_invocation_is_fully_copy_inserted(DynamicNodeInvocation const &i) { auto require_node_is_copy_inserted = [](DynamicNodeAttrs const &) { return; }; @@ -272,44 +294,97 @@ std::map require_graph_is_ready_for_copy_insertion(g); - auto slots_to_map_for_replicate = - [](dynamic_invocation_id_t const &invocation_id, - DynamicNodeInvocation const &invocation) + auto slots_to_map_for_parallel_op = + [&](dynamic_invocation_id_t const &invocation_id, + DynamicNodeInvocation const &invocation, + ParallelOpMovementKind const kind) -> std::set { - TrainingOpType op_type = dynamic_node_invocation_get_op_type(invocation); - - ASSERT(op_type == TrainingOpType{OperatorType::REPLICATE}); - - std::set slot_sites = - (invocation.node_attrs.task_type == DynamicTaskType::BWD) - ? get_incoming_dynamic_slot_sites_for_invocation(invocation_id, - invocation) - : get_output_dynamic_slot_sites_for_invocation(invocation_id, - invocation); - - { - InternalDynamicSlotSite slot_site = get_only(slot_sites); - ASSERT(slot_site.slot_name.slot_name == TensorSlotName::OUTPUT); + auto output_slot_has_dup_coords = + [&](DynamicNodeInvocation const &inv) -> bool { + TensorSlotName const sn = get_only(keys(inv.outputs)).slot_name; + return slot_name_has_dup_coords(inv, sn); }; - return slot_sites; + switch (kind) { + case ParallelOpMovementKind::BROADCAST: + return get_output_dynamic_slot_sites_for_invocation(invocation_id, + invocation); + case ParallelOpMovementKind::GATHER: + case ParallelOpMovementKind::SUM_REDUCE: + return get_incoming_dynamic_slot_sites_for_invocation(invocation_id, + invocation); + case ParallelOpMovementKind::RESHUFFLE: { + if (!output_slot_has_dup_coords(invocation)) { + return get_output_dynamic_slot_sites_for_invocation(invocation_id, + invocation); + } else { + return get_incoming_dynamic_slot_sites_for_invocation(invocation_id, + invocation); + } + } + default: + PANIC("Unhandled ParallelOpMovementKind", kind); + } }; auto get_mappings_for_invocation = [&](DynamicNodeInvocation const &invocation) -> std::map { - TrainingOpType op_type = dynamic_node_invocation_get_op_type(invocation); - dynamic_invocation_id_t invocation_id = + dynamic_invocation_id_t const invocation_id = dynamic_graph_get_id_for_invocation(g, invocation); - std::set slot_sites_to_resolve = [&] { - if (op_type == TrainingOpType{OperatorType::REPLICATE}) { - return slots_to_map_for_replicate(invocation_id, invocation); + TrainingOperationAttrs const op_attrs = + assert_unwrap(invocation.node_attrs.op_attrs); + + std::optional const kind_opt = [&]() { + if (!is_parallel_training_op(op_attrs)) { + return std::optional{std::nullopt}; + } + DynamicTaskType const task_type = + invocation.node_attrs.task_type.value_or(DynamicTaskType::FWD); + return get_parallel_op_movement_kind_for_training_op(op_attrs, task_type); + }(); + + std::set const slot_sites_to_resolve = [&]() { + if (kind_opt.has_value()) { + return slots_to_map_for_parallel_op( + invocation_id, invocation, kind_opt.value()); } else { return get_dynamic_slot_sites_for_invocation(invocation_id, invocation); } }(); + // For RESHUFFLE, the invocation output slot may have a different + // TensorSlotName than the unique side of the node mapping (swapped in BWD). + // Find the unique slot name dynamically by checking coord uniqueness. + if (kind_opt.has_value() && + kind_opt.value() == ParallelOpMovementKind::RESHUFFLE) { + DynamicNodeMapping const &nm = + assert_unwrap(invocation.node_attrs.mapping); + + TensorSlotName const out_slot_name = + get_only(keys(invocation.outputs)).slot_name; + + TensorSlotName const unique_slot_name = [&]() { + if (!slot_name_has_dup_coords(invocation, out_slot_name)) { + return out_slot_name; + } + TensorSlotName const in_slot_name = + get_only(keys(invocation.inputs)).slot_name; + ASSERT(!slot_name_has_dup_coords(invocation, in_slot_name), + "Neither slot name has unique coords in RESHUFFLE node mapping"); + return in_slot_name; + }(); + + ParallelTensorMapping const unique_mapping{ + dynamic_node_mapping_bindings_for_slot_name(nm, unique_slot_name)}; + + return generate_map( + slot_sites_to_resolve, + [&](InternalDynamicSlotSite const &) -> ParallelTensorMapping { + return unique_mapping; + }); + } return generate_map( slot_sites_to_resolve, [&](InternalDynamicSlotSite const &s) -> ParallelTensorMapping { @@ -335,45 +410,109 @@ std::map &resolved_mappings) { require_graph_is_ready_for_copy_insertion(g); - std::set all_internal_slot_sites = + std::set const all_internal_slot_sites = get_internal_dynamic_slot_sites(g); - std::set missing_mappings = + std::set const missing_mappings = set_minus(all_internal_slot_sites, keys(resolved_mappings)); auto get_mapping_for_slot_site_from_adjacent_values = [&](InternalDynamicSlotSite const &slot_site) -> ParallelTensorMapping { - DynamicNodeInvocation invocation = + DynamicNodeInvocation const invocation = dynamic_graph_get_invocation_for_id(g, slot_site.invocation_id); - TrainingOpType op_type = dynamic_node_invocation_get_op_type(invocation); - std::optional task_type = invocation.node_attrs.task_type; - - TrainingOpType replicate_op_type = TrainingOpType{OperatorType::REPLICATE}; - - if (op_type == replicate_op_type && task_type == DynamicTaskType::BWD) { - ASSERT(slot_site.direction == TensorDirection::OUTPUT); - - InternalDynamicSlotSite slot_site_sink = - get_only(dynamic_graph_find_sinks_of_slot_site(g, slot_site)); - - ASSERT(contains_key(resolved_mappings, slot_site_sink)); - - return resolved_mappings.at(slot_site_sink); - } else if (op_type == replicate_op_type && - (task_type == std::nullopt || - task_type == DynamicTaskType::FWD)) { - ASSERT(slot_site.direction == TensorDirection::INCOMING); - - InternalDynamicSlotSite slot_site_src = - dynamic_graph_find_source_of_slot_site(g, slot_site) - .require_internal(); - - ASSERT(contains_key(resolved_mappings, slot_site_src)); + TrainingOperationAttrs const op_attrs = + assert_unwrap(invocation.node_attrs.op_attrs); + DynamicTaskType const task_type = + invocation.node_attrs.task_type.value_or(DynamicTaskType::FWD); + + std::optional const kind_opt = + get_parallel_op_movement_kind_for_training_op(op_attrs, task_type); + + ASSERT(kind_opt.has_value(), + "Only parallel ops should have missing mappings"); + ParallelOpMovementKind const kind = kind_opt.value(); + + auto find_fwd_sink = + [&](InternalDynamicSlotSite const &s) -> InternalDynamicSlotSite { + std::set const sinks = + dynamic_graph_find_sinks_of_slot_site(g, s); + for (InternalDynamicSlotSite const &sink : sinks) { + DynamicNodeInvocation const sink_invocation = + dynamic_graph_get_invocation_for_id(g, sink.invocation_id); + std::optional const task_type = + sink_invocation.node_attrs.task_type; + // Accept FWD, UPD, LOSS, or pre-pass-expansion (nullopt) sinks. LOSS + // is included because a parallel op's FWD output may feed directly + // into the loss (e.g. it is the model's terminal output), in which + // case the loss node is the only non-BWD sink. + if (!task_type.has_value() || + task_type.value() == DynamicTaskType::FWD || + task_type.value() == DynamicTaskType::UPD || + task_type.value() == DynamicTaskType::LOSS) { + return sink; + } + } + PANIC("No non-BWD sink found for parallel op missing mapping resolution"); + }; - return resolved_mappings.at(slot_site_src); - } else { - PANIC("Unhandled case"); + switch (kind) { + case ParallelOpMovementKind::BROADCAST: { + if (slot_site.direction == TensorDirection::INCOMING) { + // FWD: missing INPUT (1-side) — resolve from source + InternalDynamicSlotSite const slot_site_src = + dynamic_graph_find_source_of_slot_site(g, slot_site) + .require_internal(); + ASSERT(contains_key(resolved_mappings, slot_site_src)); + return resolved_mappings.at(slot_site_src); + } else { + // BWD: missing OUTPUT (N-side = INPUT grad) — resolve from FWD sink + // Multiple sinks may exist if BWD also uses this value as activation + ASSERT(slot_site.direction == TensorDirection::OUTPUT); + InternalDynamicSlotSite const slot_site_sink = + find_fwd_sink(slot_site); + ASSERT(contains_key(resolved_mappings, slot_site_sink)); + return resolved_mappings.at(slot_site_sink); + } + } + case ParallelOpMovementKind::GATHER: + case ParallelOpMovementKind::SUM_REDUCE: { + if (slot_site.direction == TensorDirection::OUTPUT) { + // FWD: missing OUTPUT (1-side) — resolve from FWD sink + // Multiple sinks may exist if BWD also uses this value as activation + InternalDynamicSlotSite const slot_site_sink = + find_fwd_sink(slot_site); + ASSERT(contains_key(resolved_mappings, slot_site_sink)); + return resolved_mappings.at(slot_site_sink); + } else { + // BWD: missing INPUT (N-side = OUTPUT grad) — resolve from source + ASSERT(slot_site.direction == TensorDirection::INCOMING); + InternalDynamicSlotSite const slot_site_src = + dynamic_graph_find_source_of_slot_site(g, slot_site) + .require_internal(); + ASSERT(contains_key(resolved_mappings, slot_site_src)); + return resolved_mappings.at(slot_site_src); + } + } + case ParallelOpMovementKind::RESHUFFLE: { + if (slot_site.direction == TensorDirection::INCOMING) { + // Normal RESHUFFLE: missing INPUT — resolve from source + InternalDynamicSlotSite const slot_site_src = + dynamic_graph_find_source_of_slot_site(g, slot_site) + .require_internal(); + ASSERT(contains_key(resolved_mappings, slot_site_src)); + return resolved_mappings.at(slot_site_src); + } else { + // Scatter BWD: missing OUTPUT — resolve from FWD/UPD sink + ASSERT(slot_site.direction == TensorDirection::OUTPUT); + InternalDynamicSlotSite const slot_site_sink = + find_fwd_sink(slot_site); + ASSERT(contains_key(resolved_mappings, slot_site_sink)); + return resolved_mappings.at(slot_site_sink); + } + } + default: + PANIC("Unhandled ParallelOpMovementKind", kind); } }; diff --git a/lib/task-spec/src/task-spec/dynamic_graph/parallel_op_data_movement.cc b/lib/task-spec/src/task-spec/dynamic_graph/parallel_op_data_movement.cc new file mode 100644 index 0000000000..8dd0a1b33a --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/parallel_op_data_movement.cc @@ -0,0 +1,81 @@ +#include "task-spec/dynamic_graph/parallel_op_data_movement.h" +#include "op-attrs/pcg_operator_attrs.h" +#include "utils/exception.h" + +namespace FlexFlow { + +bool is_parallel_op(OperatorType op_type) { + switch (op_type) { + case OperatorType::REPLICATE: + case OperatorType::COMBINE: + case OperatorType::REDUCTION: + case OperatorType::REPARTITION: + return true; + default: + return false; + } +} + +ParallelOpMovementKind + get_parallel_op_movement_kind(OperatorType const op_type, + DynamicTaskType const task_type) { + ASSERT(is_parallel_op(op_type)); + switch (op_type) { + case OperatorType::REPLICATE: + switch (task_type) { + case DynamicTaskType::FWD: + return ParallelOpMovementKind::BROADCAST; + case DynamicTaskType::BWD: + return ParallelOpMovementKind::SUM_REDUCE; + default: + PANIC("Unexpected task type for REPLICATE", task_type); + } + case OperatorType::COMBINE: + switch (task_type) { + case DynamicTaskType::FWD: + return ParallelOpMovementKind::GATHER; + case DynamicTaskType::BWD: + return ParallelOpMovementKind::BROADCAST; + default: + PANIC("Unexpected task type for COMBINE", task_type); + } + case OperatorType::REDUCTION: + switch (task_type) { + case DynamicTaskType::FWD: + return ParallelOpMovementKind::SUM_REDUCE; + case DynamicTaskType::BWD: + return ParallelOpMovementKind::BROADCAST; + default: + PANIC("Unexpected task type for REDUCTION", task_type); + } + case OperatorType::REPARTITION: + return ParallelOpMovementKind::RESHUFFLE; + default: + PANIC("Not a parallel op", op_type); + } +} + +static std::optional + try_get_pcg_op_type(TrainingOperationAttrs const &op_attrs) { + if (!op_attrs.is_pcg_op()) { + return std::nullopt; + } + return pcg_op_attrs_get_op_type(op_attrs.require_pcg_op()); +} + +std::optional + get_parallel_op_movement_kind_for_training_op( + TrainingOperationAttrs const &op_attrs, + DynamicTaskType const task_type) { + std::optional const op_type = try_get_pcg_op_type(op_attrs); + if (!op_type.has_value() || !is_parallel_op(op_type.value())) { + return std::nullopt; + } + return get_parallel_op_movement_kind(op_type.value(), task_type); +} + +bool is_parallel_training_op(TrainingOperationAttrs const &op_attrs) { + std::optional const op_type = try_get_pcg_op_type(op_attrs); + return op_type.has_value() && is_parallel_op(op_type.value()); +} +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc index a095880118..28996e5753 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc @@ -2,6 +2,7 @@ #include "task-spec/dynamic_graph/dynamic_node_invocation.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" #include "task-spec/dynamic_graph/dynamic_tensor_role.h" +#include "task-spec/dynamic_graph/parallel_op_data_movement.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" #include "utils/containers/are_all_same.h" #include "utils/containers/flatmap.h" @@ -290,8 +291,9 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( /*node_attrs=*/invocation.node_attrs, /*outputs=*/map_values(invocation.inputs, to_grad_value), }; - } else if (training_op_attrs_has_op_type(op_attrs, - OperatorType::REPLICATE)) { + } else if (is_parallel_training_op(op_attrs)) { + // All parallel ops: BWD is purely gradient data movement. + // No fwd activations are needed — just swap input/output grad roles. return DynamicNodeInvocation{ /*inputs=*/{ transform(invocation.outputs, to_grad), diff --git a/lib/task-spec/src/task-spec/dynamic_graph/shard_expansion.cc b/lib/task-spec/src/task-spec/dynamic_graph/shard_expansion.cc index 98df030438..bc384cdcf6 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/shard_expansion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/shard_expansion.cc @@ -5,11 +5,11 @@ #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" #include "task-spec/dynamic_graph/dynamic_tensor_role.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" +#include "task-spec/dynamic_graph/parallel_op_data_movement.h" #include "task-spec/dynamic_graph/parallel_tensor_mapping.h" #include "task-spec/dynamic_graph/serializable_dynamic_node_invocation.h" #include "task-spec/dynamic_graph/shard_expansion.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" -#include "utils/bidict/algorithms/bidict_filter_keys.h" #include "utils/bidict/algorithms/bidict_filter_values.h" #include "utils/binary_relation/binary_relation_from_map.h" #include "utils/binary_relation/binary_relation_transform_right2.h" @@ -21,7 +21,6 @@ #include "utils/containers/map_from_pairs.h" #include "utils/containers/map_from_unordered.h" #include "utils/containers/map_keys.h" -#include "utils/containers/map_values2.h" #include "utils/containers/merge_disjoint_maps.h" #include "utils/containers/require_only_key.h" #include "utils/containers/require_same.h" @@ -127,55 +126,6 @@ static DynamicNodeInvocationShardingInfo invocation_sharding_info_for_binding( return result; } -static bidict - restrict_tensor_mapping_keys_to_coord( - bidict const - &mapping, - ParallelTensorSpaceCoordinate const ¶llel_tensor_coord) { - return bidict_filter_keys(mapping, - [&](ParallelTensorSpaceCoordinate const &p) { - return p == parallel_tensor_coord; - }); -} - -static DynamicNodeInvocation shard_invocation_for_binding( - DynamicNodeInvocation const &i, - global_device_id_t const &device_id, - OperatorAtomicTaskShardBinding const &binding) { - - auto shard_expand_value_attrs = - [&](DynamicTensorSlot const &s, - DynamicValueAttrs const &v) -> DynamicValueAttrs { - ParallelTensorSpaceCoordinate parallel_tensor_coord = - binding.tensor_coords.at(s.slot_name); - - DynamicValueAttrs result = v; - result.shard_coord = parallel_tensor_coord; - result.mapping = transform( - v.mapping, - [&](ParallelTensorMapping const &mapping) -> ParallelTensorMapping { - return ParallelTensorMapping{ - restrict_tensor_mapping_keys_to_coord(mapping.raw, - parallel_tensor_coord), - }; - }); - return result; - }; - - DynamicNodeAttrs expanded_node_attrs = [&]() { - DynamicNodeAttrs result = i.node_attrs; - result.device_ids = nonempty_set{device_id}; - ; - return result; - }(); - - return DynamicNodeInvocation{ - /*inputs=*/map_values2(i.inputs, shard_expand_value_attrs), - /*node_attrs=*/expanded_node_attrs, - /*outputs=*/map_values2(i.outputs, shard_expand_value_attrs), - }; -} - static std::set generate_shard_expansion_for_copy(DynamicNodeInvocation const &i) { auto [input_slot, input] = get_only(i.inputs); @@ -212,226 +162,238 @@ static std::set }); } -// TODO(@lockshaw): There is a lot of code duplication between -// generate_shard_expansion_for_fwd_replicate and -// generate_shard_expansion_for_bwd_replicate that should eventually be -// factored out. static std::set - generate_shard_expansion_for_fwd_replicate(DynamicNodeInvocation const &i) { - ASSERT(i.node_attrs.task_type == DynamicTaskType::FWD); + generate_shard_expansion_using_node_mapping( + DynamicNodeInvocation const &i, + DynamicTensorSlot const &one_side_slot, + DynamicTensorSlot const &n_side_slot, + ParallelTensorMapping const &one_side_mapping, + ParallelTensorMapping const &n_side_mapping) { - DynamicNodeMapping node_mapping = assert_unwrap(i.node_attrs.mapping); + DynamicNodeMapping const node_mapping = assert_unwrap(i.node_attrs.mapping); + bidict const + shard_bindings = dynamic_node_mapping_get_shard_bindings(node_mapping); - DynamicTensorSlot expected_input_slot = DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::INPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), - /*task_shard=*/std::nullopt, - }; - - DynamicValueAttrs input = require_only_key(i.inputs, expected_input_slot); - - DynamicTensorSlot expected_output_slot = DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), - /*task_shard=*/std::nullopt, - }; - - DynamicValueAttrs output = require_only_key(i.outputs, expected_output_slot); - - ParallelTensorMapping input_value_mapping = assert_unwrap(input.mapping); - - std::set input_tensor_shards = - pt_mapping_get_coord_set(input_value_mapping); - - ParallelTensorMapping output_value_mapping = assert_unwrap(output.mapping); - - auto get_task_shard_device_ids_for_input_tensor_shard = - [&](ParallelTensorSpaceCoordinate const &input_tensor_shard) - -> nonempty_set { - bidict - dependent_on_input_tensor_shard = bidict_filter_values( - dynamic_node_mapping_get_shard_bindings(node_mapping), - [&](OperatorAtomicTaskShardBinding const &b) -> bool { - return ptensor_space_coord_for_slot_name( - b, TensorSlotName::INPUT) == input_tensor_shard; - }); + std::set const one_side_coords = + pt_mapping_get_coord_set(one_side_mapping); - return nonempty_set(dependent_on_input_tensor_shard.left_values()); - }; - - auto invocation_sharding_info_for_input_tensor_shard = - [&](ParallelTensorSpaceCoordinate const &c) - -> DynamicNodeInvocationShardingInfo { - nonempty_set task_shard_device_ids = - get_task_shard_device_ids_for_input_tensor_shard(c); - - std::map - output_sharding_infos = - generate_map(task_shard_device_ids.unwrap_as_set(), - [&](global_device_id_t const &device_id) - -> DynamicValueAttrsShardingInfo { - ParallelTensorSpaceCoordinate pc = - pt_mapping_get_coord_for_device( - output_value_mapping, device_id); - - return DynamicValueAttrsShardingInfo{ - /*shard_coord=*/pc, - /*mapping=*/device_id, - }; - }); - - std::map - keyed_output_sharding_infos = map_keys( - output_sharding_infos, - [&](global_device_id_t const &device_id) -> DynamicTensorSlot { - return DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), - /*task_shard=*/device_id.coord, - }; - }); - - DynamicTensorSlot input_slot = DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::INPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), - /*task_shard=*/std::nullopt, - }; - - DynamicValueAttrsShardingInfo input_sharding_info = - DynamicValueAttrsShardingInfo{ - /*shard_coord=*/c, - /*mapping=*/pt_mapping_get_device_for_coord(input_value_mapping, c), - }; - - std::map sharding_infos = - binary_merge_disjoint_maps( - keyed_output_sharding_infos, - std::map{ + return transform( + one_side_coords, + [&](ParallelTensorSpaceCoordinate const &one_coord) + -> DynamicNodeInvocationShardingInfo { + // Find all devices whose shard binding has this one-side coord + bidict const + for_one_coord = bidict_filter_values( + shard_bindings, + [&](OperatorAtomicTaskShardBinding const &b) -> bool { + return ptensor_space_coord_for_slot_name( + b, one_side_slot.slot_name) == one_coord; + }); + + nonempty_set const group_device_ids = + nonempty_set(for_one_coord.left_values()); + + // Build n-side sharding infos — one per device in the group + std::map + n_side_sharding_infos = map_from_pairs(transform( + group_device_ids.unwrap_as_set(), + [&](global_device_id_t const &device) + -> std::pair { + ParallelTensorSpaceCoordinate const n_coord = + ptensor_space_coord_for_slot_name( + shard_bindings.at_l(device), n_side_slot.slot_name); + return { + DynamicTensorSlot{ + /*slot_name=*/n_side_slot.slot_name, + /*slot_tensor_role=*/n_side_slot.slot_tensor_role, + /*task_shard=*/device.coord, + }, + DynamicValueAttrsShardingInfo{ + /*shard_coord=*/n_coord, + /*mapping=*/device, + }, + }; + })); + + // Build one-side sharding info — single entry, task_shard=nullopt + std::map const + one_side_sharding_infos = { { - input_slot, - input_sharding_info, + DynamicTensorSlot{ + /*slot_name=*/one_side_slot.slot_name, + /*slot_tensor_role=*/one_side_slot.slot_tensor_role, + /*task_shard=*/std::nullopt, + }, + DynamicValueAttrsShardingInfo{ + /*shard_coord=*/one_coord, + /*mapping=*/ + pt_mapping_get_device_for_coord(one_side_mapping, + one_coord), + }, }, - }); - - return DynamicNodeInvocationShardingInfo{ - /*device_ids=*/task_shard_device_ids, - /*value_sharding=*/binary_relation_from_map(sharding_infos), - }; - }; + }; - return transform(input_tensor_shards, - invocation_sharding_info_for_input_tensor_shard); + return DynamicNodeInvocationShardingInfo{ + /*device_ids=*/group_device_ids, + /*value_sharding=*/ + binary_relation_from_map(binary_merge_disjoint_maps( + one_side_sharding_infos, n_side_sharding_infos)), + }; + }); } static std::set - generate_shard_expansion_for_bwd_replicate(DynamicNodeInvocation const &i) { - ASSERT(i.node_attrs.task_type == DynamicTaskType::BWD); - - DynamicNodeMapping node_mapping = assert_unwrap(i.node_attrs.mapping); - - DynamicTensorSlot expected_output_grad_slot = DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_bwd(), - /*task_shard=*/std::nullopt, - }; - - DynamicValueAttrs output_grad = - require_only_key(i.inputs, expected_output_grad_slot); + generate_shard_expansion_for_parallel_op( + DynamicNodeInvocation const &i, ParallelOpMovementKind const kind) { + + DynamicTensorSlot const input_slot = get_only(keys(i.inputs)); + DynamicTensorSlot const output_slot = get_only(keys(i.outputs)); + DynamicValueAttrs const &input = i.inputs.at(input_slot); + DynamicValueAttrs const &output = i.outputs.at(output_slot); + + ParallelTensorMapping const input_mapping = assert_unwrap(input.mapping); + ParallelTensorMapping const output_mapping = assert_unwrap(output.mapping); + + std::set const input_coords = + pt_mapping_get_coord_set(input_mapping); + std::set const output_coords = + pt_mapping_get_coord_set(output_mapping); + + if (kind == ParallelOpMovementKind::BROADCAST) { + // 1 input → N outputs: one-side is input, n-side is output + return generate_shard_expansion_using_node_mapping( + i, input_slot, output_slot, input_mapping, output_mapping); + } - DynamicTensorSlot expected_input_grad_slot = DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::INPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_bwd(), - /*task_shard=*/std::nullopt, - }; + if (kind == ParallelOpMovementKind::SUM_REDUCE) { + // N inputs → 1 output: one-side is output, n-side is input + return generate_shard_expansion_using_node_mapping( + i, output_slot, input_slot, output_mapping, input_mapping); + } - DynamicValueAttrs input_grad = - require_only_key(i.outputs, expected_input_grad_slot); - - ParallelTensorMapping output_grad_value_mapping = - assert_unwrap(output_grad.mapping); - ParallelTensorMapping input_grad_value_mapping = - assert_unwrap(input_grad.mapping); - - std::set input_grad_tensor_shards = - pt_mapping_get_coord_set(input_grad_value_mapping); - - auto get_task_shard_device_ids_for_input_grad_tensor_shard = - [&](ParallelTensorSpaceCoordinate const &input_grad_tensor_shard) - -> nonempty_set { - bidict - produce_input_grad_tensor_shard = bidict_filter_values( - dynamic_node_mapping_get_shard_bindings(node_mapping), - [&](OperatorAtomicTaskShardBinding const &b) -> bool { - return ptensor_space_coord_for_slot_name( - b, TensorSlotName::INPUT) == input_grad_tensor_shard; - }); + if (kind == ParallelOpMovementKind::RESHUFFLE) { + DynamicNodeMapping const node_mapping = assert_unwrap(i.node_attrs.mapping); + bidict const + shard_bindings = dynamic_node_mapping_get_shard_bindings(node_mapping); + + return transform( + output_coords, + [&](ParallelTensorSpaceCoordinate const &out_coord) + -> DynamicNodeInvocationShardingInfo { + global_device_id_t const out_device = + pt_mapping_get_device_for_coord(output_mapping, out_coord); + + // Look up the input coord for this device directly by slot name, + // rather than scanning for a coord value that happens to appear in + // input_mapping — coordinate values can collide across the input + // and output coordinate spaces, which would otherwise risk pairing + // the wrong device. + OperatorAtomicTaskShardBinding const &binding = + shard_bindings.at_l(out_device); + ParallelTensorSpaceCoordinate const in_coord = + ptensor_space_coord_for_slot_name(binding, input_slot.slot_name); + + global_device_id_t const in_device = + pt_mapping_get_device_for_coord(input_mapping, in_coord); + + return DynamicNodeInvocationShardingInfo{ + /*device_ids=*/nonempty_set{out_device}, + /*value_sharding=*/ + binary_relation_from_map( + std::map{ + { + DynamicTensorSlot{ + /*slot_name=*/input_slot.slot_name, + /*slot_tensor_role=*/input_slot.slot_tensor_role, + /*task_shard=*/std::nullopt, + }, + DynamicValueAttrsShardingInfo{ + /*shard_coord=*/in_coord, + /*mapping=*/in_device, + }, + }, + { + DynamicTensorSlot{ + /*slot_name=*/output_slot.slot_name, + /*slot_tensor_role=*/output_slot.slot_tensor_role, + /*task_shard=*/std::nullopt, + }, + DynamicValueAttrsShardingInfo{ + /*shard_coord=*/out_coord, + /*mapping=*/out_device, + }, + }, + }), + }; + }); + } - return nonempty_set(produce_input_grad_tensor_shard.left_values()); + // GATHER: N inputs → 1 output (copy). + // Anchor on the single output coord; all input coords feed into it. + // input_coords are unique across shard bindings (resolved from node mapping). + ASSERT(kind == ParallelOpMovementKind::GATHER); + + auto make_sharding_infos = + [](std::set const &coords, + ParallelTensorMapping const &mapping, + DynamicTensorSlot const &base_slot) + -> std::map { + bool const needs_task_shard = coords.size() > 1; + return map_from_pairs(transform( + coords, + [&](ParallelTensorSpaceCoordinate const &coord) + -> std::pair { + global_device_id_t const device = + pt_mapping_get_device_for_coord(mapping, coord); + DynamicTensorSlot const slot = DynamicTensorSlot{ + /*slot_name=*/base_slot.slot_name, + /*slot_tensor_role=*/base_slot.slot_tensor_role, + /*task_shard=*/ + needs_task_shard ? std::optional{device.coord} : std::nullopt, + }; + return {slot, + DynamicValueAttrsShardingInfo{ + /*shard_coord=*/coord, + /*mapping=*/device, + }}; + })); }; - auto invocation_sharding_info_for_input_grad_tensor_shard = - [&](ParallelTensorSpaceCoordinate const &c) - -> DynamicNodeInvocationShardingInfo { - nonempty_set task_shard_device_ids = - get_task_shard_device_ids_for_input_grad_tensor_shard(c); - - std::map - output_grad_sharding_infos = - generate_map(task_shard_device_ids.unwrap_as_set(), - [&](global_device_id_t const &device_id) - -> DynamicValueAttrsShardingInfo { - ParallelTensorSpaceCoordinate pc = - pt_mapping_get_coord_for_device( - output_grad_value_mapping, device_id); - - return DynamicValueAttrsShardingInfo{ - /*shard_coord=*/pc, - /*mapping=*/device_id, - }; - }); - - std::map - keyed_output_grad_sharding_infos = map_keys( - output_grad_sharding_infos, - [&](global_device_id_t const &device_id) -> DynamicTensorSlot { - return DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_bwd(), - /*task_shard=*/device_id.coord, - }; + // One sharding info per output coord (there is only 1 for GATHER) + return transform( + output_coords, + [&](ParallelTensorSpaceCoordinate const &anchor) + -> DynamicNodeInvocationShardingInfo { + global_device_id_t const task_device = + pt_mapping_get_device_for_coord(output_mapping, anchor); + + std::map const + input_sharding_infos = + make_sharding_infos(input_coords, input_mapping, input_slot); + std::map const + output_sharding_infos = + make_sharding_infos({anchor}, output_mapping, output_slot); + + // device_ids must include every device that participates in this + // invocation (all input-side devices plus the output device), not + // just the output device — matching the SUM_REDUCE/BROADCAST cases, + // which likewise include the full N-side device group. This matters + // for consumers (e.g. machine slicing) that use device_ids to decide + // which devices must execute/observe an invocation. + std::set device_id_set = transform( + input_coords, [&](ParallelTensorSpaceCoordinate const &coord) { + return pt_mapping_get_device_for_coord(input_mapping, coord); }); + device_id_set.insert(task_device); - DynamicTensorSlot input_grad_slot = DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::INPUT, - /*slot_tensor_role=*/mk_dynamic_tensor_role_bwd(), - /*task_shard=*/std::nullopt, - }; - - DynamicValueAttrsShardingInfo input_grad_sharding_info = - DynamicValueAttrsShardingInfo{ - /*shard_coord=*/c, - /*mapping=*/ - pt_mapping_get_device_for_coord(input_grad_value_mapping, c), + return DynamicNodeInvocationShardingInfo{ + /*device_ids=*/nonempty_set(device_id_set), + /*value_sharding=*/ + binary_relation_from_map(binary_merge_disjoint_maps( + input_sharding_infos, output_sharding_infos)), }; - - std::map sharding_infos = - binary_merge_disjoint_maps( - keyed_output_grad_sharding_infos, - std::map{ - { - input_grad_slot, - input_grad_sharding_info, - }, - }); - - return DynamicNodeInvocationShardingInfo{ - /*device_ids=*/task_shard_device_ids, - /*value_sharding=*/binary_relation_from_map(sharding_infos), - }; - }; - - return transform(input_grad_tensor_shards, - invocation_sharding_info_for_input_grad_tensor_shard); + }); } std::set @@ -568,37 +530,35 @@ std::set generate_shard_expansion_for_invocation(DynamicNodeInvocation const &i) { require_invocation_is_ready_for_shard_expansion(i); - std::set result = [&]() { - if (i.node_attrs.op_attrs.value().is_copy()) { - return generate_shard_expansion_for_copy(i); - } else if (training_op_attrs_has_op_type(i.node_attrs.op_attrs.value(), - OperatorType::REPLICATE)) { - DynamicTaskType task_type = assert_unwrap(i.node_attrs.task_type); - switch (task_type) { - case DynamicTaskType::FWD: - return set_of(generate_shard_expansion_for_fwd_replicate(i)); - case DynamicTaskType::BWD: - return set_of(generate_shard_expansion_for_bwd_replicate(i)); - default: - PANIC("Unexpected task type for Replicate: {}", task_type); - } - } else { - DynamicNodeMapping mapping = assert_unwrap(i.node_attrs.mapping); + std::set const result = [&]() { + TrainingOperationAttrs const op_attrs = + assert_unwrap(i.node_attrs.op_attrs); - std::set shard_machine_coords = - target_devices_of_dynamic_node_mapping(mapping); - - return transform(shard_machine_coords, - [&](global_device_id_t const &device_id) - -> DynamicNodeInvocationShardingInfo { - OperatorAtomicTaskShardBinding slot_bindings = - dynamic_node_mapping_get_shard_binding_for_device( - mapping, device_id); + if (op_attrs.is_copy()) { + return generate_shard_expansion_for_copy(i); + } - return invocation_sharding_info_for_binding( - i, device_id, slot_bindings); - }); + if (is_parallel_training_op(op_attrs)) { + DynamicTaskType const task_type = assert_unwrap(i.node_attrs.task_type); + ParallelOpMovementKind const kind = + get_parallel_op_movement_kind_for_training_op(op_attrs, task_type) + .value(); + return generate_shard_expansion_for_parallel_op(i, kind); } + + DynamicNodeMapping const mapping = assert_unwrap(i.node_attrs.mapping); + std::set const shard_machine_coords = + target_devices_of_dynamic_node_mapping(mapping); + + return transform(shard_machine_coords, + [&](global_device_id_t const &device_id) + -> DynamicNodeInvocationShardingInfo { + OperatorAtomicTaskShardBinding const slot_bindings = + dynamic_node_mapping_get_shard_binding_for_device( + mapping, device_id); + return invocation_sharding_info_for_binding( + i, device_id, slot_bindings); + }); }(); return result; From 29869bce8ba4292404550de304ed15e8beedcb17 Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 13 Aug 2026 14:59:35 -0700 Subject: [PATCH 2/7] Register CUDA-visible reduction kernels for realm-execution redops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Renames redops/realm_redop_registry.cc to .cu (added to the CMake target as a CUDA-language source, mirroring lib/kernels), adds apply_cuda/fold_cuda kernel methods to each SumReduction, and rebuilds registration around create_reduction_op() + add_cuda_redop_kernels() + register_reduction. Two secondary issues had to be worked around to get the new .cu file to compile under nvcc at all: PRealm's prealm.h (realm-execution/realm.h's namespace Realm = ::PRealm alias, used everywhere else in this codebase) fails a static_assert when parsed by nvcc, unrelated to reductions — the .cu file avoids it entirely and uses ::Realm:: fully-qualified throughout, so register_all_redops() also dropped its Realm::Runtime parameter (fetches ::Realm::Runtime::get_runtime() itself instead). And the borrowed-from- Legion SumReduction/SumReduction non-exclusive apply/fold paths called __uint2bool/__bool2uint/__longlong_as_ulonglong/ __ulonglong_as_longlong, which aren't real CUDA builtins (Legion-internal helpers) — implemented locally. Confirmed on real hardware (2x Tesla P100): all four parallel-op GPU e2e tests pass. --- lib/realm-execution/CMakeLists.txt | 12 ++ .../redops/realm_redop_registry.h | 13 +- .../src/realm-execution/realm_manager.cc | 2 +- .../redops/realm_redop_registry.cu | 131 +++++++++++++++--- 4 files changed, 131 insertions(+), 27 deletions(-) diff --git a/lib/realm-execution/CMakeLists.txt b/lib/realm-execution/CMakeLists.txt index 49fbcfa4e0..378f28c8e4 100644 --- a/lib/realm-execution/CMakeLists.txt +++ b/lib/realm-execution/CMakeLists.txt @@ -1,8 +1,14 @@ +# realm_redop_registry.cu needs to be compiled by nvcc (not g++) so that its +# GPU reduction kernels (SumReduction::apply_cuda/fold_cuda) actually get +# generated; see the comment on register_sum_redop for why that matters. +project(realm-execution LANGUAGES CXX CUDA) + ff_add_library( NAME realm-execution SRC_PATTERNS src/*.cc + src/*.cu PUBLIC_INCLUDE include/ PRIVATE_INCLUDE @@ -19,4 +25,10 @@ ff_add_library( deps::realm ) +set_target_properties( + realm-execution + PROPERTIES + CUDA_STANDARD 17 +) + add_subdirectory(test) diff --git a/lib/realm-execution/include/realm-execution/redops/realm_redop_registry.h b/lib/realm-execution/include/realm-execution/redops/realm_redop_registry.h index e7e51326e1..539a45e5ca 100644 --- a/lib/realm-execution/include/realm-execution/redops/realm_redop_registry.h +++ b/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 diff --git a/lib/realm-execution/src/realm-execution/realm_manager.cc b/lib/realm-execution/src/realm-execution/realm_manager.cc index 04c5ec6bc1..89554fd5dc 100644 --- a/lib/realm-execution/src/realm-execution/realm_manager.cc +++ b/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() { diff --git a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu index ab3304836a..b137a15ae8 100644 --- a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu +++ b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu @@ -1,5 +1,15 @@ #include "realm-execution/redops/realm_redop_registry.h" -#include "realm-execution/redops/redop_id_t.h" +// Deliberately raw Realm headers rather than +// "realm-execution/redops/redop_id_t.h"/realm-execution/realm.h: those pull +// in PRealm's prealm.h, which fails a static_assert when parsed by nvcc, and +// we don't need any PRealm-specific behavior here anyway (redop registration +// isn't something PRealm wraps/instruments). +#include "realm-execution/redops/redop_id_t.dtg.h" +#include +#include +#if defined(__CUDACC__) || defined(__HIPCC__) +#include +#endif namespace FlexFlow { @@ -8,16 +18,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 class Pointer { public: @@ -121,6 +122,56 @@ private: #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(v); +} +__device__ __forceinline__ long long int + __ulonglong_as_longlong(unsigned long long int v) { + return static_cast(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((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(value) & 0xFFu) + << shift; + return (word & ~mask) | byte_val; +} + +// apply_cuda/fold_cuda are identical (just delegate to apply/fold) for every +// SumReduction 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 \ + __device__ static void apply_cuda(LHS &lhs, RHS rhs) { \ + apply(lhs, rhs); \ + } \ + template \ + __device__ static void fold_cuda(RHS &rhs1, RHS rhs2) { \ + fold(rhs1, rhs2); \ + } +#else +#define __LEGION_CUDA_REDOP_METHODS__ +#endif + template class SumReduction { // Empty definition @@ -139,6 +190,7 @@ public: __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs); template __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2); + __LEGION_CUDA_REDOP_METHODS__ }; template <> @@ -153,6 +205,7 @@ public: __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs); template __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2); + __LEGION_CUDA_REDOP_METHODS__ }; template <> @@ -167,6 +220,7 @@ public: __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs); template __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2); + __LEGION_CUDA_REDOP_METHODS__ }; template <> @@ -181,6 +235,7 @@ public: __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs); template __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2); + __LEGION_CUDA_REDOP_METHODS__ }; template <> @@ -195,6 +250,7 @@ public: __LEGION_CUDA_HD__ static void apply(LHS &lhs, RHS rhs); template __LEGION_CUDA_HD__ static void fold(RHS &rhs1, RHS rhs2); + __LEGION_CUDA_REDOP_METHODS__ }; template <> @@ -523,18 +579,51 @@ __LEGION_CUDA_HD__ inline void SumReduction::fold(RHS &rhs1, #endif } -void register_all_redops(Realm::Runtime rt) { +namespace { + +// Building the ReductionOpUntyped by hand (rather than via the +// ::Realm::Runtime::register_reduction convenience template) lets us +// call add_cuda_redop_kernels first, so the redop carries real GPU kernel +// function pointers (REDOP::apply_cuda/fold_cuda) instead of just its CPU +// apply/fold. Without this, Realm's GPUreduceChannel::supports_redop finds +// no CUDA-capable redop and refuses to build any GPU reduction-copy path +// (including same-device ones), aborting with "no path found ... (redop=N)" +// the first time a Reduction op actually runs on a GPU. +// +// Fully qualified as ::Realm:: throughout (rather than the FlexFlow::Realm +// alias to PRealm used elsewhere in this codebase) since this file +// deliberately avoids depending on PRealm at all — see the comment on the +// includes above. +template +void register_sum_redop(::Realm::Runtime &rt, ::Realm::ReductionOpID id) { + ::Realm::ReductionOpUntyped *redop = + ::Realm::ReductionOpUntyped::create_reduction_op(); +#if defined(__CUDACC__) || defined(__HIPCC__) + ::Realm::Cuda::add_cuda_redop_kernels(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>( - get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_BOOL_REDOP_ID)); - rt.register_reduction>( - get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_INT32_REDOP_ID)); - rt.register_reduction>( - get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_INT64_REDOP_ID)); - rt.register_reduction>( - get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_FLOAT_REDOP_ID)); - rt.register_reduction>( - get_realm_reduction_op_id_for_redop_id(redop_id_t::SUM_DOUBLE_REDOP_ID)); + register_sum_redop>( + rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_BOOL_REDOP_ID)); + register_sum_redop>( + rt, + static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT32_REDOP_ID)); + register_sum_redop>( + rt, + static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT64_REDOP_ID)); + register_sum_redop>( + rt, + static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_FLOAT_REDOP_ID)); + register_sum_redop>( + rt, + static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_DOUBLE_REDOP_ID)); } } // namespace FlexFlow From 359fc7c9d2111de6f9a1e424825ab3f0cb0ec7c1 Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 13 Aug 2026 14:59:35 -0700 Subject: [PATCH 3/7] Add parallel-op tests: gather-shaped RESHUFFLE, GPU e2e, base-test fix resolve_tensor_mappings previously only tested a pure-shuffle-shaped Repartition (both INPUT and OUTPUT unique); copy_insertion's RESHUFFLE handling determines which side is unique dynamically per-invocation, and that shape never exercised the OUTPUT-has-dup-coords branch. Adds a gather-shaped case that does. Adds GPU (cuda-realm-execution-tests) e2e training test cases for Replicate/Combine/Reduction/Repartition, and fixes the pre-existing (unrelated to this branch) base "GPU Model Parallelism" test, which requested only 1 GPU from Realm despite its config using two device coordinates. --- .../test/src/realm-execution/test_e2e.cc | 781 +++++++++++++----- .../task-spec/dynamic_graph/copy_insertion.cc | 647 +++++++++++++++ 2 files changed, 1202 insertions(+), 226 deletions(-) diff --git a/lib/realm-execution/test/src/realm-execution/test_e2e.cc b/lib/realm-execution/test/src/realm-execution/test_e2e.cc index 9ba4886b4b..c973573d64 100644 --- a/lib/realm-execution/test/src/realm-execution/test_e2e.cc +++ b/lib/realm-execution/test/src/realm-execution/test_e2e.cc @@ -4,7 +4,11 @@ #include "kernels/copy_tensor_accessor.h" #include "kernels/format_accessor_contents.h" #include "kernels/tensor_accessor_reductions.h" +#include "op-attrs/ff_dim_t.dtg.h" +#include "op-attrs/ops/combine_attrs.dtg.h" #include "op-attrs/ops/element_unary.h" +#include "op-attrs/ops/reduction_attrs.dtg.h" +#include "op-attrs/ops/repartition_attrs.dtg.h" #include "op-attrs/parallel_tensor_shape.h" #include "op-attrs/tensor_shape.dtg.h" #include "op-attrs/tensor_slot_name.dtg.h" @@ -116,16 +120,10 @@ static E2ETrainingConfig create_e2e_test_case() { std::nullopt}}, std::nullopt}, { - { - TensorSlotName::INPUT, - t_input, - }, + {TensorSlotName::INPUT, t_input}, }, { - { - TensorSlotName::WEIGHT, - t_weights_1, - }, + {TensorSlotName::WEIGHT, t_weights_1}, }); parallel_tensor_guid_t t_linear_1 = require_only_key(linear_operator_1.outputs, TensorSlotName::OUTPUT); @@ -139,16 +137,10 @@ static E2ETrainingConfig create_e2e_test_case() { std::nullopt}}, std::nullopt}, { - { - TensorSlotName::INPUT, - t_linear_1, - }, + {TensorSlotName::INPUT, t_linear_1}, }, { - { - TensorSlotName::WEIGHT, - t_weights_2, - }, + {TensorSlotName::WEIGHT, t_weights_2}, }); parallel_tensor_guid_t t_linear_2 = require_only_key(linear_operator_2.outputs, TensorSlotName::OUTPUT); @@ -223,24 +215,51 @@ static E2ETrainingConfig create_e2e_test_case() { }; } +static OptimizerAttrs make_sgd_optimizer() { + return OptimizerAttrs{SGDOptimizerAttrs{ + /*lr=*/0.001, + /*momentum=*/0.9, + /*nesterov=*/false, + /*weight_decay=*/0.001, + }}; +} + +static void run_one_epoch(RealmContext &ctx, + MappedParallelComputationGraph const &mpcg, + DeviceType device_type) { + std::map input_tensors; + + DistributedFfHandle device_handle = + create_distributed_ff_handle(ctx, + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + + PCGInstance pcg_instance = create_pcg_instance( + /*ctx=*/ctx, + /*mpcg=*/mpcg, + /*optimizer=*/make_sgd_optimizer(), + /*loss=*/std::nullopt, + /*input_tensors=*/input_tensors, + /*profiling_settings=*/ProfilingSettings{0, 0}, + /*device_handle=*/device_handle, + /*device_type=*/device_type); + + perform_all_passes_for_pcg_instance( + /*instance=*/pcg_instance, + /*profiling_settings=*/ProfilingSettings{0, 0}, + /*device_handle=*/device_handle); +} + MappedParallelComputationGraph make_test_replicate_mpcg_for_device_type(DeviceType device_type) { positive_int batch_size = 10_p; positive_int data_dim = 16_p; - positive_int hidden_dim = 32_p; - positive_int output_dim = 1_p; - - TensorShape output_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; - - TensorShape label_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; - - ParallelComputationGraph pcg = empty_parallel_computation_graph(); TensorShape input_tensor_shape = TensorShape{TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; + ParallelComputationGraph pcg = empty_parallel_computation_graph(); + ParallelLayerAddedResult inputs_layer = pcg_add_input_layer(pcg, input_tensor_shape); parallel_tensor_guid_t t_input = @@ -251,61 +270,32 @@ MappedParallelComputationGraph parallel_tensor_guid_t t_input_2 = require_only_key(inputs_layer_2.outputs, TensorSlotName::OUTPUT); - ElementBinaryAttrs add_attrs = ElementBinaryAttrs{ - OperatorType::EW_ADD, - DataType::FLOAT, - false, - false, - }; + ElementBinaryAttrs add_attrs = + ElementBinaryAttrs{OperatorType::EW_ADD, DataType::FLOAT, false, false}; ParallelLayerAddedResult add_operator_1 = add_parallel_layer(pcg, make_layer_attrs(add_attrs), - { - { - TensorSlotName::LHS_INPUT, - t_input, - }, - { - TensorSlotName::RHS_INPUT, - t_input_2, - }, - }, - /*weights=*/{}); - + {{TensorSlotName::LHS_INPUT, t_input}, + {TensorSlotName::RHS_INPUT, t_input_2}}, + {}); parallel_tensor_guid_t t_add_1 = require_only_key(add_operator_1.outputs, TensorSlotName::OUTPUT); - positive_int replicate_degree = 2_p; - ReplicateAttrs repl_attrs = ReplicateAttrs{replicate_degree}; + ReplicateAttrs repl_attrs = ReplicateAttrs{/*replicate_degree=*/2_p}; ParallelLayerAddedResult repl_operator_1 = add_parallel_layer(pcg, make_layer_attrs(repl_attrs), - { - { - TensorSlotName::INPUT, - t_add_1, - }, - }, - /*weight=*/{}); - + {{TensorSlotName::INPUT, t_add_1}}, + {}); parallel_tensor_guid_t t_repl_1 = require_only_key(repl_operator_1.outputs, TensorSlotName::OUTPUT); ParallelLayerAddedResult relu_operator_1 = add_parallel_layer(pcg, make_layer_attrs(make_relu_attrs()), - /*inputs=*/ - { - { - TensorSlotName::INPUT, - t_repl_1, - }, - }, - /*weights=*/{}); - - parallel_tensor_guid_t t_relu_1 = - require_only_key(relu_operator_1.outputs, TensorSlotName::OUTPUT); + {{TensorSlotName::INPUT, t_repl_1}}, + {}); MachineSpaceCoordinate mc0{0_n, 0_n}; MachineSpaceCoordinate mc1{0_n, 1_n}; @@ -319,96 +309,413 @@ MappedParallelComputationGraph /*discard_copy_component=*/1_n, /*shard_component=*/FFOrdered{0_n}}; - MappedParallelComputationGraph mpcg = - mapped_pcg_from_pcg_and_mapped_op_task_groups( - /*pcg=*/pcg, - /*mapped_op_task_groups=*/{ - { - inputs_layer.parallel_layer, - MappedOperatorTaskGroup{ - { - { - mc0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::OUTPUT, tensor_coord0}, - }}, - }, - }, - }, - }, - { - inputs_layer_2.parallel_layer, - MappedOperatorTaskGroup{ - { - { - mc0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::OUTPUT, tensor_coord0}, - }}, - }, - }, - }, - }, - { - add_operator_1.parallel_layer, - MappedOperatorTaskGroup{ - { - { - mc0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::LHS_INPUT, tensor_coord0}, - {TensorSlotName::RHS_INPUT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}, - }, - }, - }, - }, - { - repl_operator_1.parallel_layer, - MappedOperatorTaskGroup{ - { - { - mc0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}, - }, - { - mc1, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord1}, - }}, - }, - }, - }, - }, - { - relu_operator_1.parallel_layer, - MappedOperatorTaskGroup{ - { - { - mc0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}, - }, - { - mc1, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord1}, - {TensorSlotName::OUTPUT, tensor_coord1}, - }}, - }, - }, - }, - }, - }); + return mapped_pcg_from_pcg_and_mapped_op_task_groups( + pcg, + { + {inputs_layer.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, + {inputs_layer_2.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, + {add_operator_1.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::LHS_INPUT, tensor_coord0}, + {TensorSlotName::RHS_INPUT, tensor_coord0}, + {TensorSlotName::OUTPUT, tensor_coord0}, + }}}}}}, + {repl_operator_1.parallel_layer, + MappedOperatorTaskGroup{ + { + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord0}, + {TensorSlotName::OUTPUT, tensor_coord0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord0}, + {TensorSlotName::OUTPUT, tensor_coord1}, + }}}, + }, + }}, + {relu_operator_1.parallel_layer, + MappedOperatorTaskGroup{ + { + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord0}, + {TensorSlotName::OUTPUT, tensor_coord0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord1}, + {TensorSlotName::OUTPUT, tensor_coord1}, + }}}, + }, + }}, + }); +} +MappedParallelComputationGraph + make_test_combine_mpcg_for_device_type(DeviceType device_type) { + positive_int batch_size = 10_p; + positive_int data_dim = 16_p; + + TensorShape input_tensor_shape = + TensorShape{TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; + + ParallelComputationGraph pcg = empty_parallel_computation_graph(); + + ParallelLayerAddedResult inputs_layer = + pcg_add_input_layer(pcg, input_tensor_shape); + parallel_tensor_guid_t t_input = + require_only_key(inputs_layer.outputs, TensorSlotName::OUTPUT); + + // Repartition along dim 0 (batch) to get shard_degree[0]=2 + RepartitionAttrs repartition_attrs = RepartitionAttrs{ff_dim_t{0_n}, 2_p}; + ParallelLayerAddedResult repartition_op = + add_parallel_layer(pcg, + make_layer_attrs(repartition_attrs), + {{TensorSlotName::INPUT, t_input}}, + {}); + parallel_tensor_guid_t t_repartitioned = + require_only_key(repartition_op.outputs, TensorSlotName::OUTPUT); + + // Combine along dim 0 to merge shards back + CombineAttrs combine_attrs = CombineAttrs{ff_dim_t{0_n}, 2_p}; + ParallelLayerAddedResult combine_op = + add_parallel_layer(pcg, + make_layer_attrs(combine_attrs), + {{TensorSlotName::INPUT, t_repartitioned}}, + {}); + parallel_tensor_guid_t t_combine = + require_only_key(combine_op.outputs, TensorSlotName::OUTPUT); + + ParallelLayerAddedResult relu_op = + add_parallel_layer(pcg, + make_layer_attrs(make_relu_attrs()), + {{TensorSlotName::INPUT, t_combine}}, + {}); + + MachineSpaceCoordinate mc0{0_n, 0_n}; + MachineSpaceCoordinate mc1{0_n, 1_n}; + + // Input: single shard on mc0 + ParallelTensorSpaceCoordinate coord0{0_n, 0_n, FFOrdered{0_n}}; + // After repartition: two shards along batch dim + ParallelTensorSpaceCoordinate repart_coord0{0_n, 0_n, FFOrdered{0_n}}; + ParallelTensorSpaceCoordinate repart_coord1{0_n, 0_n, FFOrdered{1_n}}; + // After combine: back to single shard + ParallelTensorSpaceCoordinate out_coord{0_n, 0_n, FFOrdered{0_n}}; + + return mapped_pcg_from_pcg_and_mapped_op_task_groups( + pcg, + { + {inputs_layer.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, coord0}}}}}}}, + {repartition_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, coord0}, + {TensorSlotName::OUTPUT, repart_coord0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, coord0}, + {TensorSlotName::OUTPUT, repart_coord1}, + }}}, + }}}, + {combine_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, repart_coord0}, + {TensorSlotName::OUTPUT, out_coord}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, repart_coord1}, + {TensorSlotName::OUTPUT, out_coord}, + }}}, + }}}, + {relu_op.parallel_layer, + MappedOperatorTaskGroup{{{mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, out_coord}, + {TensorSlotName::OUTPUT, out_coord}, + }}}}}}, + }); +} + +MappedParallelComputationGraph + make_test_repartition_mpcg_for_device_type(DeviceType device_type) { + positive_int batch_size = 10_p; + positive_int data_dim = 16_p; + + TensorShape input_tensor_shape = + TensorShape{TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; + + ParallelComputationGraph pcg = empty_parallel_computation_graph(); + + ParallelLayerAddedResult inputs_layer = + pcg_add_input_layer(pcg, input_tensor_shape); + parallel_tensor_guid_t t_input = + require_only_key(inputs_layer.outputs, TensorSlotName::OUTPUT); + + // Repartition along dim 0 + RepartitionAttrs repartition_attrs = RepartitionAttrs{ff_dim_t{0_n}, 2_p}; + ParallelLayerAddedResult repartition_op = + add_parallel_layer(pcg, + make_layer_attrs(repartition_attrs), + {{TensorSlotName::INPUT, t_input}}, + {}); + parallel_tensor_guid_t t_repartitioned = + require_only_key(repartition_op.outputs, TensorSlotName::OUTPUT); + + // Relu on repartitioned tensor + ParallelLayerAddedResult relu_op = + add_parallel_layer(pcg, + make_layer_attrs(make_relu_attrs()), + {{TensorSlotName::INPUT, t_repartitioned}}, + {}); + + MachineSpaceCoordinate mc0{0_n, 0_n}; + MachineSpaceCoordinate mc1{0_n, 1_n}; + + ParallelTensorSpaceCoordinate coord0{0_n, 0_n, FFOrdered{0_n}}; + ParallelTensorSpaceCoordinate repart_coord0{0_n, 0_n, FFOrdered{0_n}}; + ParallelTensorSpaceCoordinate repart_coord1{0_n, 0_n, FFOrdered{1_n}}; + + return mapped_pcg_from_pcg_and_mapped_op_task_groups( + pcg, + { + {inputs_layer.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, coord0}}}}}}}, + {repartition_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, coord0}, + {TensorSlotName::OUTPUT, repart_coord0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, coord0}, + {TensorSlotName::OUTPUT, repart_coord1}, + }}}, + }}}, + {relu_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, repart_coord0}, + {TensorSlotName::OUTPUT, repart_coord0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, repart_coord1}, + {TensorSlotName::OUTPUT, repart_coord1}, + }}}, + }}}, + }); +} + +MappedParallelComputationGraph + make_test_reduction_mpcg_for_device_type(DeviceType device_type) { + positive_int batch_size = 10_p; + positive_int data_dim = 16_p; + positive_int hidden_dim = 8_p; + + TensorShape input_tensor_shape = + TensorShape{TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; + TensorShape weight_shape = + TensorShape{TensorDims{FFOrdered{hidden_dim, data_dim}}, DataType::FLOAT}; + + ParallelComputationGraph pcg = empty_parallel_computation_graph(); + + ParallelLayerAddedResult inputs_layer = + pcg_add_input_layer(pcg, input_tensor_shape); + parallel_tensor_guid_t t_input = + require_only_key(inputs_layer.outputs, TensorSlotName::OUTPUT); + + ParallelLayerAddedResult weights_layer = add_parallel_layer( + pcg, + ParallelLayerAttrs{ + PCGOperatorAttrs{WeightAttrs{weight_shape, + InitializerAttrs{GlorotNormalAttrs{0}}}}, + std::nullopt}, + {}, + {}); + parallel_tensor_guid_t t_weight = + require_only_key(weights_layer.outputs, TensorSlotName::OUTPUT); + + // Repartition input along last dim (data_dim) → shard_degrees[-1]=2 + RepartitionAttrs input_repartition_attrs = + RepartitionAttrs{ff_dim_t{1_n}, 2_p}; + ParallelLayerAddedResult input_repartition_op = + add_parallel_layer(pcg, + make_layer_attrs(input_repartition_attrs), + {{TensorSlotName::INPUT, t_input}}, + {}); + parallel_tensor_guid_t t_input_repartitioned = + require_only_key(input_repartition_op.outputs, TensorSlotName::OUTPUT); + + // Repartition weight along last dim (data_dim) → shard_degrees[1]=2 + // Required by Linear when input has shard_degrees[-1]=2 + RepartitionAttrs weight_repartition_attrs = + RepartitionAttrs{ff_dim_t{1_n}, 2_p}; + ParallelLayerAddedResult weight_repartition_op = + add_parallel_layer(pcg, + make_layer_attrs(weight_repartition_attrs), + {{TensorSlotName::INPUT, t_weight}}, + {}); + parallel_tensor_guid_t t_weight_repartitioned = + require_only_key(weight_repartition_op.outputs, TensorSlotName::OUTPUT); + + // Linear: input shard_degrees[-1]=2, weight shard_degrees[1]=2 + // → output sum_degree=2 + ParallelLayerAddedResult linear_op = add_parallel_layer( + pcg, + ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{hidden_dim, + false, + DataType::FLOAT, + Activation::RELU, + std::nullopt}}, + std::nullopt}, + {{TensorSlotName::INPUT, t_input_repartitioned}}, + {{TensorSlotName::WEIGHT, t_weight_repartitioned}}); + parallel_tensor_guid_t t_linear = + require_only_key(linear_op.outputs, TensorSlotName::OUTPUT); + + // Reduction: sum_degree 2→1 + ReductionAttrs reduction_attrs = ReductionAttrs{2_p}; + ParallelLayerAddedResult reduction_op = + add_parallel_layer(pcg, + make_layer_attrs(reduction_attrs), + {{TensorSlotName::INPUT, t_linear}}, + {}); + parallel_tensor_guid_t t_reduction = + require_only_key(reduction_op.outputs, TensorSlotName::OUTPUT); + + ParallelLayerAddedResult relu_op = + add_parallel_layer(pcg, + make_layer_attrs(make_relu_attrs()), + {{TensorSlotName::INPUT, t_reduction}}, + {}); - return mpcg; + MachineSpaceCoordinate mc0{0_n, 0_n}; + MachineSpaceCoordinate mc1{0_n, 1_n}; + + // Coords for input + ParallelTensorSpaceCoordinate input_coord{0_n, 0_n, FFOrdered{0_n}}; + // Coords for weight (unsharded on mc0) + ParallelTensorSpaceCoordinate weight_coord{0_n, 0_n, FFOrdered{0_n, 0_n}}; + // After repartitioning input along dim 1: two shards + ParallelTensorSpaceCoordinate input_repart0{0_n, 0_n, FFOrdered{0_n, 0_n}}; + ParallelTensorSpaceCoordinate input_repart1{0_n, 0_n, FFOrdered{0_n, 1_n}}; + // After repartitioning weight along dim 1: two shards + ParallelTensorSpaceCoordinate weight_repart0{0_n, 0_n, FFOrdered{0_n, 0_n}}; + ParallelTensorSpaceCoordinate weight_repart1{0_n, 0_n, FFOrdered{0_n, 1_n}}; + // Linear output: partial sums (sum_component distinguishes) + ParallelTensorSpaceCoordinate linear_coord0{0_n, 0_n, FFOrdered{0_n}}; + ParallelTensorSpaceCoordinate linear_coord1{1_n, 0_n, FFOrdered{0_n}}; + // After reduction + ParallelTensorSpaceCoordinate out_coord{0_n, 0_n, FFOrdered{0_n}}; + + return mapped_pcg_from_pcg_and_mapped_op_task_groups( + pcg, + { + // Input on mc0 only + {inputs_layer.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, input_coord}}}}}}}, + // Weight on mc0 only (unsharded) + {weights_layer.parallel_layer, + MappedOperatorTaskGroup{ + {{mc0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, weight_coord}}}}}}}, + // Input repartition: scatter input to both devices + {input_repartition_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, input_coord}, + {TensorSlotName::OUTPUT, input_repart0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, input_coord}, + {TensorSlotName::OUTPUT, input_repart1}, + }}}, + }}}, + // Weight repartition: scatter weight to both devices + {weight_repartition_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, weight_coord}, + {TensorSlotName::OUTPUT, weight_repart0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, weight_coord}, + {TensorSlotName::OUTPUT, weight_repart1}, + }}}, + }}}, + // Linear: each device computes partial matmul + {linear_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, input_repart0}, + {TensorSlotName::WEIGHT, weight_repart0}, + {TensorSlotName::OUTPUT, linear_coord0}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, input_repart1}, + {TensorSlotName::WEIGHT, weight_repart1}, + {TensorSlotName::OUTPUT, linear_coord1}, + }}}, + }}}, + // Reduction: sum partial results + {reduction_op.parallel_layer, + MappedOperatorTaskGroup{{ + {mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, linear_coord0}, + {TensorSlotName::OUTPUT, out_coord}, + }}}, + {mc1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, linear_coord1}, + {TensorSlotName::OUTPUT, out_coord}, + }}}, + }}}, + // Relu on single output + {relu_op.parallel_layer, + MappedOperatorTaskGroup{{{mc0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, out_coord}, + {TensorSlotName::OUTPUT, out_coord}, + }}}}}}, + }); } TEST_SUITE(FF_TEST_SUITE) { @@ -453,7 +760,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*device_handle=*/device_handle, /*device_type=*/DeviceType::CPU); - // begin training loop int num_epochs = 5; std::vector loss_values; @@ -473,8 +779,6 @@ TEST_SUITE(FF_TEST_SUITE) { allocator)); } - // Assert that each sample in the batch has a lower loss in last epoch - // than the first epoch GenericTensorAccessorR first_epoch_loss = loss_values.at(0); GenericTensorAccessorR last_epoch_loss = loss_values.back(); CHECK_MESSAGE( @@ -495,45 +799,58 @@ TEST_SUITE(FF_TEST_SUITE) { RealmManager manager = RealmManager{&fake_argc, &fake_argv}; ControllerTaskResult result = manager.start_controller([](RealmContext &ctx) { - Allocator allocator = ctx.get_current_device_allocator(); - MappedParallelComputationGraph mpcg = make_test_replicate_mpcg_for_device_type(DeviceType::CPU); + run_one_epoch(ctx, mpcg, DeviceType::CPU); + }); + result.wait(); + } - std::map input_tensors; + TEST_CASE("RealmBackend e2e Training Combine Op (CPU Model Parallelism)") { + std::vector fake_args = + make_fake_realm_args(/*num_cpus=*/2_p, /*num_gpus=*/0_n); + int fake_argc = fake_args.size(); + char **fake_argv = fake_args.data(); - OptimizerAttrs optimizer_attrs = OptimizerAttrs{ - SGDOptimizerAttrs{ - /*lr=*/0.001, - /*momentum=*/0.9, - /*nesterov=*/false, - /*weight_decay=*/0.001, - }, - }; + RealmManager manager = RealmManager{&fake_argc, &fake_argv}; + ControllerTaskResult result = + manager.start_controller([](RealmContext &ctx) { + MappedParallelComputationGraph mpcg = + make_test_combine_mpcg_for_device_type(DeviceType::CPU); + run_one_epoch(ctx, mpcg, DeviceType::CPU); + }); + result.wait(); + } - DistributedFfHandle device_handle = create_distributed_ff_handle( - ctx, - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); + TEST_CASE("RealmBackend e2e Training Reduction Op (CPU Model Parallelism)") { + std::vector fake_args = + make_fake_realm_args(/*num_cpus=*/2_p, /*num_gpus=*/0_n); + int fake_argc = fake_args.size(); + char **fake_argv = fake_args.data(); - PCGInstance pcg_instance = create_pcg_instance( - /*ctx=*/ctx, - /*mpcg=*/mpcg, - /*optimizer=*/optimizer_attrs, - /*loss=*/std::nullopt, - /*input_tensors=*/input_tensors, - /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle, - /*device_type=*/DeviceType::CPU); + RealmManager manager = RealmManager{&fake_argc, &fake_argv}; + ControllerTaskResult result = + manager.start_controller([](RealmContext &ctx) { + MappedParallelComputationGraph mpcg = + make_test_reduction_mpcg_for_device_type(DeviceType::CPU); + run_one_epoch(ctx, mpcg, DeviceType::CPU); + }); + result.wait(); + } - // begin training loop - int num_epochs = 1; - for (int i = 0; i < num_epochs; i++) { - perform_all_passes_for_pcg_instance( - /*instance=*/pcg_instance, - /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle); - } + TEST_CASE( + "RealmBackend e2e Training Repartition Op (CPU Model Parallelism)") { + std::vector fake_args = + make_fake_realm_args(/*num_cpus=*/2_p, /*num_gpus=*/0_n); + int fake_argc = fake_args.size(); + char **fake_argv = fake_args.data(); + + RealmManager manager = RealmManager{&fake_argc, &fake_argv}; + ControllerTaskResult result = + manager.start_controller([](RealmContext &ctx) { + MappedParallelComputationGraph mpcg = + make_test_repartition_mpcg_for_device_type(DeviceType::CPU); + run_one_epoch(ctx, mpcg, DeviceType::CPU); }); result.wait(); } @@ -544,8 +861,11 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { E2ETrainingConfig cfg = create_e2e_test_case(); //! [realm-execution example] + // create_e2e_test_case() maps weights_layer_2/linear_operator_2 to a + // second device (MachineSpaceCoordinate{0, 1}), so this needs 2 real + // GPUs, not 1 — matching every other GPU TEST_CASE in this suite. std::vector fake_args = - make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/1_n); + make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/2_n); int fake_argc = fake_args.size(); char **fake_argv = fake_args.data(); @@ -583,7 +903,6 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*device_handle=*/device_handle, /*device_type=*/DeviceType::GPU); - // begin training loop int num_epochs = 5; std::vector loss_values; @@ -604,8 +923,6 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { allocator)); } - // Assert that each sample in the batch has a lower loss in last epoch - // than the first epoch GenericTensorAccessorR first_epoch_loss = loss_values.at(0); GenericTensorAccessorR last_epoch_loss = loss_values.back(); CHECK_MESSAGE( @@ -627,48 +944,60 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { char **fake_argv = fake_args.data(); RealmManager manager = RealmManager{&fake_argc, &fake_argv}; - ControllerTaskResult result = manager.start_controller([](RealmContext &ctx) { - Allocator allocator = ctx.get_current_device_allocator(); - MappedParallelComputationGraph mpcg = make_test_replicate_mpcg_for_device_type(DeviceType::GPU); + run_one_epoch(ctx, mpcg, DeviceType::GPU); + }); + result.wait(); + } - OptimizerAttrs optimizer_attrs = OptimizerAttrs{ - SGDOptimizerAttrs{ - /*lr=*/0.001, - /*momentum=*/0.9, - /*nesterov=*/false, - /*weight_decay=*/0.001, - }, - }; + TEST_CASE("RealmBackend e2e Training Combine Op (GPU Model Parallelism)") { + std::vector fake_args = + make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/2_n); + int fake_argc = fake_args.size(); + char **fake_argv = fake_args.data(); - std::map input_tensors; + RealmManager manager = RealmManager{&fake_argc, &fake_argv}; + ControllerTaskResult result = + manager.start_controller([](RealmContext &ctx) { + MappedParallelComputationGraph mpcg = + make_test_combine_mpcg_for_device_type(DeviceType::GPU); + run_one_epoch(ctx, mpcg, DeviceType::GPU); + }); + result.wait(); + } - DistributedFfHandle device_handle = create_distributed_ff_handle( - ctx, - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); + TEST_CASE("RealmBackend e2e Training Reduction Op (GPU Model Parallelism)") { + std::vector fake_args = + make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/2_n); + int fake_argc = fake_args.size(); + char **fake_argv = fake_args.data(); - PCGInstance pcg_instance = create_pcg_instance( - /*ctx=*/ctx, - /*mpcg=*/mpcg, - /*optimizer=*/optimizer_attrs, - /*loss=*/std::nullopt, - /*input_tensors=*/input_tensors, - /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle, - /*device_type=*/DeviceType::GPU); + RealmManager manager = RealmManager{&fake_argc, &fake_argv}; + ControllerTaskResult result = + manager.start_controller([](RealmContext &ctx) { + MappedParallelComputationGraph mpcg = + make_test_reduction_mpcg_for_device_type(DeviceType::GPU); + run_one_epoch(ctx, mpcg, DeviceType::GPU); + }); + result.wait(); + } - // begin training loop - int num_epochs = 1; - for (int i = 0; i < num_epochs; i++) { - perform_all_passes_for_pcg_instance( - /*instance=*/pcg_instance, - /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle); - } + TEST_CASE( + "RealmBackend e2e Training Repartition Op (GPU Model Parallelism)") { + std::vector fake_args = + make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/2_n); + int fake_argc = fake_args.size(); + char **fake_argv = fake_args.data(); + + RealmManager manager = RealmManager{&fake_argc, &fake_argv}; + ControllerTaskResult result = + manager.start_controller([](RealmContext &ctx) { + MappedParallelComputationGraph mpcg = + make_test_repartition_mpcg_for_device_type(DeviceType::GPU); + run_one_epoch(ctx, mpcg, DeviceType::GPU); }); result.wait(); } diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc index 055586664a..b53575aabf 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc @@ -1,5 +1,9 @@ #include "task-spec/dynamic_graph/copy_insertion.h" +#include "op-attrs/ff_dim_t.dtg.h" +#include "op-attrs/ops/combine_attrs.dtg.h" #include "op-attrs/ops/element_unary.h" +#include "op-attrs/ops/reduction_attrs.dtg.h" +#include "op-attrs/ops/repartition_attrs.dtg.h" #include "op-attrs/tensor_slot_name.dtg.h" #include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" @@ -1922,4 +1926,647 @@ TEST_SUITE(FF_TEST_SUITE) { } } } + TEST_CASE("resolve_tensor_mappings - combine operator") { + MachineSpaceCoordinate mc1 = mk_machine_coord(0_n); + MachineSpaceCoordinate mc2 = mk_machine_coord(1_n); + MachineSpaceCoordinate mc3 = mk_machine_coord(2_n); + + auto mk_pt_coord = + [](nonnegative_int idx) -> ParallelTensorSpaceCoordinate { + return ParallelTensorSpaceCoordinate{ + /*sum_component=*/0_n, + /*discard_copy_component=*/idx, + /*shard_components=*/FFOrdered{0_n, 0_n}, + }; + }; + + auto mk_node_mapping = + [](MappedOperatorTaskGroup const &op_task_group) -> DynamicNodeMapping { + return DynamicNodeMapping{op_task_group, DeviceType::GPU}; + }; + + auto mk_single_mapping = + [&](MachineSpaceCoordinate const &mc, + nonnegative_int pt_idx) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(pt_idx), mk_device_id(mc)}, + }, + }; + }; + + auto mk_two_mapping = + [&](MachineSpaceCoordinate const &mc1_, + MachineSpaceCoordinate const &mc2_) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(0_n), mk_device_id(mc1_)}, + {mk_pt_coord(1_n), mk_device_id(mc2_)}, + }, + }; + }; + + // Graph: input → relu → combine → relu + // Two input shards combined into one output shard + + DynamicValueAttrs input_output = + mk_value_attrs(101, TensorSlotName::OUTPUT, std::nullopt); + DynamicValueAttrs relu1_output = + mk_value_attrs(102, TensorSlotName::OUTPUT, std::nullopt); + DynamicValueAttrs combine_output = + mk_value_attrs(103, TensorSlotName::OUTPUT, std::nullopt); + + MappedOperatorTaskGroup relu1_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(1_n)}}}}, + }, + }; + + MappedOperatorTaskGroup combine_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + MappedOperatorTaskGroup relu2_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc3, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + DynamicNodeInvocation relu1_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), input_output}}, + /*node_attrs=*/ + mk_node_attrs(102, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu1_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), relu1_output}}, + }; + + DynamicNodeInvocation combine_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), relu1_output}}, + /*node_attrs=*/ + mk_node_attrs(103, + PCGOperatorAttrs{CombineAttrs{ff_dim_t{0_n}, 2_p}}, + mk_node_mapping(combine_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), combine_output}}, + }; + + DynamicNodeInvocation relu2_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), combine_output}}, + /*node_attrs=*/ + mk_node_attrs(104, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu2_mapping_group)), + /*outputs=*/ + {{mk_slot(TensorSlotName::OUTPUT), + mk_value_attrs(104, TensorSlotName::OUTPUT, std::nullopt)}}, + }; + + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set( + {relu1_invocation, combine_invocation, relu2_invocation}); + + dynamic_invocation_id_t relu1_id = + dynamic_graph_get_id_for_invocation(g, relu1_invocation); + dynamic_invocation_id_t combine_id = + dynamic_graph_get_id_for_invocation(g, combine_invocation); + dynamic_invocation_id_t relu2_id = + dynamic_graph_get_id_for_invocation(g, relu2_invocation); + + SUBCASE("combine output resolved from node mapping") { + std::map result = + resolve_tensor_mappings(g); + + // Combine output (1-side) must be resolved from node mapping + InternalDynamicSlotSite const combine_out = InternalDynamicSlotSite{ + combine_id, + TensorDirection::OUTPUT, + mk_slot(TensorSlotName::OUTPUT), + }; + + ParallelTensorMapping const expected_combine_out = + mk_single_mapping(mc3, 0_n); + + REQUIRE(result.count(combine_out) == 1); + CHECK(result.at(combine_out) == expected_combine_out); + } + + SUBCASE("combine input resolved from adjacent relu output") { + std::map result = + resolve_tensor_mappings(g); + + // Combine input (N-side) resolved from adjacent relu output mapping + InternalDynamicSlotSite const combine_in = InternalDynamicSlotSite{ + combine_id, + TensorDirection::INCOMING, + mk_slot(TensorSlotName::INPUT), + }; + + ParallelTensorMapping const expected_combine_in = + mk_two_mapping(mc1, mc2); + + REQUIRE(result.count(combine_in) == 1); + CHECK(result.at(combine_in) == expected_combine_in); + } + + SUBCASE("relu after combine gets single-shard mapping") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const relu2_in = InternalDynamicSlotSite{ + relu2_id, + TensorDirection::INCOMING, + mk_slot(TensorSlotName::INPUT), + }; + + ParallelTensorMapping const expected = mk_single_mapping(mc3, 0_n); + + REQUIRE(result.count(relu2_in) == 1); + CHECK(result.at(relu2_in) == expected); + } + } + + TEST_CASE("resolve_tensor_mappings - reduction operator") { + MachineSpaceCoordinate mc1 = mk_machine_coord(0_n); + MachineSpaceCoordinate mc2 = mk_machine_coord(1_n); + MachineSpaceCoordinate mc3 = mk_machine_coord(2_n); + + auto mk_pt_coord = + [](nonnegative_int idx) -> ParallelTensorSpaceCoordinate { + return ParallelTensorSpaceCoordinate{ + 0_n, + idx, + FFOrdered{0_n, 0_n}, + }; + }; + + auto mk_node_mapping = + [](MappedOperatorTaskGroup const &op_task_group) -> DynamicNodeMapping { + return DynamicNodeMapping{op_task_group, DeviceType::GPU}; + }; + + auto mk_single_mapping = + [&](MachineSpaceCoordinate const &mc, + nonnegative_int pt_idx) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(pt_idx), mk_device_id(mc)}, + }, + }; + }; + + auto mk_two_mapping = + [&](MachineSpaceCoordinate const &mc1_, + MachineSpaceCoordinate const &mc2_) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(0_n), mk_device_id(mc1_)}, + {mk_pt_coord(1_n), mk_device_id(mc2_)}, + }, + }; + }; + + // Same topology as combine: N inputs → 1 output, but with sum semantics + DynamicValueAttrs relu1_output = + mk_value_attrs(201, TensorSlotName::OUTPUT, std::nullopt); + DynamicValueAttrs reduction_output = + mk_value_attrs(202, TensorSlotName::OUTPUT, std::nullopt); + + MappedOperatorTaskGroup relu1_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(1_n)}}}}, + }, + }; + + MappedOperatorTaskGroup reduction_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + MappedOperatorTaskGroup relu2_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc3, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + DynamicNodeInvocation relu1_invocation = DynamicNodeInvocation{ + /*inputs=*/{ + {mk_slot(TensorSlotName::INPUT), + mk_value_attrs(200, TensorSlotName::OUTPUT, std::nullopt)}}, + /*node_attrs=*/ + mk_node_attrs(201, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu1_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), relu1_output}}, + }; + + DynamicNodeInvocation reduction_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), relu1_output}}, + /*node_attrs=*/ + mk_node_attrs(202, + PCGOperatorAttrs{ReductionAttrs{2_p}}, + mk_node_mapping(reduction_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), reduction_output}}, + }; + + DynamicNodeInvocation relu2_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), reduction_output}}, + /*node_attrs=*/ + mk_node_attrs(203, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu2_mapping_group)), + /*outputs=*/ + {{mk_slot(TensorSlotName::OUTPUT), + mk_value_attrs(203, TensorSlotName::OUTPUT, std::nullopt)}}, + }; + + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set( + {relu1_invocation, reduction_invocation, relu2_invocation}); + + dynamic_invocation_id_t reduction_id = + dynamic_graph_get_id_for_invocation(g, reduction_invocation); + + SUBCASE("reduction output resolved from node mapping") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const reduction_out = InternalDynamicSlotSite{ + reduction_id, + TensorDirection::OUTPUT, + mk_slot(TensorSlotName::OUTPUT), + }; + + REQUIRE(result.count(reduction_out) == 1); + CHECK(result.at(reduction_out) == mk_single_mapping(mc3, 0_n)); + } + + SUBCASE("reduction input resolved from adjacent relu output") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const reduction_in = InternalDynamicSlotSite{ + reduction_id, + TensorDirection::INCOMING, + mk_slot(TensorSlotName::INPUT), + }; + + REQUIRE(result.count(reduction_in) == 1); + CHECK(result.at(reduction_in) == mk_two_mapping(mc1, mc2)); + } + } + + TEST_CASE("resolve_tensor_mappings - repartition operator") { + MachineSpaceCoordinate mc1 = mk_machine_coord(0_n); + MachineSpaceCoordinate mc2 = mk_machine_coord(1_n); + + auto mk_pt_coord = + [](nonnegative_int idx) -> ParallelTensorSpaceCoordinate { + return ParallelTensorSpaceCoordinate{ + 0_n, + idx, + FFOrdered{0_n, 0_n}, + }; + }; + + auto mk_node_mapping = + [](MappedOperatorTaskGroup const &op_task_group) -> DynamicNodeMapping { + return DynamicNodeMapping{op_task_group, DeviceType::GPU}; + }; + + auto mk_single_mapping = + [&](MachineSpaceCoordinate const &mc, + nonnegative_int pt_idx) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(pt_idx), mk_device_id(mc)}, + }, + }; + }; + + // Pure-shuffle-shaped repartition: both INPUT and OUTPUT are unique + // (input pt0→mc1, pt1→mc2 remapped to output pt0→mc2, pt1→mc1), but + // copy insertion only ever reads OUTPUT directly from the node mapping + // (checked first) and derives INPUT via adjacent-value resolution + // instead — it never actually checks whether INPUT is *also* unique. + // See the "gather-shaped" TEST_CASE below for the case where OUTPUT is + // NOT unique, which this topology never exercises. + + DynamicValueAttrs relu1_output = + mk_value_attrs(301, TensorSlotName::OUTPUT, std::nullopt); + DynamicValueAttrs repartition_output = + mk_value_attrs(302, TensorSlotName::OUTPUT, std::nullopt); + + MappedOperatorTaskGroup relu1_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(1_n)}}}}, + }, + }; + + // Repartition swaps: mc1 gets pt1 output, mc2 gets pt0 output + MappedOperatorTaskGroup repartition_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(1_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + DynamicNodeInvocation relu1_invocation = DynamicNodeInvocation{ + /*inputs=*/{ + {mk_slot(TensorSlotName::INPUT), + mk_value_attrs(300, TensorSlotName::OUTPUT, std::nullopt)}}, + /*node_attrs=*/ + mk_node_attrs(301, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu1_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), relu1_output}}, + }; + + DynamicNodeInvocation repartition_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), relu1_output}}, + /*node_attrs=*/ + mk_node_attrs(302, + PCGOperatorAttrs{RepartitionAttrs{ff_dim_t{0_n}, 2_p}}, + mk_node_mapping(repartition_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), repartition_output}}, + }; + + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set( + {relu1_invocation, repartition_invocation}); + + dynamic_invocation_id_t repartition_id = + dynamic_graph_get_id_for_invocation(g, repartition_invocation); + + SUBCASE("repartition input resolved from adjacent values") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const repartition_in = InternalDynamicSlotSite{ + repartition_id, + TensorDirection::INCOMING, + mk_slot(TensorSlotName::INPUT), + }; + + // Repartition INPUT resolved from adjacent relu1 output + // relu1 has mc1→pt0, mc2→pt1 + ParallelTensorMapping const expected_in = ParallelTensorMapping{ + bidict{ + {mk_pt_coord(0_n), mk_device_id(mc1)}, + {mk_pt_coord(1_n), mk_device_id(mc2)}, + }, + }; + + REQUIRE(result.count(repartition_in) == 1); + CHECK(result.at(repartition_in) == expected_in); + } + + SUBCASE("repartition output resolved from node mapping") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const repartition_out = InternalDynamicSlotSite{ + repartition_id, + TensorDirection::OUTPUT, + mk_slot(TensorSlotName::OUTPUT), + }; + + // After repartition: mc1 holds pt1, mc2 holds pt0 + ParallelTensorMapping const expected_out = ParallelTensorMapping{ + bidict{ + {mk_pt_coord(1_n), mk_device_id(mc1)}, + {mk_pt_coord(0_n), mk_device_id(mc2)}, + }, + }; + + REQUIRE(result.count(repartition_out) == 1); + CHECK(result.at(repartition_out) == expected_out); + } + + SUBCASE("repartition input missing from partial — resolved from adjacent") { + // RESHUFFLE resolves only OUTPUT from node mapping. + // INPUT is resolved from adjacent values. + std::map partial = + resolve_partial_tensor_mappings_from_node_mappings(g); + + // relu1 in + relu1 out + repartition out = 3 + // repartition in is missing (resolved from adjacent) + CHECK(partial.size() == 3); + } + } + + TEST_CASE("resolve_tensor_mappings - repartition operator (gather-shaped)") { + // Gather-shaped repartition: 2 input shards (unique per device) collapse + // into 1 output shard (both devices write pt_coord(0) — duplicate). + // This exercises the OUTPUT-has-dup-coords branch of the RESHUFFLE + // handling: here INPUT (not OUTPUT) is the side resolved directly from + // the node mapping, and OUTPUT is instead resolved from the adjacent + // downstream sink (relu2) — the mirror image of the (already-tested, + // above) scatter/shuffle-shaped case, and previously untested. + MachineSpaceCoordinate mc1 = mk_machine_coord(0_n); + MachineSpaceCoordinate mc2 = mk_machine_coord(1_n); + MachineSpaceCoordinate mc3 = mk_machine_coord(2_n); + + auto mk_pt_coord = + [](nonnegative_int idx) -> ParallelTensorSpaceCoordinate { + return ParallelTensorSpaceCoordinate{ + 0_n, + idx, + FFOrdered{0_n, 0_n}, + }; + }; + + auto mk_node_mapping = + [](MappedOperatorTaskGroup const &op_task_group) -> DynamicNodeMapping { + return DynamicNodeMapping{op_task_group, DeviceType::GPU}; + }; + + auto mk_single_mapping = + [&](MachineSpaceCoordinate const &mc, + nonnegative_int pt_idx) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(pt_idx), mk_device_id(mc)}, + }, + }; + }; + + auto mk_two_mapping = + [&](MachineSpaceCoordinate const &mc1_, + MachineSpaceCoordinate const &mc2_) -> ParallelTensorMapping { + return ParallelTensorMapping{ + bidict{ + {mk_pt_coord(0_n), mk_device_id(mc1_)}, + {mk_pt_coord(1_n), mk_device_id(mc2_)}, + }, + }; + }; + + DynamicValueAttrs relu1_output = + mk_value_attrs(401, TensorSlotName::OUTPUT, std::nullopt); + DynamicValueAttrs repartition_output = + mk_value_attrs(402, TensorSlotName::OUTPUT, std::nullopt); + + MappedOperatorTaskGroup relu1_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(1_n)}}}}, + }, + }; + + // OUTPUT collapses both devices onto the same coordinate — duplicate. + MappedOperatorTaskGroup repartition_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + {mc2, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(1_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + MappedOperatorTaskGroup relu2_mapping_group = MappedOperatorTaskGroup{ + bidict{ + {mc3, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::INPUT, mk_pt_coord(0_n)}, + {TensorSlotName::OUTPUT, mk_pt_coord(0_n)}}}}, + }, + }; + + DynamicNodeInvocation relu1_invocation = DynamicNodeInvocation{ + /*inputs=*/{ + {mk_slot(TensorSlotName::INPUT), + mk_value_attrs(400, TensorSlotName::OUTPUT, std::nullopt)}}, + /*node_attrs=*/ + mk_node_attrs(401, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu1_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), relu1_output}}, + }; + + DynamicNodeInvocation repartition_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), relu1_output}}, + /*node_attrs=*/ + mk_node_attrs(402, + PCGOperatorAttrs{RepartitionAttrs{ff_dim_t{0_n}, 2_p}}, + mk_node_mapping(repartition_mapping_group)), + /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), repartition_output}}, + }; + + DynamicNodeInvocation relu2_invocation = DynamicNodeInvocation{ + /*inputs=*/{{mk_slot(TensorSlotName::INPUT), repartition_output}}, + /*node_attrs=*/ + mk_node_attrs(403, + PCGOperatorAttrs{make_relu_attrs()}, + mk_node_mapping(relu2_mapping_group)), + /*outputs=*/ + {{mk_slot(TensorSlotName::OUTPUT), + mk_value_attrs(403, TensorSlotName::OUTPUT, std::nullopt)}}, + }; + + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set( + {relu1_invocation, repartition_invocation, relu2_invocation}); + + dynamic_invocation_id_t repartition_id = + dynamic_graph_get_id_for_invocation(g, repartition_invocation); + + SUBCASE("repartition input resolved from node mapping") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const repartition_in = InternalDynamicSlotSite{ + repartition_id, + TensorDirection::INCOMING, + mk_slot(TensorSlotName::INPUT), + }; + + REQUIRE(result.count(repartition_in) == 1); + CHECK(result.at(repartition_in) == mk_two_mapping(mc1, mc2)); + } + + SUBCASE("repartition output resolved from adjacent relu2 input") { + std::map result = + resolve_tensor_mappings(g); + + InternalDynamicSlotSite const repartition_out = InternalDynamicSlotSite{ + repartition_id, + TensorDirection::OUTPUT, + mk_slot(TensorSlotName::OUTPUT), + }; + + REQUIRE(result.count(repartition_out) == 1); + CHECK(result.at(repartition_out) == mk_single_mapping(mc3, 0_n)); + } + + SUBCASE("repartition output missing from partial — resolved from adjacent") { + // Gather-shaped RESHUFFLE resolves only INPUT from node mapping. + // OUTPUT is resolved from adjacent values (the downstream relu2 sink). + std::map partial = + resolve_partial_tensor_mappings_from_node_mappings(g); + + // relu1 in + relu1 out + repartition in + relu2 in + relu2 out = 5 + // repartition out is missing (resolved from adjacent) + CHECK(partial.size() == 5); + } + } } From 7bd5a6d185f4aaaad8b63097906c2c1ae3421e7e Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 13 Aug 2026 16:15:12 -0700 Subject: [PATCH 4/7] Fix realm-execution CUDA redo file --- .../redops/{realm_redop_registry.cu => realm_redop_registry.cc} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename lib/realm-execution/src/realm-execution/redops/{realm_redop_registry.cu => realm_redop_registry.cc} (100%) diff --git a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc similarity index 100% rename from lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cu rename to lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc From 9e15171724a2068e78da49c16089126df54c7c28 Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 13 Aug 2026 16:16:02 -0700 Subject: [PATCH 5/7] update filename --- lib/realm-execution/CMakeLists.txt | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/lib/realm-execution/CMakeLists.txt b/lib/realm-execution/CMakeLists.txt index 378f28c8e4..149cadbd59 100644 --- a/lib/realm-execution/CMakeLists.txt +++ b/lib/realm-execution/CMakeLists.txt @@ -1,6 +1,3 @@ -# realm_redop_registry.cu needs to be compiled by nvcc (not g++) so that its -# GPU reduction kernels (SumReduction::apply_cuda/fold_cuda) actually get -# generated; see the comment on register_sum_redop for why that matters. project(realm-execution LANGUAGES CXX CUDA) ff_add_library( @@ -8,7 +5,6 @@ ff_add_library( realm-execution SRC_PATTERNS src/*.cc - src/*.cu PUBLIC_INCLUDE include/ PRIVATE_INCLUDE @@ -25,6 +21,19 @@ 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::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 From 5cb613291433422e58aa90de94c0bea03506ed69 Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 13 Aug 2026 16:27:08 -0700 Subject: [PATCH 6/7] formatting updates --- .../redops/realm_redop_registry.cc | 30 ++++++++----------- .../task-spec/dynamic_graph/copy_insertion.cc | 3 +- .../task-spec/dynamic_graph/copy_insertion.cc | 3 +- 3 files changed, 16 insertions(+), 20 deletions(-) diff --git a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc index b137a15ae8..c249406590 100644 --- a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc +++ b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc @@ -151,7 +151,7 @@ __device__ __forceinline__ unsigned int unsigned int const shift = offset * 8; unsigned int const mask = 0xFFu << shift; unsigned int const byte_val = (static_cast(value) & 0xFFu) - << shift; + << shift; return (word & ~mask) | byte_val; } @@ -159,14 +159,14 @@ __device__ __forceinline__ unsigned int // SumReduction 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 \ - __device__ static void apply_cuda(LHS &lhs, RHS rhs) { \ - apply(lhs, rhs); \ - } \ - template \ - __device__ static void fold_cuda(RHS &rhs1, RHS rhs2) { \ - fold(rhs1, rhs2); \ +#define __LEGION_CUDA_REDOP_METHODS__ \ + template \ + __device__ static void apply_cuda(LHS &lhs, RHS rhs) { \ + apply(lhs, rhs); \ + } \ + template \ + __device__ static void fold_cuda(RHS &rhs1, RHS rhs2) { \ + fold(rhs1, rhs2); \ } #else #define __LEGION_CUDA_REDOP_METHODS__ @@ -613,17 +613,13 @@ void register_all_redops() { register_sum_redop>( rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_BOOL_REDOP_ID)); register_sum_redop>( - rt, - static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT32_REDOP_ID)); + rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT32_REDOP_ID)); register_sum_redop>( - rt, - static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT64_REDOP_ID)); + rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_INT64_REDOP_ID)); register_sum_redop>( - rt, - static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_FLOAT_REDOP_ID)); + rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_FLOAT_REDOP_ID)); register_sum_redop>( - rt, - static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_DOUBLE_REDOP_ID)); + rt, static_cast<::Realm::ReductionOpID>(redop_id_t::SUM_DOUBLE_REDOP_ID)); } } // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index b9cffc2370..ab60404d1c 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc @@ -83,8 +83,7 @@ void require_value_is_copy_inserted(DynamicValueAttrs const &v) { // bindings (i.e. cannot form a valid bidict). static bool slot_name_has_dup_coords(DynamicNodeInvocation const &invocation, TensorSlotName const &slot_name) { - DynamicNodeMapping const &nm = - assert_unwrap(invocation.node_attrs.mapping); + DynamicNodeMapping const &nm = assert_unwrap(invocation.node_attrs.mapping); bidict const shard_bindings = dynamic_node_mapping_get_shard_bindings(nm); std::set seen; diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc index b53575aabf..cb2fe0f154 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc @@ -2558,7 +2558,8 @@ TEST_SUITE(FF_TEST_SUITE) { CHECK(result.at(repartition_out) == mk_single_mapping(mc3, 0_n)); } - SUBCASE("repartition output missing from partial — resolved from adjacent") { + SUBCASE( + "repartition output missing from partial — resolved from adjacent") { // Gather-shaped RESHUFFLE resolves only INPUT from node mapping. // OUTPUT is resolved from adjacent values (the downstream relu2 sink). std::map partial = From d8183d0504eed9e7c168ac6558cb4ff5c3717666 Mon Sep 17 00:00:00 2001 From: Seema Mirchandaney Date: Thu, 20 Aug 2026 14:42:39 -0700 Subject: [PATCH 7/7] code cleanup --- .../redops/realm_redop_registry.cc | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc index c249406590..4c2331147d 100644 --- a/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc +++ b/lib/realm-execution/src/realm-execution/redops/realm_redop_registry.cc @@ -1,9 +1,4 @@ #include "realm-execution/redops/realm_redop_registry.h" -// Deliberately raw Realm headers rather than -// "realm-execution/redops/redop_id_t.h"/realm-execution/realm.h: those pull -// in PRealm's prealm.h, which fails a static_assert when parsed by nvcc, and -// we don't need any PRealm-specific behavior here anyway (redop registration -// isn't something PRealm wraps/instruments). #include "realm-execution/redops/redop_id_t.dtg.h" #include #include @@ -580,20 +575,6 @@ __LEGION_CUDA_HD__ inline void SumReduction::fold(RHS &rhs1, } namespace { - -// Building the ReductionOpUntyped by hand (rather than via the -// ::Realm::Runtime::register_reduction convenience template) lets us -// call add_cuda_redop_kernels first, so the redop carries real GPU kernel -// function pointers (REDOP::apply_cuda/fold_cuda) instead of just its CPU -// apply/fold. Without this, Realm's GPUreduceChannel::supports_redop finds -// no CUDA-capable redop and refuses to build any GPU reduction-copy path -// (including same-device ones), aborting with "no path found ... (redop=N)" -// the first time a Reduction op actually runs on a GPU. -// -// Fully qualified as ::Realm:: throughout (rather than the FlexFlow::Realm -// alias to PRealm used elsewhere in this codebase) since this file -// deliberately avoids depending on PRealm at all — see the comment on the -// includes above. template void register_sum_redop(::Realm::Runtime &rt, ::Realm::ReductionOpID id) { ::Realm::ReductionOpUntyped *redop =