From 768c1ebf9fcbca4828b4ef913e72c2c7f5bf379f Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Fri, 5 Jun 2026 01:18:09 -0700 Subject: [PATCH 1/5] Minor include fixes (#1655) --- .../src/compiler/machine_mapping/allowed_machine_views.cc | 2 +- lib/compiler/test/src/compiler/mcmc/generic_mcmc_algorithm.cc | 2 +- lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc | 2 +- .../test/src/compiler/unity_algorithm/graph_optimize_state.cc | 2 +- .../test/src/compiler/unity_algorithm/unity_algorithm.cc | 2 +- lib/op-attrs/include/op-attrs/ff_dim_t.h | 2 +- lib/op-attrs/src/op-attrs/relative_ff_dim_t.cc | 2 +- .../realm-execution/tasks/impl/controller_task_args.dtg.toml | 2 +- .../serializer/serializable_device_specific_ptr.dtg.toml | 4 ++-- lib/utils/include/utils/variant.h | 2 +- lib/utils/test/common/include/test/utils/rapidcheck/doctest.h | 2 +- lib/utils/test/common/include/test/utils/rapidcheck/gen.h | 2 +- .../graph/series_parallel/non_normal_sp_decomposition.cc | 2 +- 13 files changed, 14 insertions(+), 14 deletions(-) diff --git a/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc b/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc index 2a0402a791..d280e929c5 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc @@ -1,5 +1,5 @@ #include "compiler/machine_mapping/allowed_machine_views.h" -#include "doctest/doctest.h" +#include #include "utils/containers/extend.h" #include "utils/containers/range.h" #include "utils/containers/transform.h" diff --git a/lib/compiler/test/src/compiler/mcmc/generic_mcmc_algorithm.cc b/lib/compiler/test/src/compiler/mcmc/generic_mcmc_algorithm.cc index b21ee4333f..d3c96dee92 100644 --- a/lib/compiler/test/src/compiler/mcmc/generic_mcmc_algorithm.cc +++ b/lib/compiler/test/src/compiler/mcmc/generic_mcmc_algorithm.cc @@ -1,5 +1,5 @@ #include "compiler/mcmc/generic_mcmc_algorithm.h" -#include "doctest/doctest.h" +#include using namespace FlexFlow; diff --git a/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc b/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc index 2584f6b3a6..4c03af5382 100644 --- a/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc +++ b/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc @@ -1,6 +1,6 @@ #include "compiler/mcmc/mcmc_over_mapped_pcg.h" #include "compiler/task_graph_simulator/task_simulator.h" -#include "doctest/doctest.h" +#include #include "internal/runtime_only_cost_estimator_for_test.h" #include "op-attrs/parallel_tensor_dims.h" #include "op-attrs/parallel_tensor_shape.dtg.h" diff --git a/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc b/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc index bf7884df29..cc7ca9425f 100644 --- a/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc +++ b/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc @@ -3,7 +3,7 @@ #include "compiler/machine_mapping/machine_mapping.h" #include "compiler/machine_mapping/machine_view.dtg.h" #include "compiler/machine_mapping/machine_view.h" -#include "doctest/doctest.h" +#include #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.h" #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h" #include "test/utils/doctest/check_without_stringify.h" diff --git a/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc b/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc index f5278612aa..0d4123d381 100644 --- a/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc +++ b/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc @@ -1,6 +1,6 @@ #include "compiler/unity_algorithm/unity_algorithm.h" #include "compiler/cost_estimator/runtime_only_cost_estimator_from_cost_estimator.h" -#include "doctest/doctest.h" +#include #include "internal/cost_estimator_for_test.h" #include "op-attrs/parallel_tensor_dims.h" #include "op-attrs/parallel_tensor_shape.dtg.h" diff --git a/lib/op-attrs/include/op-attrs/ff_dim_t.h b/lib/op-attrs/include/op-attrs/ff_dim_t.h index 1411886eee..ae1f3a5cc0 100644 --- a/lib/op-attrs/include/op-attrs/ff_dim_t.h +++ b/lib/op-attrs/include/op-attrs/ff_dim_t.h @@ -3,7 +3,7 @@ #include "op-attrs/ff_dim_t.dtg.h" #include "op-attrs/relative_ff_dim_t.dtg.h" -#include "rapidcheck.h" +#include namespace FlexFlow { diff --git a/lib/op-attrs/src/op-attrs/relative_ff_dim_t.cc b/lib/op-attrs/src/op-attrs/relative_ff_dim_t.cc index 91caa03f36..1dcc8a845e 100644 --- a/lib/op-attrs/src/op-attrs/relative_ff_dim_t.cc +++ b/lib/op-attrs/src/op-attrs/relative_ff_dim_t.cc @@ -1,5 +1,5 @@ #include "op-attrs/relative_ff_dim_t.h" -#include "rapidcheck.h" +#include namespace FlexFlow { ff_dim_t ff_dim_t_from_relative_ff_dim_t(relative_ff_dim_t ff_dim, diff --git a/lib/realm-execution/include/realm-execution/tasks/impl/controller_task_args.dtg.toml b/lib/realm-execution/include/realm-execution/tasks/impl/controller_task_args.dtg.toml index 0c0bd7b96c..9b63503aff 100644 --- a/lib/realm-execution/include/realm-execution/tasks/impl/controller_task_args.dtg.toml +++ b/lib/realm-execution/include/realm-execution/tasks/impl/controller_task_args.dtg.toml @@ -5,7 +5,7 @@ features = [] includes = [ "realm-execution/realm_context.h", - "functional", + "", ] [[fields]] diff --git a/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml b/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml index 07cf61f7e1..6a5675772e 100644 --- a/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml +++ b/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml @@ -10,8 +10,8 @@ features = [ includes = [ "pcg/device_id_t.dtg.h", - "cstdint", - "optional", + "", + "", ] src_includes = [ diff --git a/lib/utils/include/utils/variant.h b/lib/utils/include/utils/variant.h index c5947413e1..ba689d57be 100644 --- a/lib/utils/include/utils/variant.h +++ b/lib/utils/include/utils/variant.h @@ -1,7 +1,7 @@ #ifndef _FLEXFLOW_UTILS_VARIANT_H #define _FLEXFLOW_UTILS_VARIANT_H -#include "rapidcheck.h" +#include #include "utils/type_traits.h" #include #include diff --git a/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h b/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h index 0121f03926..15ea9fc663 100644 --- a/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h +++ b/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h @@ -1,7 +1,7 @@ #ifndef _FLEXFLOW_UTILS_TEST_COMMON_INCLUDE_TEST_UTILS_RAPIDCHECK_DOCTEST_H #define _FLEXFLOW_UTILS_TEST_COMMON_INCLUDE_TEST_UTILS_RAPIDCHECK_DOCTEST_H -#include "rapidcheck.h" +#include #include namespace FlexFlow { diff --git a/lib/utils/test/common/include/test/utils/rapidcheck/gen.h b/lib/utils/test/common/include/test/utils/rapidcheck/gen.h index 4fd7c0e069..ad9879d6e7 100644 --- a/lib/utils/test/common/include/test/utils/rapidcheck/gen.h +++ b/lib/utils/test/common/include/test/utils/rapidcheck/gen.h @@ -1,7 +1,7 @@ #ifndef _FLEXFLOW_UTILS_LIB_TEST_COMMON_INCLUDE_UTILS_TEST_RAPIDCHECK_GEN_H #define _FLEXFLOW_UTILS_LIB_TEST_COMMON_INCLUDE_UTILS_TEST_RAPIDCHECK_GEN_H -#include "rapidcheck.h" +#include #include namespace rc { diff --git a/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc b/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc index 5f41d77305..4d91d13c05 100644 --- a/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc +++ b/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc @@ -1,5 +1,5 @@ #include "utils/graph/series_parallel/non_normal_sp_decomposition.h" -#include "doctest/doctest.h" +#include #include "utils/graph/series_parallel/series_parallel_decomposition.dtg.h" using namespace ::FlexFlow; From bfbb318dc0414d4d0b3ca9d3b62a92b1cbdacc3a Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Fri, 29 May 2026 17:43:32 -0700 Subject: [PATCH 2/5] Remove DeviceType from MachineSpaceCoord, move task-spec to use device_id_t, flesh out copy_insertion logic --- .../cost_estimator/op_cost_estimate_key.h | 1 - .../machine_mapping/allowed_machine_views.h | 3 +- .../machine_mapping/machine_mapping.h | 1 - .../machine_mapping_mutation_set.h | 6 +- .../compiler/machine_mapping/machine_view.h | 8 - .../start_invariant_machine_view.dtg.toml | 6 - .../start_invariant_machine_view.h | 2 - .../unstructured_device_mapping.dtg.toml | 27 - .../mcmc/mcmc_over_mapped_pcg_config.dtg.toml | 5 - .../pcg_task_graph.dtg.toml | 4 +- .../cost_estimator/op_cost_estimate_key.cc | 1 - .../machine_mapping/allowed_machine_views.cc | 16 +- .../machine_mapping_mutation_set.cc | 10 +- .../machine_mapping/machine_resource_split.cc | 2 - .../compiler/machine_mapping/machine_view.cc | 19 +- .../start_invariant_machine_view.cc | 13 +- .../unstructured_device_mapping.cc | 26 - .../src/compiler/mcmc/mcmc_over_mapped_pcg.cc | 4 +- .../task_graph_simulator/pcg_task_graph.cc | 8 +- .../task_graph_simulator/task_simulator.cc | 11 +- .../unity_algorithm/unity_algorithm.cc | 2 +- .../machine_mapping/allowed_machine_views.cc | 9 +- .../get_optimal_machine_mapping.cc | 3 - .../get_tensor_set_movement_across_split.cc | 6 - .../machine_mapping/machine_mapping.cc | 4 - .../machine_mapping/machine_mapping_result.cc | 7 - .../compiler/machine_mapping/machine_view.cc | 148 +--- ...get_optimal_machine_mapping_with_memory.cc | 5 - .../machine_mapping_with_memory_result.cc | 8 - .../start_invariant_machine_view.cc | 47 +- .../src/compiler/mcmc/mcmc_over_mapped_pcg.cc | 6 +- .../task_graph_simulator/task_simulator.cc | 16 +- .../computation_graph_instance.h | 2 +- .../cost_estimator/local_cost_estimator.h | 2 +- .../local_task_argument_accessor.h | 2 +- .../per_device_op_state_initialization.h | 2 +- .../cost_estimator/local_cost_estimator.cc | 3 +- .../local_task_argument_accessor.cc | 4 +- .../computation_graph_instance.cc | 30 +- .../cost_estimator/local_cost_estimator.cc | 24 +- .../local_task_argument_accessor.cc | 10 +- lib/pcg/include/pcg/device_id.h | 22 - lib/pcg/include/pcg/device_id_t.dtg.toml | 23 - lib/pcg/include/pcg/device_id_t.h | 14 - .../machine_compute_resource_slice.dtg.toml | 2 + .../pcg/machine_compute_specification.h | 7 - .../pcg/machine_space_coordinate.dtg.toml | 7 +- .../include/pcg/machine_space_offset.dtg.toml | 8 +- lib/pcg/src/pcg/device_id.cc | 51 -- lib/pcg/src/pcg/device_id_t.cc | 17 - .../src/pcg/machine_compute_resource_slice.cc | 1 - .../src/pcg/machine_compute_specification.cc | 17 - lib/pcg/src/pcg/machine_space_offset.cc | 3 - .../src/pcg/machine_compute_specification.cc | 23 - .../mapped_operator_task_group.cc | 23 +- .../mapped_parallel_computation_graph.cc | 1 - ...ce_specific_managed_per_device_ff_handle.h | 2 +- .../realm-execution/device_specific_ptr.h | 3 +- .../realm-execution/distributed_ff_handle.h | 6 +- ...buted_per_device_op_state_initialization.h | 3 +- .../realm-execution/instance_allocation.h | 10 +- .../include/realm-execution/pcg_instance.h | 5 +- .../include/realm-execution/processor_kind.h | 14 + .../include/realm-execution/realm_context.h | 17 +- .../include/realm-execution/realm_manager.h | 2 +- ...uted_per_device_op_state_initialization.cc | 3 +- .../realm-execution/instance_allocation.cc | 16 +- .../src/realm-execution/pcg_instance.cc | 42 +- .../src/realm-execution/processor_kind.cc | 29 + .../src/realm-execution/realm_allocator.cc | 10 +- .../src/realm-execution/realm_context.cc | 98 +-- .../src/realm-execution/realm_manager.cc | 1 + .../tasks/impl/controller_task.cc | 6 +- .../test/src/realm-execution/test_e2e.cc | 529 +++++-------- .../include/task-spec/concrete_arg_spec.h | 1 - .../include/task-spec/device_id_t.dtg.toml | 23 + .../include/task-spec/device_specific.h | 16 +- .../dynamic_graph/dynamic_graph_edge.dtg.toml | 29 + .../dynamic_graph/dynamic_graph_edge.h | 15 + .../dynamic_graph/dynamic_node_attrs.dtg.toml | 13 +- .../dynamic_graph/dynamic_node_invocation.h | 22 + .../dynamic_node_mapping.dtg.toml | 26 + .../dynamic_graph/dynamic_node_mapping.h | 18 + .../dynamic_open_dataflow_graph.h | 27 + .../dynamic_graph/dynamic_slot_site.dtg.toml | 21 + .../dynamic_graph/dynamic_slot_site.h | 13 + .../dynamic_value_attrs.dtg.toml | 9 +- .../dynamic_graph/dynamic_value_attrs.h | 5 + .../external_dynamic_slot_site.dtg.toml | 19 + .../internal_dynamic_slot_site.dtg.toml | 29 + .../task-spec/dynamic_graph/loss_insertion.h | 2 +- .../task-spec/dynamic_graph/machine_slicing.h | 4 +- ...amic_open_dataflow_graph_from_mapped_pcg.h | 2 +- .../parallel_tensor_mapping.dtg.toml | 22 + .../serializable_dynamic_node_attrs.dtg.toml | 8 +- .../serializable_dynamic_value_attrs.dtg.toml | 3 +- .../training_only_op_type.dtg.toml | 15 + .../dynamic_graph/training_op_type.dtg.toml | 22 + .../dynamic_graph/training_operation_attrs.h | 13 + .../include/task-spec/serialization.h | 113 --- .../itask_argument_accessor.h | 2 +- .../task-spec/dynamic_graph/copy_insertion.cc | 263 ++++--- .../dynamic_graph/dynamic_graph_edge.cc | 20 + .../dynamic_graph/dynamic_node_invocation.cc | 62 ++ .../dynamic_graph/dynamic_node_mapping.cc | 31 + .../dynamic_open_dataflow_graph.cc | 181 ++++- .../dynamic_graph/dynamic_slot_site.cc | 26 + .../dynamic_graph/dynamic_value_attrs.cc | 9 + .../task-spec/dynamic_graph/loss_insertion.cc | 2 +- .../dynamic_graph/machine_slicing.cc | 4 +- ...mic_open_dataflow_graph_from_mapped_pcg.cc | 8 +- .../dynamic_graph/shard_expansion.cc | 39 +- .../dynamic_graph/training_operation_attrs.cc | 24 + .../dynamic_graph/update_insertion.cc | 4 +- lib/task-spec/src/task-spec/serialization.cc | 1 - lib/task-spec/test/CMakeLists.txt | 1 - .../test/src/task-spec/device_specific.cc | 10 +- .../task-spec/dynamic_graph/copy_insertion.cc | 731 ++++++++++-------- .../dynamic_open_dataflow_graph.cc | 501 ++++++++++-- .../dynamic_graph/machine_slicing.cc | 21 +- ...mic_open_dataflow_graph_from_mapped_pcg.cc | 1 + .../dynamic_graph/shard_expansion.cc | 508 ++++++------ .../dynamic_graph/update_insertion.cc | 157 ++++ 123 files changed, 2679 insertions(+), 1945 deletions(-) delete mode 100644 lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.dtg.toml delete mode 100644 lib/compiler/src/compiler/machine_mapping/unstructured_device_mapping.cc delete mode 100644 lib/pcg/include/pcg/device_id.h delete mode 100644 lib/pcg/include/pcg/device_id_t.dtg.toml delete mode 100644 lib/pcg/include/pcg/device_id_t.h delete mode 100644 lib/pcg/src/pcg/device_id.cc delete mode 100644 lib/pcg/src/pcg/device_id_t.cc create mode 100644 lib/realm-execution/include/realm-execution/processor_kind.h create mode 100644 lib/realm-execution/src/realm-execution/processor_kind.cc create mode 100644 lib/task-spec/include/task-spec/device_id_t.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.h create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_invocation.h create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.h create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/external_dynamic_slot_site.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/training_op_type.dtg.toml create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.h delete mode 100644 lib/task-spec/include/task-spec/serialization.h create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_graph_edge.cc create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_invocation.cc create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_slot_site.cc create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc delete mode 100644 lib/task-spec/src/task-spec/serialization.cc create mode 100644 lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc create mode 100644 lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc diff --git a/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h b/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h index d905abeb77..dae6989da8 100644 --- a/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h +++ b/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h @@ -3,7 +3,6 @@ #include "compiler/cost_estimator/op_cost_estimate_key.dtg.h" #include "compiler/cost_estimator/runtime_only_op_cost_estimate_key.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/machine_specification.dtg.h" #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h" #include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h" diff --git a/lib/compiler/include/compiler/machine_mapping/allowed_machine_views.h b/lib/compiler/include/compiler/machine_mapping/allowed_machine_views.h index 5201f7fa31..de899ab74d 100644 --- a/lib/compiler/include/compiler/machine_mapping/allowed_machine_views.h +++ b/lib/compiler/include/compiler/machine_mapping/allowed_machine_views.h @@ -14,8 +14,7 @@ bool is_valid_machine_view(MachineView const &mv, std::unordered_set get_allowed_machine_views(MachineComputeResourceSlice const &machine_spec, - OperatorTaskSpace const &task, - DeviceType device_type); + OperatorTaskSpace const &task); } // namespace FlexFlow diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping.h b/lib/compiler/include/compiler/machine_mapping/machine_mapping.h index 28d1fc943b..7ace9989fb 100644 --- a/lib/compiler/include/compiler/machine_mapping/machine_mapping.h +++ b/lib/compiler/include/compiler/machine_mapping/machine_mapping.h @@ -4,7 +4,6 @@ #include "compiler/machine_mapping/machine_mapping.dtg.h" #include "compiler/machine_mapping/machine_mapping_result.h" #include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/machine_specification.dtg.h" #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h" #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h" diff --git a/lib/compiler/include/compiler/machine_mapping/machine_mapping_mutation_set.h b/lib/compiler/include/compiler/machine_mapping/machine_mapping_mutation_set.h index d05e3fab7c..005e2a1ee5 100644 --- a/lib/compiler/include/compiler/machine_mapping/machine_mapping_mutation_set.h +++ b/lib/compiler/include/compiler/machine_mapping/machine_mapping_mutation_set.h @@ -7,13 +7,11 @@ namespace FlexFlow { std::optional get_random_mapping(ParallelComputationGraph const &pcg, - MachineComputeSpecification const &resources, - DeviceType const &device_type); + MachineComputeSpecification const &resources); std::optional get_random_mutation(SearchResult const &mapped_pcg, - MachineComputeSpecification const &resource, - DeviceType const &device_type); + MachineComputeSpecification const &resource); } // namespace FlexFlow #endif diff --git a/lib/compiler/include/compiler/machine_mapping/machine_view.h b/lib/compiler/include/compiler/machine_mapping/machine_view.h index 6888fa6b94..4a9e12215c 100644 --- a/lib/compiler/include/compiler/machine_mapping/machine_view.h +++ b/lib/compiler/include/compiler/machine_mapping/machine_view.h @@ -6,7 +6,6 @@ #include "op-attrs/operator_task_space.dtg.h" #include "op-attrs/parallel_tensor_dim_degrees.dtg.h" #include "op-attrs/task_space_coordinate.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/machine_compute_specification.dtg.h" #include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h" #include "pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.dtg.h" @@ -20,8 +19,6 @@ namespace FlexFlow { nonnegative_int mv_get_expected_task_space_num_dims(MachineView const &mv); -DeviceType get_device_type(MachineView const &mv); - std::vector get_strides(MachineView const &mv); std::vector @@ -50,11 +47,6 @@ std::unordered_set get_machine_space_coordinates(OperatorTaskSpace const &task, MachineView const &mv); -std::unordered_set - get_device_ids(OperatorTaskSpace const &task, - MachineView const &mv, - MachineComputeSpecification const &ms); - MachineView make_1d_machine_view(MachineSpaceCoordinate const &start, MachineSpecificationDimension const &dim, stride_t stride); diff --git a/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.dtg.toml b/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.dtg.toml index b271ed1959..7a0460a3b6 100644 --- a/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.dtg.toml +++ b/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.dtg.toml @@ -12,7 +12,6 @@ features = [ includes = [ "compiler/machine_mapping/machine_view_dimension.dtg.h", - "pcg/device_type.dtg.h" ] src_includes = [ @@ -23,8 +22,3 @@ src_includes = [ [[fields]] name = "dimensions" type = "std::vector<::FlexFlow::MachineViewDimension>" - - -[[fields]] -name = "device_type" -type = "::FlexFlow::DeviceType" diff --git a/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.h b/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.h index 631d23b07c..7b5a4f57f0 100644 --- a/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.h +++ b/lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.h @@ -19,8 +19,6 @@ StartInvariantMachineView nonnegative_int num_dims(StartInvariantMachineView const &mv); -DeviceType get_device_type(StartInvariantMachineView const &mv); - std::vector get_strides(StartInvariantMachineView const &mv); std::vector diff --git a/lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.dtg.toml b/lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.dtg.toml deleted file mode 100644 index 28391eddc0..0000000000 --- a/lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.dtg.toml +++ /dev/null @@ -1,27 +0,0 @@ -namespace = "FlexFlow" -name = "UnstructuredDeviceMapping" -type = "struct" -features = [ - "eq", - # "ord", - "hash", - # "json", - # "rapidcheck", - "fmt", -] - -includes = [ - "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h", - "pcg/device_id_t.dtg.h" -] - -src_includes = [ - "utils/hash/unordered_map.h", - "utils/fmt/unordered_map.h", - "utils/hash/unordered_set.h", - "utils/fmt/unordered_set.h" -] - -[[fields]] -name = "raw_device_map" -type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, std::unordered_set<::FlexFlow::device_id_t>>" diff --git a/lib/compiler/include/compiler/mcmc/mcmc_over_mapped_pcg_config.dtg.toml b/lib/compiler/include/compiler/mcmc/mcmc_over_mapped_pcg_config.dtg.toml index 99320f735e..cf8476e882 100644 --- a/lib/compiler/include/compiler/mcmc/mcmc_over_mapped_pcg_config.dtg.toml +++ b/lib/compiler/include/compiler/mcmc/mcmc_over_mapped_pcg_config.dtg.toml @@ -8,7 +8,6 @@ features = [ ] includes = [ - "pcg/device_type.dtg.h", "utils/nonnegative_int/nonnegative_int.h" ] @@ -23,7 +22,3 @@ type = "::FlexFlow::nonnegative_int" [[fields]] name = "substitution_frequency" type = "float" - -[[fields]] -name = "device_type" -type = "::FlexFlow::DeviceType" \ No newline at end of file diff --git a/lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.dtg.toml b/lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.dtg.toml index 2c5b5f56fc..14de2d9455 100644 --- a/lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.dtg.toml +++ b/lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.dtg.toml @@ -9,8 +9,8 @@ includes = [ "utils/graph/digraph/digraph_view.h", "utils/bidict/bidict.h", "compiler/task_graph_simulator/pcg_task.dtg.h", - "pcg/device_id_t.dtg.h", "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h", + "pcg/machine_space_coordinate.dtg.h", "", "" ] @@ -32,4 +32,4 @@ type = "::FlexFlow::bidict<::FlexFlow::Node, ::FlexFlow::PCGTask>" [[fields]] name = "node_to_devices" -type = "std::unordered_map<::FlexFlow::Node, std::unordered_set<::FlexFlow::device_id_t>>" +type = "std::unordered_map<::FlexFlow::Node, std::unordered_set<::FlexFlow::MachineSpaceCoordinate>>" diff --git a/lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc b/lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc index f3edd6a69a..62b645e3da 100644 --- a/lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc +++ b/lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc @@ -4,7 +4,6 @@ #include "compiler/machine_mapping/machine_view.dtg.h" #include "compiler/machine_mapping/machine_view.h" #include "op-attrs/parallel_tensor_shape.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/machine_specification.dtg.h" #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h" #include diff --git a/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc b/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc index c3d9ae7bfb..9194c2e982 100644 --- a/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc +++ b/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc @@ -49,8 +49,7 @@ bool is_valid_machine_view(MachineView const &mv, */ static std::unordered_set get_candidate_machine_views(MachineComputeResourceSlice const &machine_spec, - OperatorTaskSpace const &task_space, - DeviceType const &device_type) { + OperatorTaskSpace const &task_space) { auto get_max_stride_upper_bound = [](std::vector const &tensor_dims, @@ -90,17 +89,15 @@ static std::unordered_set return strides; }; - auto get_candidate_starts = [](MachineComputeResourceSlice const &slice, - DeviceType const &device_type) + auto get_candidate_starts = [](MachineComputeResourceSlice const &slice) -> std::unordered_set { - ASSERT(device_type == DeviceType::GPU); std::unordered_set result; for (nonnegative_int node_idx : nonnegative_range(slice.num_nodes)) { for (nonnegative_int device_idx : nonnegative_range(slice.num_gpus_per_node)) { result.insert( - MachineSpaceCoordinate{node_idx, device_idx, device_type}); + MachineSpaceCoordinate{node_idx, device_idx}); } } return result; @@ -127,7 +124,7 @@ static std::unordered_set ASSERT(candidate_strides.size() > 0); std::unordered_set candidate_starts = - get_candidate_starts(machine_spec, device_type); + get_candidate_starts(machine_spec); ASSERT(candidate_starts.size() > 0); std::unordered_multiset> @@ -151,11 +148,10 @@ static std::unordered_set std::unordered_set get_allowed_machine_views(MachineComputeResourceSlice const &machine_spec, - OperatorTaskSpace const &task_space, - DeviceType device_type) { + OperatorTaskSpace const &task_space) { std::unordered_set views = - get_candidate_machine_views(machine_spec, task_space, device_type); + get_candidate_machine_views(machine_spec, task_space); return filter(views, [&](MachineView const &mv) { return is_valid_machine_view(mv, task_space, machine_spec); }); diff --git a/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc b/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc index 47639ff88a..d6cdca97d1 100644 --- a/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc +++ b/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc @@ -11,15 +11,14 @@ namespace FlexFlow { std::optional get_random_mapping(ParallelComputationGraph const &pcg, - MachineComputeSpecification const &resources, - DeviceType const &device_type) { + MachineComputeSpecification const &resources) { std::vector layers = topological_ordering(pcg); std::unordered_map machine_views; for (parallel_layer_guid_t layer : layers) { OperatorTaskSpace task = get_operator_task_space(pcg, layer); std::unordered_set allowed_machine_views = get_allowed_machine_views( - compute_slice_from_specification(resources), task, DeviceType::GPU); + compute_slice_from_specification(resources), task); if (allowed_machine_views.empty()) { return std::nullopt; } @@ -31,8 +30,7 @@ std::optional std::optional get_random_mutation(SearchResult const &mapped_pcg, - MachineComputeSpecification const &resources, - DeviceType const &device_type) { + MachineComputeSpecification const &resources) { ParallelComputationGraph pcg = mapped_pcg.pcg; std::vector layers = topological_ordering(pcg); if (layers.size() == 0) { @@ -46,7 +44,7 @@ std::optional std::vector allowed_machine_views = vector_of(get_allowed_machine_views( - compute_slice_from_specification(resources), task, device_type)); + compute_slice_from_specification(resources), task)); MachineView random_new_machine_view = select_random(allowed_machine_views); machine_mapping.machine_views.at(random_layer) = random_new_machine_view; diff --git a/lib/compiler/src/compiler/machine_mapping/machine_resource_split.cc b/lib/compiler/src/compiler/machine_mapping/machine_resource_split.cc index 875f44a0c9..370b764a69 100644 --- a/lib/compiler/src/compiler/machine_mapping/machine_resource_split.cc +++ b/lib/compiler/src/compiler/machine_mapping/machine_resource_split.cc @@ -88,7 +88,6 @@ MachineSpaceCoordinate /*node_idx=*/(coord.node_idx + split.offset) .nonnegative_int_from_positive_int(), /*device_idx=*/coord.device_idx, - /*device_type=*/coord.device_type, }; } else { ASSERT(split.dimension == MachineSpecificationDimension::INTRA_NODE); @@ -97,7 +96,6 @@ MachineSpaceCoordinate /*node_idx=*/coord.node_idx, /*device_idx=*/ (coord.device_idx + split.offset).nonnegative_int_from_positive_int(), - /*device_type=*/coord.device_type, }; } } diff --git a/lib/compiler/src/compiler/machine_mapping/machine_view.cc b/lib/compiler/src/compiler/machine_mapping/machine_view.cc index 090dec5845..5c38a66901 100644 --- a/lib/compiler/src/compiler/machine_mapping/machine_view.cc +++ b/lib/compiler/src/compiler/machine_mapping/machine_view.cc @@ -33,10 +33,6 @@ nonnegative_int mv_get_expected_task_space_num_dims(MachineView const &mv) { return num_elements(get_strides(mv)); } -DeviceType get_device_type(MachineView const &mv) { - return mv.start.device_type; -} - std::vector get_strides(MachineView const &mv) { return transform(mv.dimensions, [](MachineViewDimension const &dim) { return dim.stride; }); @@ -130,7 +126,7 @@ MachineSpaceCoordinate nonnegative_int device_idx = compute_index(machine_view.start.device_idx, intra_dimension_indices); MachineSpaceCoordinate ms_coord = MachineSpaceCoordinate{ - node_idx, device_idx, get_device_type(machine_view)}; + node_idx, device_idx}; return ms_coord; } @@ -177,19 +173,6 @@ std::unordered_set }); } -std::unordered_set - get_device_ids(OperatorTaskSpace const &task_space, - MachineView const &mv, - MachineComputeSpecification const &ms) { - ASSERT(op_task_space_num_dims(task_space) == - mv_get_expected_task_space_num_dims(mv)); - - return transform(get_machine_space_coordinates(task_space, mv), - [&](MachineSpaceCoordinate const &coord) { - return get_device_id(ms, coord); - }); -} - MachineView make_1d_machine_view(MachineSpaceCoordinate const &start, MachineSpecificationDimension const &dim, stride_t stride) { diff --git a/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc b/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc index cbb64d5bcf..4a2d66acc1 100644 --- a/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc +++ b/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc @@ -18,17 +18,13 @@ MachineView machine_view_from_start_invariant( StartInvariantMachineView start_invariant_from_machine_view(MachineView const &mv) { - return StartInvariantMachineView{mv.dimensions, get_device_type(mv)}; + return StartInvariantMachineView{mv.dimensions}; } nonnegative_int num_dims(StartInvariantMachineView const &start_inv_mv) { return num_elements(start_inv_mv.dimensions); } -DeviceType get_device_type(StartInvariantMachineView const &start_inv_mv) { - return start_inv_mv.device_type; -} - std::vector get_strides(StartInvariantMachineView const &start_inv_mv) { return transform(start_inv_mv.dimensions, @@ -45,13 +41,12 @@ std::vector StartInvariantMachineView start_invariant_machine_view_from_strides_and_machine_spec_dimensions( std::vector const &strides, - std::vector const &dims, - DeviceType device_type) { + std::vector const &dims) { std::vector dimensions = transform(zip(strides, dims), [&](auto const &p) { return MachineViewDimension{p.first, p.second}; }); - return StartInvariantMachineView{dimensions, device_type}; + return StartInvariantMachineView{dimensions}; } MachineSpaceOffset get_machine_space_offset( @@ -60,7 +55,7 @@ MachineSpaceOffset get_machine_space_offset( TaskSpaceCoordinate const &coord) { MachineSpaceCoordinate dummy_start = - MachineSpaceCoordinate{0_n, 0_n, get_device_type(start_inv_machine_view)}; + MachineSpaceCoordinate{0_n, 0_n}; MachineView mv = machine_view_from_start_invariant(start_inv_machine_view, dummy_start); diff --git a/lib/compiler/src/compiler/machine_mapping/unstructured_device_mapping.cc b/lib/compiler/src/compiler/machine_mapping/unstructured_device_mapping.cc deleted file mode 100644 index 80c09d2dba..0000000000 --- a/lib/compiler/src/compiler/machine_mapping/unstructured_device_mapping.cc +++ /dev/null @@ -1,26 +0,0 @@ -#include "compiler/machine_mapping/unstructured_device_mapping.h" -#include "compiler/machine_mapping/machine_view.h" -#include "compiler/machine_mapping/unstructured_device_mapping.dtg.h" -#include "op-attrs/operator_task_space.dtg.h" -#include "op-attrs/operator_task_space.h" -#include "pcg/parallel_computation_graph/parallel_computation_graph.h" -#include "utils/containers/keys.h" -#include "utils/containers/map_values.h" - -namespace FlexFlow { - -UnstructuredDeviceMapping get_unstructured_device_mapping( - MachineMapping const &machine_mapping, - MachineComputeSpecification const &machine_spec, - ParallelComputationGraph const &pcg) { - std::unordered_map> - device_mapping; - for (auto const &[layer, machine_view] : machine_mapping.machine_views) { - OperatorTaskSpace op = get_operator_task_space(pcg, layer); - device_mapping.insert( - {layer, get_device_ids(op, machine_view, machine_spec)}); - } - return UnstructuredDeviceMapping{device_mapping}; -} - -} // namespace FlexFlow diff --git a/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc b/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc index 583a60b1ad..0d2c1e4455 100644 --- a/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc +++ b/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc @@ -22,7 +22,7 @@ SearchResult MachineComputeSpecification compute_spec = machine_spec.compute_specification; std::vector substitutions = get_substitution_set(compute_spec); MachineMapping random_mapping = assert_unwrap( - get_random_mapping(pcg, compute_spec, search_config.device_type)); + get_random_mapping(pcg, compute_spec)); SearchResult starting_state = SearchResult{pcg, random_mapping}; auto sampler = [&](SearchResult mapped_pcg) -> std::optional { @@ -43,7 +43,7 @@ SearchResult }); } else { MachineMapping new_machine_mapping = assert_unwrap(get_random_mutation( - mapped_pcg, compute_spec, search_config.device_type)); + mapped_pcg, compute_spec)); return SearchResult{mapped_pcg.pcg, new_machine_mapping}; } }; diff --git a/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc b/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc index d4d5a78d6a..058f5b72bb 100644 --- a/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc +++ b/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc @@ -5,7 +5,6 @@ #include "compiler/machine_mapping/machine_view.dtg.h" #include "compiler/machine_mapping/machine_view.h" #include "op-attrs/operator_task_space.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/machine_specification.dtg.h" #include "pcg/parallel_computation_graph/parallel_computation_graph.h" #include "pcg/parallel_computation_graph/parallel_computation_graph_edge.dtg.h" @@ -25,7 +24,7 @@ PCGTaskGraph DiGraph digraph = DiGraph::create(); bidict node_to_task; bidict node_to_layer; - std::unordered_map> node_to_devices; + std::unordered_map> node_to_devices; for (parallel_layer_guid_t const &layer : get_parallel_layers(pcg)) { MachineView mv = machine_mapping.machine_views.at(layer); @@ -35,9 +34,8 @@ PCGTaskGraph node_to_task.equate(node, PCGTask{op_key}); node_to_layer.equate(node, layer); node_to_devices[node] = - get_device_ids(get_operator_task_space(pcg, layer), - machine_mapping.machine_views.at(layer), - machine_spec); + get_machine_space_coordinates(get_operator_task_space(pcg, layer), + machine_mapping.machine_views.at(layer)); } for (ParallelComputationGraphEdge const &edge : get_edges(pcg)) { diff --git a/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc b/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc index e514e6a753..28a6a3efae 100644 --- a/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc +++ b/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc @@ -1,8 +1,6 @@ #include "compiler/task_graph_simulator/task_simulator.h" #include "compiler/cost_estimator/cost_estimator.h" #include "compiler/cost_estimator/op_cost_estimate_key.h" -#include "compiler/machine_mapping/unstructured_device_mapping.dtg.h" -#include "compiler/machine_mapping/unstructured_device_mapping.h" #include "compiler/task_graph_simulator/pcg_task.dtg.h" #include "compiler/task_graph_simulator/pcg_task_graph.h" #include "compiler/task_graph_simulator/simulate_task_graph_execution.h" @@ -49,21 +47,18 @@ milliseconds_t task_simulator_estimate_forward_pass_time( std::unordered_set const &finished_tasks) -> bool { PCGTask current_task = task_graph.node_to_task.at_l(task); - UnstructuredDeviceMapping device_map = get_unstructured_device_mapping( - machine_mapping, machine_spec.compute_specification, pcg); - if (current_task.is_tensor_movement()) { return true; } assert(current_task.is_operator()); - auto get_devices = [&](Node const &n) { + auto get_devices = [&](Node const &n) -> std::unordered_set { return task_graph.node_to_devices.at(n); }; - std::unordered_set devices_occupied = + std::unordered_set devices_occupied = set_union(transform(in_progress_tasks, get_devices)); - std::unordered_set required_devices = get_devices(task); + std::unordered_set required_devices = get_devices(task); return set_intersection(devices_occupied, required_devices).empty(); }; diff --git a/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc b/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc index be8c7c4f98..240ad0b1c4 100644 --- a/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc +++ b/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc @@ -68,7 +68,7 @@ SearchResult graph_optimize(ParallelComputationGraph &pcg, get_operator_task_space_for_runtime_only_op_cost_estimate_key(key); return get_allowed_machine_views( - resources, op_task_space, DeviceType::GPU); + resources, op_task_space); }, }; diff --git a/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc b/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc index d280e929c5..6a867b16f3 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc @@ -40,7 +40,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate{ start_node_idx, start_device_idx, - DeviceType::GPU, }, strides, }; @@ -65,7 +64,7 @@ TEST_SUITE(FF_TEST_SUITE) { }; std::unordered_set result = - get_allowed_machine_views(ms, task, DeviceType::GPU); + get_allowed_machine_views(ms, task); CHECK(correct == result); } @@ -96,7 +95,7 @@ TEST_SUITE(FF_TEST_SUITE) { }; std::unordered_set result = - get_allowed_machine_views(ms, task, DeviceType::GPU); + get_allowed_machine_views(ms, task); CHECK(correct == result); } @@ -110,7 +109,7 @@ TEST_SUITE(FF_TEST_SUITE) { OperatorTaskSpace task = OperatorTaskSpace{MinimalOrthotope{{}}}; std::unordered_set result = - get_allowed_machine_views(full_machine_spec, task, DeviceType::GPU); + get_allowed_machine_views(full_machine_spec, task); std::unordered_set correct = { make_machine_view(0_n, 0_n), @@ -129,7 +128,7 @@ TEST_SUITE(FF_TEST_SUITE) { OperatorTaskSpace task = OperatorTaskSpace{MinimalOrthotope{{2_ge2}}}; std::unordered_set result = - get_allowed_machine_views(full_machine_spec, task, DeviceType::GPU); + get_allowed_machine_views(full_machine_spec, task); std::unordered_set correct = { make_machine_view(0_n, 0_n, /*stride_1=*/1_p, intra), diff --git a/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc b/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc index 392e16bec5..59ed83d6fa 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc @@ -53,7 +53,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -68,7 +67,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -557,7 +555,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/2_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { diff --git a/lib/compiler/test/src/compiler/machine_mapping/get_tensor_set_movement_across_split.cc b/lib/compiler/test/src/compiler/machine_mapping/get_tensor_set_movement_across_split.cc index b3901e08ca..e2a87204f0 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/get_tensor_set_movement_across_split.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/get_tensor_set_movement_across_split.cc @@ -93,7 +93,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -108,7 +107,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/1_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -123,7 +121,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/2_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -138,7 +135,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/3_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -160,13 +156,11 @@ TEST_SUITE(FF_TEST_SUITE) { /*src=*/MachineSpaceCoordinate{ /*node_idx=*/src_mv.start.node_idx, /*device_idx=*/src_task_idx, - /*device_type=*/DeviceType::GPU, }, /*dst=*/ MachineSpaceCoordinate{ /*node_idx=*/dst_mv.start.node_idx, /*device_idx=*/dst_task_idx, - /*device_type=*/DeviceType::GPU, }, }; }; diff --git a/lib/compiler/test/src/compiler/machine_mapping/machine_mapping.cc b/lib/compiler/test/src/compiler/machine_mapping/machine_mapping.cc index 8af07a032c..6cae9a2e50 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/machine_mapping.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/machine_mapping.cc @@ -11,7 +11,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -26,7 +25,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -57,7 +55,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -72,7 +69,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { diff --git a/lib/compiler/test/src/compiler/machine_mapping/machine_mapping_result.cc b/lib/compiler/test/src/compiler/machine_mapping/machine_mapping_result.cc index 75dd63cccb..230602be98 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/machine_mapping_result.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/machine_mapping_result.cc @@ -10,7 +10,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -25,7 +24,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -191,7 +189,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -206,7 +203,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -285,7 +281,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/3_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -335,7 +330,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -350,7 +344,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { diff --git a/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc b/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc index 2ea8312991..29fdfff3f6 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc @@ -16,7 +16,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - DeviceType::GPU, }, { MachineViewDimension{ @@ -33,28 +32,6 @@ TEST_SUITE(FF_TEST_SUITE) { CHECK(mv_get_expected_task_space_num_dims(mv) == 2_n); } - TEST_CASE("get_device_type") { - MachineView mv = MachineView{ - MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - DeviceType::GPU, - }, - { - MachineViewDimension{ - stride_t{2_p}, - MachineSpecificationDimension::INTER_NODE, - }, - MachineViewDimension{ - stride_t{2_p}, - MachineSpecificationDimension::INTER_NODE, - }, - }, - }; - - CHECK(get_device_type(mv) == DeviceType::GPU); - } - TEST_CASE("get_machine_space_coordinate") { SUBCASE("1D case") { /** @@ -81,7 +58,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/1_n, - DeviceType::GPU, }, { MachineViewDimension{ @@ -100,7 +76,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate correct = MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/1_n, - DeviceType::GPU, }; CHECK(result == correct); @@ -115,7 +90,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate correct = MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/3_n, - DeviceType::GPU, }; CHECK(result == correct); @@ -130,7 +104,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate correct = MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/5_n, - DeviceType::GPU, }; CHECK(result == correct); @@ -174,7 +147,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate{ /*node_idx=*/1_n, /*device_idx=*/2_n, - DeviceType::GPU, }, { MachineViewDimension{ @@ -191,7 +163,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 0_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/2_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/2_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -200,7 +172,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/4_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/4_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -209,7 +181,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/2_n, /*device_idx=*/2_n, DeviceType::GPU}; + /*node_idx=*/2_n, /*device_idx=*/2_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -218,7 +190,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/2_n, /*device_idx=*/4_n, DeviceType::GPU}; + /*node_idx=*/2_n, /*device_idx=*/4_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -250,7 +222,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate{ /*node_idx=*/1_n, /*device_idx=*/0_n, - DeviceType::GPU, }, { MachineViewDimension{ @@ -267,7 +238,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 0_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/0_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/0_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -276,7 +247,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/4_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/4_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -285,7 +256,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/1_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/1_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -294,7 +265,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/5_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/5_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -332,7 +303,7 @@ TEST_SUITE(FF_TEST_SUITE) { }; MachineView mv = MachineView{ MachineSpaceCoordinate{ - /*node_idx=*/0_n, /*device_idx=*/1_n, DeviceType::GPU}, + /*node_idx=*/0_n, /*device_idx=*/1_n}, {MachineViewDimension{stride_t{1_p}, MachineSpecificationDimension::INTER_NODE}, MachineViewDimension{stride_t{2_p}, @@ -343,7 +314,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n, 0_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/0_n, /*device_idx=*/3_n, DeviceType::GPU}; + /*node_idx=*/0_n, /*device_idx=*/3_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -352,7 +323,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n, 1_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/5_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/5_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -361,106 +332,11 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n, 1_n}); MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/7_n, DeviceType::GPU}; + /*node_idx=*/1_n, /*device_idx=*/7_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); } } } - - TEST_CASE("get_device_ids") { - SUBCASE("1D machine view") { - /** - * This operator has shape (3,), and thus 3 tasks. - * The (only) dimension is projected onto the INTRA (device) dimension - * with a stride of 2. The start of the projection defined by MachineView - * is at MachineSpaceCoordinate (0, 1), and the machine space has 1 node - * and 6 devices per node. - * - * The tasks will thus be distributed like this: - * +-------+-------+-------+-------+-------+-------+ - * | 0 | ((1)) | 2 | ((3)) | 4 | ((5)) | - * +-------+-------+-------+-------+-------+-------+ - * Where the integers are the device ids and ((x)) are the devices we - * select - */ - MachineComputeSpecification ms = MachineComputeSpecification{ - /*num_nodes=*/1_p, - /*num_cpus_per_node=*/6_p, - /*num_gpus_per_node=*/6_p, - }; - - OperatorTaskSpace task = OperatorTaskSpace{ - MinimalOrthotope{{ - 3_ge2, - }}, - }; - MachineView mv = MachineView{ - MachineSpaceCoordinate{ - /*node_idx=*/0_n, /*device_idx=*/1_n, DeviceType::GPU}, - {MachineViewDimension{stride_t{2_p}, - MachineSpecificationDimension::INTRA_NODE}}}; - - std::unordered_set correct = { - device_id_t{gpu_id_t{1_n}}, - device_id_t{gpu_id_t{3_n}}, - device_id_t{gpu_id_t{5_n}}, - }; - std::unordered_set result = get_device_ids(task, mv, ms); - CHECK(result == correct); - } - - SUBCASE("2D machine view") { - /** - * This operator has shape (2, 2), and thus 2 * 2 = 4 tasks. - * - The first dimension is projected onto the INTER (node) dimension with - * stride 1, - * - The second dimension is projected onto the INTRA (device) dimension - * with stride 2. The start of the projection defined by MachineView is at - * MachineSpaceCoordinate (1, 2), and the machine space has 3 nodes and 5 - * devices per node. - * - * The tasks will thus be distributed like this: - * +-------+-------+-------+-------+-------+ - * | 0 | 1 | 2 | 3 | 4 | - * +-------+-------+-------+-------+-------+ - * | 5 | 6 | ((7)) | 8 | ((9)) | - * +-------+-------+-------+-------+-------+ - * | 10 | 11 | ((12))| 13 | ((14))| - * +-------+-------+-------+-------+-------+ - * Where the integers are the device ids and ((x)) are the devices we - * select - */ - - MachineComputeSpecification ms = MachineComputeSpecification{ - /*num_nodes=*/3_p, - /*num_cpus_per_node=*/5_p, - /*num_gpus_per_node=*/5_p, - }; - - OperatorTaskSpace task = OperatorTaskSpace{ - MinimalOrthotope{{ - 2_ge2, - 2_ge2, - }}, - }; - MachineView mv = MachineView{ - MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/2_n, DeviceType::GPU}, - {MachineViewDimension{stride_t{1_p}, - MachineSpecificationDimension::INTER_NODE}, - MachineViewDimension{stride_t{2_p}, - MachineSpecificationDimension::INTRA_NODE}}}; - - std::unordered_set correct = { - device_id_t{gpu_id_t{7_n}}, - device_id_t{gpu_id_t{9_n}}, - device_id_t{gpu_id_t{12_n}}, - device_id_t{gpu_id_t{14_n}}, - }; - std::unordered_set result = get_device_ids(task, mv, ms); - CHECK(result == correct); - } - } } diff --git a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc index 54717d6699..8d21131c7d 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc @@ -52,7 +52,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/{}, }; @@ -61,7 +60,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -76,7 +74,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -91,7 +88,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/1_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -724,7 +720,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/2_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { diff --git a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.cc b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.cc index 402dbe66d7..086c1ab984 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/memory_optimization/machine_mapping_with_memory_result.cc @@ -104,7 +104,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -119,7 +118,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -307,7 +305,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -322,7 +319,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -410,7 +406,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/3_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -464,7 +459,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -479,7 +473,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { @@ -494,7 +487,6 @@ TEST_SUITE(FF_TEST_SUITE) { /*start=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }, /*dimensions=*/ { diff --git a/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc b/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc index 3159f49118..e3b08e3805 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc @@ -12,8 +12,7 @@ TEST_SUITE(FF_TEST_SUITE) { {MachineViewDimension{stride_t{2_p}, MachineSpecificationDimension::INTER_NODE}, MachineViewDimension{stride_t{2_p}, - MachineSpecificationDimension::INTER_NODE}}, - DeviceType::GPU}; + MachineSpecificationDimension::INTER_NODE}}}; SUBCASE("num_dims") { nonnegative_int result = num_dims(simv); @@ -21,12 +20,6 @@ TEST_SUITE(FF_TEST_SUITE) { CHECK(result == correct); } - SUBCASE("get_device_type") { - DeviceType result = get_device_type(simv); - DeviceType correct = DeviceType::GPU; - CHECK(result == correct); - } - SUBCASE("get_strides") { std::vector result = get_strides(simv); std::vector correct = {stride_t{2_p}, stride_t{2_p}}; @@ -44,7 +37,7 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("StartInvariantMachineView - conversions") { MachineSpaceCoordinate start = - MachineSpaceCoordinate{1_n, 2_n, DeviceType::GPU}; + MachineSpaceCoordinate{1_n, 2_n}; std::vector dimensions = { MachineViewDimension{stride_t{2_p}, MachineSpecificationDimension::INTER_NODE}, @@ -53,7 +46,7 @@ TEST_SUITE(FF_TEST_SUITE) { MachineView mv = MachineView{start, dimensions}; StartInvariantMachineView simv = - StartInvariantMachineView{dimensions, DeviceType::GPU}; + StartInvariantMachineView{dimensions}; SUBCASE("start_invariant_from_machine_view") { StartInvariantMachineView result = start_invariant_from_machine_view(mv); @@ -102,8 +95,7 @@ TEST_SUITE(FF_TEST_SUITE) { }; StartInvariantMachineView simv = StartInvariantMachineView{ {MachineViewDimension{stride_t{2_p}, - MachineSpecificationDimension::INTRA_NODE}}, - DeviceType::GPU}; + MachineSpecificationDimension::INTRA_NODE}}}; MachineComputeSpecification ms = MachineComputeSpecification{ /*num_nodes=*/1_p, /*num_cpus_per_node=*/6_p, @@ -114,7 +106,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n}); MachineSpaceOffset correct = - MachineSpaceOffset{0, 0, DeviceType::GPU}; + MachineSpaceOffset{0, 0}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -123,7 +115,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n}); MachineSpaceOffset correct = - MachineSpaceOffset{0, 2, DeviceType::GPU}; + MachineSpaceOffset{0, 2}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -132,7 +124,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (2,)") { TaskSpaceCoordinate coord = make_task_space_coordinate({2_n}); MachineSpaceOffset correct = - MachineSpaceOffset{0, 4, DeviceType::GPU}; + MachineSpaceOffset{0, 4}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -141,9 +133,9 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("get_machine_space_offsets") { std::unordered_set correct = { - MachineSpaceOffset{0, 0, DeviceType::GPU}, - MachineSpaceOffset{0, 2, DeviceType::GPU}, - MachineSpaceOffset{0, 4, DeviceType::GPU}}; + MachineSpaceOffset{0, 0}, + MachineSpaceOffset{0, 2}, + MachineSpaceOffset{0, 4}}; std::unordered_set result = get_machine_space_offsets(task, simv); CHECK(correct == result); @@ -176,8 +168,7 @@ TEST_SUITE(FF_TEST_SUITE) { {MachineViewDimension{stride_t{1_p}, MachineSpecificationDimension::INTER_NODE}, MachineViewDimension{stride_t{2_p}, - MachineSpecificationDimension::INTRA_NODE}}, - DeviceType::GPU}; + MachineSpecificationDimension::INTRA_NODE}}}; MachineComputeSpecification ms = MachineComputeSpecification{ /*num_nodes=*/2_p, /*num_cpus_per_node=*/4_p, @@ -188,7 +179,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 0_n}); MachineSpaceOffset correct = - MachineSpaceOffset{0, 0, DeviceType::GPU}; + MachineSpaceOffset{0, 0}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -197,7 +188,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n}); MachineSpaceOffset correct = - MachineSpaceOffset{0, 2, DeviceType::GPU}; + MachineSpaceOffset{0, 2}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -206,7 +197,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n}); MachineSpaceOffset correct = - MachineSpaceOffset{1, 0, DeviceType::GPU}; + MachineSpaceOffset{1, 0}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -215,7 +206,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n}); MachineSpaceOffset correct = - MachineSpaceOffset{1, 2, DeviceType::GPU}; + MachineSpaceOffset{1, 2}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -224,10 +215,10 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("get_machine_space_offsets") { std::unordered_set correct = { - MachineSpaceOffset{0, 0, DeviceType::GPU}, - MachineSpaceOffset{0, 2, DeviceType::GPU}, - MachineSpaceOffset{1, 0, DeviceType::GPU}, - MachineSpaceOffset{1, 2, DeviceType::GPU}}; + MachineSpaceOffset{0, 0}, + MachineSpaceOffset{0, 2}, + MachineSpaceOffset{1, 0}, + MachineSpaceOffset{1, 2}}; std::unordered_set result = get_machine_space_offsets(task, simv); CHECK(correct == result); diff --git a/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc b/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc index 4c03af5382..0fe4277ab6 100644 --- a/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc +++ b/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc @@ -65,8 +65,7 @@ TEST_SUITE(FF_TEST_SUITE) { MCMCOverMappedPCGConfig no_search = MCMCOverMappedPCGConfig{/*temperature=*/1.0, /*num_iterations=*/1_n, - /*substitution_frequency=*/0.2, - /*device_type=*/DeviceType::GPU}; + /*substitution_frequency=*/0.2}; SearchResult base_result = mcmc_over_mapped_pcg(pcg, cost_estimator, full_machine_spec, no_search); @@ -80,8 +79,7 @@ TEST_SUITE(FF_TEST_SUITE) { MCMCOverMappedPCGConfig search_config = MCMCOverMappedPCGConfig{/*temperature=*/1.0, /*num_iterations=*/100_n, - /*substitution_frequency=*/0.2, - /*device_type=*/DeviceType::GPU}; + /*substitution_frequency=*/0.2}; SearchResult result = mcmc_over_mapped_pcg( pcg, cost_estimator, full_machine_spec, search_config); diff --git a/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc b/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc index 2846de6559..a7f8886846 100644 --- a/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc +++ b/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc @@ -13,8 +13,6 @@ #include "op-attrs/parallel_tensor_dims.dtg.h" #include "op-attrs/parallel_tensor_shape.dtg.h" #include "op-attrs/parallel_tensor_shape.h" -#include "pcg/device_id.h" -#include "pcg/device_type.dtg.h" #include "pcg/machine_space_coordinate.dtg.h" #include "pcg/machine_specification_dimension.dtg.h" #include "pcg/parallel_computation_graph/parallel_computation_graph.h" @@ -68,9 +66,9 @@ TEST_SUITE(FF_TEST_SUITE) { std::vector dims = {}; ParallelComputationGraph pcg = b.pcg; MachineView mv1 = - MachineView{MachineSpaceCoordinate{0_n, 0_n, DeviceType::GPU}, dims}; + MachineView{MachineSpaceCoordinate{0_n, 0_n}, dims}; MachineView mv2 = - MachineView{MachineSpaceCoordinate{0_n, 1_n, DeviceType::GPU}, dims}; + MachineView{MachineSpaceCoordinate{0_n, 1_n}, dims}; MachineMapping device_mapping = MachineMapping{{ {layer0, mv1}, @@ -149,13 +147,13 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("all different devices") { MachineView mv0 = MachineView{ - MachineSpaceCoordinate{0_n, 0_n, DeviceType::GPU}, dims}; + MachineSpaceCoordinate{0_n, 0_n}, dims}; MachineView mv1 = MachineView{ - MachineSpaceCoordinate{0_n, 1_n, DeviceType::GPU}, dims}; + MachineSpaceCoordinate{0_n, 1_n}, dims}; MachineView mv2 = MachineView{ - MachineSpaceCoordinate{1_n, 0_n, DeviceType::GPU}, dims}; + MachineSpaceCoordinate{1_n, 0_n}, dims}; MachineView mv3 = MachineView{ - MachineSpaceCoordinate{1_n, 1_n, DeviceType::GPU}, dims}; + MachineSpaceCoordinate{1_n, 1_n}, dims}; MachineMapping device_mapping = MachineMapping{{ {layer0, mv0}, @@ -209,7 +207,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("all the same device") { MachineView mv = MachineView{ - MachineSpaceCoordinate{0_n, 0_n, DeviceType::GPU}, dims}; + MachineSpaceCoordinate{0_n, 0_n}, dims}; MachineMapping device_mapping = MachineMapping{{ {layer0, mv}, {layer1, mv}, diff --git a/lib/local-execution/include/local-execution/computation_graph_instance.h b/lib/local-execution/include/local-execution/computation_graph_instance.h index a4ded5edaf..aa05cbd582 100644 --- a/lib/local-execution/include/local-execution/computation_graph_instance.h +++ b/lib/local-execution/include/local-execution/computation_graph_instance.h @@ -7,7 +7,7 @@ #include "kernels/profiling_settings.dtg.h" #include "local-execution/loss_config.dtg.h" #include "pcg/computation_graph.dtg.h" -#include "pcg/device_id_t.dtg.h" +#include "task-spec/device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" diff --git a/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h b/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h index fdff1153ff..a978f3c996 100644 --- a/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h +++ b/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h @@ -5,7 +5,7 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "pcg/device_id_t.dtg.h" +#include "task-spec/device_id_t.dtg.h" #include "pcg/machine_interconnect_specification.dtg.h" namespace FlexFlow { diff --git a/lib/local-execution/include/local-execution/local_task_argument_accessor.h b/lib/local-execution/include/local-execution/local_task_argument_accessor.h index 12eab4a76d..a7df8814d0 100644 --- a/lib/local-execution/include/local-execution/local_task_argument_accessor.h +++ b/lib/local-execution/include/local-execution/local_task_argument_accessor.h @@ -2,7 +2,7 @@ #define _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_LOCAL_TASK_ARGUMENT_ACCESSOR_H #include "kernels/accessor.h" -#include "pcg/device_id_t.dtg.h" +#include "task-spec/device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" #include "task-spec/task_argument_accessor/itask_argument_accessor.h" #include "task-spec/task_argument_accessor/task_tensor_parameter.dtg.h" diff --git a/lib/local-execution/include/local-execution/per_device_op_state_initialization.h b/lib/local-execution/include/local-execution/per_device_op_state_initialization.h index ff0ffd73f8..174e51f307 100644 --- a/lib/local-execution/include/local-execution/per_device_op_state_initialization.h +++ b/lib/local-execution/include/local-execution/per_device_op_state_initialization.h @@ -4,7 +4,7 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "pcg/device_id_t.dtg.h" +#include "task-spec/device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" diff --git a/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc b/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc index 0c6107c4af..27a1c39677 100644 --- a/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc +++ b/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc @@ -11,7 +11,6 @@ #include "op-attrs/tensor_slot_name.dtg.h" #include "pcg/computation_graph.h" #include "pcg/computation_graph/layer_added_result.dtg.h" -#include "pcg/device_id.h" #include "pcg/parallel_tensor_attrs.h" #include "utils/containers/concat_vectors.h" #include "utils/containers/map_values.h" @@ -105,7 +104,7 @@ OpCostMetrics LocalCostEstimator::estimate_cost( // allocate memory std::shared_ptr tracked_allocator_ptr = std::make_shared(create_local_allocator_for_device_type( - get_device_type(this->device_idx))); + this->device_idx.device_type)); layer_guid_t layer_guid = layer_guid_t{Node{0}}; diff --git a/lib/local-execution/src/local-execution/local_task_argument_accessor.cc b/lib/local-execution/src/local-execution/local_task_argument_accessor.cc index b8feca720e..5fbe207e6c 100644 --- a/lib/local-execution/src/local-execution/local_task_argument_accessor.cc +++ b/lib/local-execution/src/local-execution/local_task_argument_accessor.cc @@ -1,7 +1,5 @@ #include "local-execution/local_task_argument_accessor.h" #include "kernels/accessor.h" -#include "pcg/device_id.h" -#include "pcg/device_id_t.h" #include "utils/exception.h" #include "utils/optional.h" #include "utils/overload.h" @@ -84,7 +82,7 @@ device_handle_t LocalTaskArgumentAccessor::get_ff_handle() const { } DeviceType LocalTaskArgumentAccessor::get_kernel_device_type() const { - return get_device_type(this->device_idx); + return this->device_idx.device_type; } PCGOperatorAttrs LocalTaskArgumentAccessor::get_op_attrs() const { diff --git a/lib/local-execution/test/src/local-execution/computation_graph_instance.cc b/lib/local-execution/test/src/local-execution/computation_graph_instance.cc index 2a4e204d59..a4049a609f 100644 --- a/lib/local-execution/test/src/local-execution/computation_graph_instance.cc +++ b/lib/local-execution/test/src/local-execution/computation_graph_instance.cc @@ -11,7 +11,6 @@ #include "op-attrs/ops/loss_functions/loss_attrs.dtg.h" #include "pcg/computation_graph.h" #include "pcg/computation_graph_builder.h" -#include "pcg/device_id_t.h" #include "pcg/device_type.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.h" @@ -140,8 +139,14 @@ TEST_SUITE(FF_TEST_SUITE) { /*nesterov=*/false, /*weight_decay=*/0.001}}; device_handle_t ff_handle = cpu_make_device_handle_t(); - device_id_t device_idx = - make_device_id_t_from_idx(nonnegative_int{0}, DeviceType::CPU); + device_id_t device_idx = device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::CPU, + }; + std::unordered_map input_tensors; @@ -307,8 +312,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*weight_decay=*/0.001, }, }; - device_id_t device_idx = - make_device_id_t_from_idx(nonnegative_int{0}, DeviceType::GPU); + device_id_t device_idx = device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::GPU, + }; device_handle_t ff_handle = gpu_make_device_handle_t(managed_handle.raw_handle()); @@ -423,8 +433,14 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }, }; - device_id_t device_idx = - make_device_id_t_from_idx(nonnegative_int{0}, DeviceType::GPU); + device_id_t device_idx = device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::GPU, + }; + device_handle_t ff_handle = gpu_make_device_handle_t(managed_handle.raw_handle()); diff --git a/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc b/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc index ac20c1ad75..34cfbb4e3d 100644 --- a/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc +++ b/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc @@ -10,7 +10,6 @@ #include "op-attrs/parallel_tensor_shape.h" #include "op-attrs/tensor_slot_name.dtg.h" #include "pcg/computation_graph_builder.h" -#include "pcg/device_id_t.h" #include using namespace ::FlexFlow; @@ -19,8 +18,13 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("LocalCostEstimator") { Allocator allocator = create_local_cpu_memory_allocator(); device_handle_t ff_handle = cpu_make_device_handle_t(); - device_id_t device_idx = - make_device_id_t_from_idx(nonnegative_int{0}, DeviceType::CPU); + device_id_t device_idx = device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::CPU, + }; OptimizerAttrs optimizer_attrs = OptimizerAttrs{ SGDOptimizerAttrs{ @@ -66,7 +70,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*optimizer_attrs=*/optimizer_attrs, /*machine_view=*/ make_1d_machine_view( - MachineSpaceCoordinate{0_n, 0_n, DeviceType::CPU}, + MachineSpaceCoordinate{0_n, 0_n}, MachineSpecificationDimension::INTRA_NODE, stride_t{1_p}), }; @@ -89,8 +93,14 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { Allocator allocator = create_local_cuda_memory_allocator(); - device_id_t device_idx = - make_device_id_t_from_idx(nonnegative_int{0}, DeviceType::GPU); + device_id_t device_idx = device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::GPU, + }; + device_handle_t ff_handle = gpu_make_device_handle_t(managed_handle.raw_handle()); @@ -159,7 +169,7 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*optimizer_attrs=*/optimizer_attrs, /*machine_view=*/ make_1d_machine_view( - MachineSpaceCoordinate{0_n, 0_n, DeviceType::GPU}, + MachineSpaceCoordinate{0_n, 0_n}, MachineSpecificationDimension::INTRA_NODE, stride_t{1_p}), }; diff --git a/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc b/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc index 07bb869d5f..9594e13336 100644 --- a/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc +++ b/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc @@ -3,7 +3,6 @@ #include "kernels/local_cpu_allocator.h" #include "kernels/profiling_settings.dtg.h" #include "op-attrs/ops/input_attrs.dtg.h" -#include "pcg/device_id_t.h" #include "task-spec/task_argument_accessor/task_tensor_parameter.h" #include "task-spec/task_impl_function.dtg.h" #include "utils/fmt/variant.h" @@ -52,8 +51,13 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; - device_id_t device_idx = - make_device_id_t_from_idx(nonnegative_int{0}, DeviceType::CPU); + device_id_t device_idx = device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::CPU, + }; LocalTaskArgumentAccessor acc = LocalTaskArgumentAccessor{ /*allocator=*/allocator, diff --git a/lib/pcg/include/pcg/device_id.h b/lib/pcg/include/pcg/device_id.h deleted file mode 100644 index 36ea9de6b3..0000000000 --- a/lib/pcg/include/pcg/device_id.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef _FLEXFLOW_PCG_INCLUDE_PCG_DEVICE_ID_H -#define _FLEXFLOW_PCG_INCLUDE_PCG_DEVICE_ID_H - -#include "pcg/cpu_id_t.dtg.h" -#include "pcg/device_id_t.dtg.h" -#include "pcg/device_type.dtg.h" -#include "pcg/gpu_id_t.dtg.h" - -namespace FlexFlow { - -device_id_t operator+(device_id_t, size_t); - -DeviceType get_device_type(device_id_t const &device_id); -gpu_id_t unwrap_gpu(device_id_t); -cpu_id_t unwrap_cpu(device_id_t); -nonnegative_int get_raw_id(device_id_t); - -device_id_t device_id_from_index(nonnegative_int, DeviceType); - -} // namespace FlexFlow - -#endif diff --git a/lib/pcg/include/pcg/device_id_t.dtg.toml b/lib/pcg/include/pcg/device_id_t.dtg.toml deleted file mode 100644 index 4efcb07975..0000000000 --- a/lib/pcg/include/pcg/device_id_t.dtg.toml +++ /dev/null @@ -1,23 +0,0 @@ -namespace = "FlexFlow" -name = "device_id_t" -type = "variant" -features = [ - "eq", - "ord", - "hash", - "json", - "fmt", -] - -includes = [ - "pcg/cpu_id_t.dtg.h", - "pcg/gpu_id_t.dtg.h", -] - -[[values]] -type = "::FlexFlow::gpu_id_t" -key = "gpu" - -[[values]] -type = "::FlexFlow::cpu_id_t" -key = "cpu" diff --git a/lib/pcg/include/pcg/device_id_t.h b/lib/pcg/include/pcg/device_id_t.h deleted file mode 100644 index e8e605b068..0000000000 --- a/lib/pcg/include/pcg/device_id_t.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_DEVICE_ID_T_H -#define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_DEVICE_ID_T_H - -#include "pcg/device_id_t.dtg.h" -#include "pcg/device_type.dtg.h" - -namespace FlexFlow { - -device_id_t make_device_id_t_from_idx(nonnegative_int idx, - DeviceType device_type); - -} // namespace FlexFlow - -#endif diff --git a/lib/pcg/include/pcg/machine_compute_resource_slice.dtg.toml b/lib/pcg/include/pcg/machine_compute_resource_slice.dtg.toml index 77d6b7558c..41693ca0dc 100644 --- a/lib/pcg/include/pcg/machine_compute_resource_slice.dtg.toml +++ b/lib/pcg/include/pcg/machine_compute_resource_slice.dtg.toml @@ -6,6 +6,8 @@ features = [ "ord", "hash", "fmt", + "rapidcheck", + "json", ] includes = [ diff --git a/lib/pcg/include/pcg/machine_compute_specification.h b/lib/pcg/include/pcg/machine_compute_specification.h index 835e9040e0..2e8f910c08 100644 --- a/lib/pcg/include/pcg/machine_compute_specification.h +++ b/lib/pcg/include/pcg/machine_compute_specification.h @@ -1,7 +1,6 @@ #ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_MACHINE_COMPUTE_SPECIFICATION_H #define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_MACHINE_COMPUTE_SPECIFICATION_H -#include "pcg/device_id_t.dtg.h" #include "pcg/device_type.dtg.h" #include "pcg/machine_compute_specification.dtg.h" #include "pcg/machine_space_coordinate.dtg.h" @@ -15,12 +14,6 @@ positive_int get_num_devices(MachineComputeSpecification const &ms, positive_int get_num_devices_per_node(MachineComputeSpecification const &ms, DeviceType const &device_type); -bool is_valid_machine_space_coordinate(MachineComputeSpecification const &ms, - MachineSpaceCoordinate const &coord); - -device_id_t get_device_id(MachineComputeSpecification const &ms, - MachineSpaceCoordinate const &coord); - } // namespace FlexFlow #endif diff --git a/lib/pcg/include/pcg/machine_space_coordinate.dtg.toml b/lib/pcg/include/pcg/machine_space_coordinate.dtg.toml index 41f4d563f3..8787b0d8f2 100644 --- a/lib/pcg/include/pcg/machine_space_coordinate.dtg.toml +++ b/lib/pcg/include/pcg/machine_space_coordinate.dtg.toml @@ -10,8 +10,7 @@ features = [ "fmt", ] -includes = [ - "pcg/device_type.dtg.h", +includes = [ "utils/nonnegative_int/nonnegative_int.h", ] @@ -22,7 +21,3 @@ type = "::FlexFlow::nonnegative_int" [[fields]] name = "device_idx" type = "::FlexFlow::nonnegative_int" - -[[fields]] -name = "device_type" -type = "::FlexFlow::DeviceType" diff --git a/lib/pcg/include/pcg/machine_space_offset.dtg.toml b/lib/pcg/include/pcg/machine_space_offset.dtg.toml index 57f884906b..54eda5cc76 100644 --- a/lib/pcg/include/pcg/machine_space_offset.dtg.toml +++ b/lib/pcg/include/pcg/machine_space_offset.dtg.toml @@ -10,9 +10,7 @@ features = [ "fmt", ] -includes = [ - "pcg/device_type.dtg.h", -] +includes = [] [[fields]] name = "node_offset" @@ -21,7 +19,3 @@ type = "int" [[fields]] name = "device_offset" type = "int" - -[[fields]] -name = "device_type" -type = "::FlexFlow::DeviceType" diff --git a/lib/pcg/src/pcg/device_id.cc b/lib/pcg/src/pcg/device_id.cc deleted file mode 100644 index 1a4f7b7d22..0000000000 --- a/lib/pcg/src/pcg/device_id.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "pcg/device_id.h" -#include "utils/exception.h" -#include - -namespace FlexFlow { - -device_id_t operator+(device_id_t, size_t) { - NOT_IMPLEMENTED(); -} - -DeviceType get_device_type(device_id_t const &device_id) { - if (device_id.has()) { - return DeviceType::GPU; - } else { - assert(device_id.has()); - return DeviceType::CPU; - } -} - -gpu_id_t unwrap_gpu(device_id_t device_id) { - return device_id.get(); -} - -cpu_id_t unwrap_cpu(device_id_t device_id) { - return device_id.get(); -} - -nonnegative_int get_raw_id(device_id_t device_id) { - switch (get_device_type(device_id)) { - case DeviceType::GPU: - return unwrap_gpu(device_id).gpu_index; - case DeviceType::CPU: - return unwrap_cpu(device_id).cpu_index; - default: - throw mk_runtime_error(fmt::format("Unsupported device {}", device_id)); - } -} - -device_id_t device_id_from_index(nonnegative_int idx, DeviceType device_type) { - switch (device_type) { - case DeviceType::GPU: - return device_id_t{gpu_id_t{idx}}; - case DeviceType::CPU: - return device_id_t{cpu_id_t{idx}}; - default: - throw mk_runtime_error( - fmt::format("Unsupported DeviceType {}", device_type)); - } -} - -} // namespace FlexFlow diff --git a/lib/pcg/src/pcg/device_id_t.cc b/lib/pcg/src/pcg/device_id_t.cc deleted file mode 100644 index eecaf1c81d..0000000000 --- a/lib/pcg/src/pcg/device_id_t.cc +++ /dev/null @@ -1,17 +0,0 @@ -#include "pcg/device_id_t.h" - -namespace FlexFlow { - -device_id_t make_device_id_t_from_idx(nonnegative_int idx, - DeviceType device_type) { - switch (device_type) { - case DeviceType::GPU: - return device_id_t{gpu_id_t{idx}}; - case DeviceType::CPU: - return device_id_t{cpu_id_t{idx}}; - default: - PANIC("Unhandled device_type", device_type); - } -} - -} // namespace FlexFlow diff --git a/lib/pcg/src/pcg/machine_compute_resource_slice.cc b/lib/pcg/src/pcg/machine_compute_resource_slice.cc index 2cb2386926..401ed3d850 100644 --- a/lib/pcg/src/pcg/machine_compute_resource_slice.cc +++ b/lib/pcg/src/pcg/machine_compute_resource_slice.cc @@ -20,7 +20,6 @@ positive_int bool is_valid_machine_space_coordinate_in_slice( MachineComputeResourceSlice const &slice, MachineSpaceCoordinate const &coord) { - ASSERT(coord.device_type == DeviceType::GPU); return (coord.node_idx < slice.num_nodes) && (coord.device_idx < slice.num_gpus_per_node); diff --git a/lib/pcg/src/pcg/machine_compute_specification.cc b/lib/pcg/src/pcg/machine_compute_specification.cc index a8dfb27524..0dbe546fff 100644 --- a/lib/pcg/src/pcg/machine_compute_specification.cc +++ b/lib/pcg/src/pcg/machine_compute_specification.cc @@ -1,5 +1,4 @@ #include "pcg/machine_compute_specification.h" -#include "pcg/device_id.h" #include "utils/containers/transform.h" #include @@ -37,20 +36,4 @@ positive_int get_num_devices_per_node(MachineComputeSpecification const &ms, } } -bool is_valid_machine_space_coordinate(MachineComputeSpecification const &ms, - MachineSpaceCoordinate const &coord) { - return (coord.node_idx < ms.num_nodes) && - (coord.device_idx < get_num_devices_per_node(ms, coord.device_type)); -} - -device_id_t get_device_id(MachineComputeSpecification const &ms, - MachineSpaceCoordinate const &coord) { - ASSERT(is_valid_machine_space_coordinate(ms, coord)); - - nonnegative_int raw_idx = - coord.node_idx * get_num_devices_per_node(ms, coord.device_type) + - coord.device_idx; - return device_id_from_index(raw_idx, coord.device_type); -} - } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/machine_space_offset.cc b/lib/pcg/src/pcg/machine_space_offset.cc index 953dc38bc6..0240d3582d 100644 --- a/lib/pcg/src/pcg/machine_space_offset.cc +++ b/lib/pcg/src/pcg/machine_space_offset.cc @@ -13,14 +13,11 @@ MachineSpaceOffset get_machine_space_offset_from_coordinate( "The start node_idx is greater than one of the coord node_idx." "Are you sure you didn't swap them?"); - ASSERT(start.device_type == coord.device_type); - return MachineSpaceOffset{ /*node_offset=*/coord.node_idx.unwrap_nonnegative() - start.node_idx.unwrap_nonnegative(), /*device_offset=*/coord.device_idx.unwrap_nonnegative() - start.device_idx.unwrap_nonnegative(), - /*device_type=*/coord.device_type, }; } diff --git a/lib/pcg/test/src/pcg/machine_compute_specification.cc b/lib/pcg/test/src/pcg/machine_compute_specification.cc index c725da80ed..8a20d8d10d 100644 --- a/lib/pcg/test/src/pcg/machine_compute_specification.cc +++ b/lib/pcg/test/src/pcg/machine_compute_specification.cc @@ -1,5 +1,4 @@ #include "pcg/machine_compute_specification.h" -#include "pcg/device_id.h" #include using namespace FlexFlow; @@ -25,27 +24,5 @@ TEST_SUITE(FF_TEST_SUITE) { CHECK(get_num_devices(ms, DeviceType::GPU) == 4 * 8); CHECK(get_num_devices(ms, DeviceType::CPU) == 16 * 4); } - - SUBCASE("get_device_id") { - SUBCASE("valid MachineSpaceCoordinate") { - MachineSpaceCoordinate coord = MachineSpaceCoordinate{ - /*node_idx=*/2_n, - /*device_idx=*/12_n, - DeviceType::CPU, - }; - device_id_t correct = - device_id_from_index(nonnegative_int{2 * 16 + 12}, DeviceType::CPU); - device_id_t result = get_device_id(ms, coord); - CHECK(correct == result); - } - SUBCASE("MachineSpaceCoordinate out of bounds for given machine spec") { - MachineSpaceCoordinate coord = MachineSpaceCoordinate{ - /*node_idx=*/2_n, - /*device_idx=*/18_n, - DeviceType::CPU, - }; - CHECK_THROWS(get_device_id(ms, coord)); - } - } } } diff --git a/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_operator_task_group.cc b/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_operator_task_group.cc index 1c3667afc7..3543ebe09f 100644 --- a/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_operator_task_group.cc +++ b/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_operator_task_group.cc @@ -13,14 +13,21 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("adl_serializer") { bidict shard_bindings{ - {MachineSpaceCoordinate{0_n, 0_n, DeviceType::CPU}, - OperatorAtomicTaskShardBinding{ - { - {TensorSlotName::INPUT, - ParallelTensorSpaceCoordinate{ - 0_n, 0_n, FFOrdered{1_n, 2_n, 3_n}}}, - }, - }}, + { + MachineSpaceCoordinate{0_n, 0_n}, + OperatorAtomicTaskShardBinding{ + { + { + TensorSlotName::INPUT, + ParallelTensorSpaceCoordinate{ + /*sum_component=*/0_n, + /*discard_copy_component=*/0_n, + /*shard_components=*/FFOrdered{1_n, 2_n, 3_n}, + }, + }, + }, + }, + }, }; MappedOperatorTaskGroup deserialized{shard_bindings}; nlohmann::json serialized = shard_bindings; diff --git a/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.cc b/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.cc index 7856d89f27..5506afc63e 100644 --- a/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.cc +++ b/lib/pcg/test/src/pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.cc @@ -53,7 +53,6 @@ TEST_SUITE(FF_TEST_SUITE) { return MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/x, - /*device_type=*/DeviceType::GPU, }; }; diff --git a/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h b/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h index b698e613b5..e5ab8b597e 100644 --- a/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h +++ b/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h @@ -3,8 +3,8 @@ #include "kernels/device_handle_t.dtg.h" #include "kernels/managed_per_device_ff_handle.h" -#include "pcg/device_id_t.dtg.h" #include "realm-execution/device_specific_ptr.h" +#include "task-spec/device_id_t.dtg.h" #include namespace FlexFlow { diff --git a/lib/realm-execution/include/realm-execution/device_specific_ptr.h b/lib/realm-execution/include/realm-execution/device_specific_ptr.h index 32f8c13272..e8de4cfdad 100644 --- a/lib/realm-execution/include/realm-execution/device_specific_ptr.h +++ b/lib/realm-execution/include/realm-execution/device_specific_ptr.h @@ -1,7 +1,8 @@ #ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_DEVICE_SPECIFIC_PTR_H #define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_DEVICE_SPECIFIC_PTR_H -#include "pcg/device_id_t.dtg.h" +#include "task-spec/device_id_t.dtg.h" +#include #include namespace FlexFlow { diff --git a/lib/realm-execution/include/realm-execution/distributed_ff_handle.h b/lib/realm-execution/include/realm-execution/distributed_ff_handle.h index 8409a234a7..b2b61ee347 100644 --- a/lib/realm-execution/include/realm-execution/distributed_ff_handle.h +++ b/lib/realm-execution/include/realm-execution/distributed_ff_handle.h @@ -9,7 +9,7 @@ namespace FlexFlow { /** - * \brief Tracks the \ref device_handle_t (i.e., FFHandle) for each %GPU, both + * \brief Tracks the \ref device_handle_t (i.e., FFHandle) for each GPU, both * local and remote. A GPU here is represented by a Realm::Processor. */ struct DistributedFfHandle { @@ -31,8 +31,8 @@ struct DistributedFfHandle { /** * \brief Launches tasks (using \ref spawn_ff_handle_init_task) to create - * the \ref device_handle_t ""s for each %GPU and packages the results into a - * DistributedFfHandle. + * the \ref device_handle_t ""s for each GPU and packages the results into a + * \ref DistributedFfHandle. * * \relates DistributedFfHandle */ diff --git a/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h b/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h index 5d52f8caaf..f160146b96 100644 --- a/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h +++ b/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h @@ -25,7 +25,8 @@ PerDeviceOpStateBacking perform_distributed_per_device_op_state_initialization( ProfilingSettings const &profiling_settings, DistributedFfHandle const &device_handle, OptimizerAttrs const &optimizer_attrs, - Realm::Event precondition); + Realm::Event precondition, + DeviceType device_type); } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/instance_allocation.h b/lib/realm-execution/include/realm-execution/instance_allocation.h index 66cc07af75..1f2a6d9134 100644 --- a/lib/realm-execution/include/realm-execution/instance_allocation.h +++ b/lib/realm-execution/include/realm-execution/instance_allocation.h @@ -12,10 +12,9 @@ namespace FlexFlow { * on the device represented by \p device_coord. */ std::pair - perform_instance_allocation_for_value( - MachineSpaceCoordinate const &device_coord, - DynamicValueAttrs const &value, - RealmContext &ctx); + perform_instance_allocation_for_value(device_id_t const &device_id, + DynamicValueAttrs const &value, + RealmContext &ctx); /** * @brief Allocates the (potentially remote) Realm instances for all of the @@ -28,7 +27,8 @@ TensorInstanceBacking perform_instance_allocation( DynamicOpenDataflowGraph const &g, std::unordered_map const &preallocated, - RealmContext &ctx); + RealmContext &ctx, + DeviceType device_type); /** * @brief Destroys all of the instances held in \p instances. diff --git a/lib/realm-execution/include/realm-execution/pcg_instance.h b/lib/realm-execution/include/realm-execution/pcg_instance.h index 7b86d6d383..1e17856999 100644 --- a/lib/realm-execution/include/realm-execution/pcg_instance.h +++ b/lib/realm-execution/include/realm-execution/pcg_instance.h @@ -4,7 +4,6 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "realm-execution/distributed_ff_handle.h" @@ -12,6 +11,7 @@ #include "realm-execution/per_device_op_state_backing.dtg.h" #include "realm-execution/realm_context.h" #include "realm-execution/tensor_instance_backing.dtg.h" +#include "task-spec/device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" @@ -86,7 +86,8 @@ PCGInstance create_pcg_instance( std::unordered_map const &input_tensors, ProfilingSettings const &profiling_settings, - DistributedFfHandle const &ff_handle); + DistributedFfHandle const &ff_handle, + DeviceType device_type); /** * \brief Dispatch a training iteration for a \ref PCGInstance. diff --git a/lib/realm-execution/include/realm-execution/processor_kind.h b/lib/realm-execution/include/realm-execution/processor_kind.h new file mode 100644 index 0000000000..01bf6f8c35 --- /dev/null +++ b/lib/realm-execution/include/realm-execution/processor_kind.h @@ -0,0 +1,14 @@ +#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_PROCESSOR_KIND_H +#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_PROCESSOR_KIND_H + +#include "pcg/device_type.dtg.h" +#include "realm-execution/realm.h" + +namespace FlexFlow { + +DeviceType device_type_from_processor_kind(Realm::Processor::Kind); +Realm::Processor::Kind processor_kind_from_device_type(DeviceType); + +} // namespace FlexFlow + +#endif diff --git a/lib/realm-execution/include/realm-execution/realm_context.h b/lib/realm-execution/include/realm-execution/realm_context.h index ab89e916c0..0f2bed7de4 100644 --- a/lib/realm-execution/include/realm-execution/realm_context.h +++ b/lib/realm-execution/include/realm-execution/realm_context.h @@ -6,10 +6,10 @@ #include "kernels/managed_per_device_ff_handle.h" #include "op-attrs/parallel_tensor_shape.dtg.h" #include "op-attrs/tensor_shape.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/machine_space_coordinate.dtg.h" #include "realm-execution/realm.h" #include "realm-execution/tasks/task_id_t.dtg.h" +#include "task-spec/device_id_t.dtg.h" #include #include @@ -32,8 +32,8 @@ struct RealmContext { /** \name Device mapping */ ///\{ - Realm::Processor - map_device_coord_to_processor(MachineSpaceCoordinate const &); + Realm::Processor map_device_coord_to_processor(device_id_t const &) const; + device_id_t map_processor_to_device_coord(Realm::Processor) const; static Realm::Memory get_nearest_memory(Realm::Processor); ///\} @@ -97,8 +97,6 @@ struct RealmContext { */ [[nodiscard]] Realm::Event merge_outstanding_events(); - void discover_machine_topology(); - static std::optional make_device_handle_for_processor(Realm::Processor processor); @@ -111,14 +109,15 @@ struct RealmContext { */ Realm::Runtime get_runtime(); -private: + void discover_machine_topology(); + +public: Realm::Runtime runtime; Realm::Processor processor; Allocator allocator; std::vector outstanding_events; - std::unordered_map, - std::vector> - processors; + std::optional> processors = + std::nullopt; }; } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/realm_manager.h b/lib/realm-execution/include/realm-execution/realm_manager.h index 3984291641..ed01f4a022 100644 --- a/lib/realm-execution/include/realm-execution/realm_manager.h +++ b/lib/realm-execution/include/realm-execution/realm_manager.h @@ -3,10 +3,10 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "realm-execution/realm.h" #include "realm-execution/realm_context.h" #include "realm-execution/tasks/impl/controller_task.h" +#include "task-spec/device_id_t.dtg.h" namespace FlexFlow { diff --git a/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc b/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc index 8cf4c21b25..015af2b1ae 100644 --- a/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc +++ b/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc @@ -22,7 +22,8 @@ PerDeviceOpStateBacking perform_distributed_per_device_op_state_initialization( ProfilingSettings const &profiling_settings, DistributedFfHandle const &device_handle, OptimizerAttrs const &optimizer_attrs, - Realm::Event precondition) { + Realm::Event precondition, + DeviceType device_type) { // Initialize all operators and save the per-device op state ASSERT(no_nodes_are_initialized(dg)); diff --git a/lib/realm-execution/src/realm-execution/instance_allocation.cc b/lib/realm-execution/src/realm-execution/instance_allocation.cc index 4ef2919b10..37bf0ca03d 100644 --- a/lib/realm-execution/src/realm-execution/instance_allocation.cc +++ b/lib/realm-execution/src/realm-execution/instance_allocation.cc @@ -22,15 +22,14 @@ namespace FlexFlow { std::pair - perform_instance_allocation_for_value( - MachineSpaceCoordinate const &device_coord, - DynamicValueAttrs const &value, - RealmContext &ctx) { + perform_instance_allocation_for_value(device_id_t const &device_id, + DynamicValueAttrs const &value, + RealmContext &ctx) { ASSERT(value.accessor == std::nullopt); TensorShape shape = get_piece_shape(value.parallel_tensor_shape.value()); - Realm::Processor proc = ctx.map_device_coord_to_processor(device_coord); + Realm::Processor proc = ctx.map_device_coord_to_processor(device_id); Realm::Memory memory = ctx.get_nearest_memory(proc); return ctx.create_instance(memory, shape, Realm::ProfilingRequestSet()); } @@ -39,7 +38,8 @@ TensorInstanceBacking perform_instance_allocation( DynamicOpenDataflowGraph const &g, std::unordered_map const &preallocated, - RealmContext &ctx) { + RealmContext &ctx, + DeviceType device_type) { ASSERT(no_tensors_are_allocated(g)); ASSERT(tensors_are_ready_for_allocation(g)); for (DynamicValueAttrs const &v : keys(preallocated)) { @@ -53,9 +53,9 @@ TensorInstanceBacking perform_instance_allocation( NOT_IMPLEMENTED(); } else { if (!contains_key(result.backing, v)) { - MachineSpaceCoordinate device_coord = assert_unwrap(n.device_coord); + device_id_t device = assert_unwrap(n.device_coord); result.backing.insert(std::pair{ - v, perform_instance_allocation_for_value(device_coord, v, ctx)}); + v, perform_instance_allocation_for_value(device, v, ctx)}); } return result.backing.at(v); } diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc index 1ef0c4270b..62747fda7a 100644 --- a/lib/realm-execution/src/realm-execution/pcg_instance.cc +++ b/lib/realm-execution/src/realm-execution/pcg_instance.cc @@ -85,20 +85,31 @@ PCGInstance create_pcg_instance( std::unordered_map const &input_tensors, ProfilingSettings const &profiling_settings, - DistributedFfHandle const &device_handle) { + DistributedFfHandle const &device_handle, + DeviceType device_type) { DynamicOpenDataflowGraph dg = - make_dynamic_open_dataflow_graph_from_mapped_pcg(mpcg); + make_dynamic_open_dataflow_graph_from_mapped_pcg(mpcg, device_type); dg = perform_pass_expansion(dg); std::unordered_map inputs = input_tensors; std::optional logit_grad_value; if (loss.has_value()) { - auto [loss_attrs, label_tensor, logit_tensor, loss_mapping] = - assert_unwrap(loss); + ParallelLossConfig loss_config = assert_unwrap(loss); + + LossAttrs loss_attrs = loss_config.loss_attrs; + GenericTensorAccessorR label_tensor = loss_config.label_tensor; + parallel_tensor_guid_t logit_tensor = loss_config.logit_tensor; + MappedOperatorTaskGroup loss_op_task_group = loss_config.loss_mapping; + + DynamicNodeMapping mapping = DynamicNodeMapping{ + /*op_task_group=*/loss_op_task_group, + /*device_type=*/device_type, + }; + auto [dg2, label_v, logit_grad_v] = perform_loss_insertion( - dg, loss_attrs, dynamic_tensor_guid_t{logit_tensor}, loss_mapping); + dg, loss_attrs, dynamic_tensor_guid_t{logit_tensor}, mapping); dg = dg2; logit_grad_value = logit_grad_v; inputs.insert(std::pair{label_v, label_tensor}); @@ -106,9 +117,11 @@ PCGInstance create_pcg_instance( dg = perform_update_insertion(dg, optimizer_attrs); dg = perform_copy_insertion(dg); + debug_print_dynamic_open_dataflow_graph_as_dot(dg); dg = perform_shard_expansion(dg); + TensorInstanceBacking tensor_instance_backing = - perform_instance_allocation(dg, inputs, ctx); + perform_instance_allocation(dg, inputs, ctx, device_type); logit_grad_value = transform(logit_grad_value, [&](DynamicValueAttrs const &lgv) { @@ -141,7 +154,8 @@ PCGInstance create_pcg_instance( profiling_settings, device_handle, optimizer_attrs, - ctx.get_outstanding_events()); + ctx.get_outstanding_events(), + device_type); // Compute the topological ordering of the graph auto [kwarg_graph, node_map] = @@ -150,12 +164,14 @@ PCGInstance create_pcg_instance( std::vector invocation_topo_order = transform( node_topo_order, [&](Node node) { return node_map.at_l(node); }); - return PCGInstance{/*ctx=*/ctx, - /*execution_order=*/invocation_topo_order, - /*tensor_instance_backing=*/tensor_instance_backing, - /*device_state_backing=*/device_state_backing, - /*optimizer_attrs=*/optimizer_attrs, - /*logit_grad_tensor=*/logit_grad_tensor}; + return PCGInstance{ + /*ctx=*/ctx, + /*execution_order=*/invocation_topo_order, + /*tensor_instance_backing=*/tensor_instance_backing, + /*device_state_backing=*/device_state_backing, + /*optimizer_attrs=*/optimizer_attrs, + /*logit_grad_tensor=*/logit_grad_tensor, + }; } /** diff --git a/lib/realm-execution/src/realm-execution/processor_kind.cc b/lib/realm-execution/src/realm-execution/processor_kind.cc new file mode 100644 index 0000000000..3a40de0a42 --- /dev/null +++ b/lib/realm-execution/src/realm-execution/processor_kind.cc @@ -0,0 +1,29 @@ +#include "realm-execution/processor_kind.h" +#include + +namespace FlexFlow { + +DeviceType + device_type_from_processor_kind(Realm::Processor::Kind processor_kind) { + switch (processor_kind) { + case Realm::Processor::Kind::LOC_PROC: + return DeviceType::CPU; + case Realm::Processor::Kind::TOC_PROC: + return DeviceType::GPU; + default: + PANIC("Unhandled Realm::Processor::Kind", processor_kind); + } +} + +Realm::Processor::Kind processor_kind_from_device_type(DeviceType device_type) { + switch (device_type) { + case DeviceType::CPU: + return Realm::Processor::Kind::LOC_PROC; + case DeviceType::GPU: + return Realm::Processor::Kind::TOC_PROC; + default: + PANIC("Unhandled DeviceType", device_type); + } +} + +} // namespace FlexFlow diff --git a/lib/realm-execution/src/realm-execution/realm_allocator.cc b/lib/realm-execution/src/realm-execution/realm_allocator.cc index 9d5f3ff6b4..842539f9c2 100644 --- a/lib/realm-execution/src/realm-execution/realm_allocator.cc +++ b/lib/realm-execution/src/realm-execution/realm_allocator.cc @@ -1,6 +1,7 @@ #include "realm-execution/realm_allocator.h" #include "kernels/device.h" #include "pcg/device_type.dtg.h" +#include "realm-execution/processor_kind.h" #include "utils/containers/contains_key.h" #include "utils/containers/values.h" @@ -45,14 +46,7 @@ void RealmAllocator::deallocate(void *ptr) { } DeviceType RealmAllocator::get_allocation_device_type() const { - switch (this->processor.kind()) { - case Realm::Processor::Kind::LOC_PROC: - return DeviceType::CPU; - case Realm::Processor::Kind::TOC_PROC: - return DeviceType::GPU; - default: - PANIC("Unhandled FwbTensorType", this->processor.kind()); - } + return device_type_from_processor_kind(this->processor.kind()); } Allocator get_realm_allocator(Realm::Processor processor, diff --git a/lib/realm-execution/src/realm-execution/realm_context.cc b/lib/realm-execution/src/realm-execution/realm_context.cc index 790c1bd613..4cd7bdc221 100644 --- a/lib/realm-execution/src/realm-execution/realm_context.cc +++ b/lib/realm-execution/src/realm-execution/realm_context.cc @@ -4,8 +4,7 @@ #include "op-attrs/datatype.h" #include "op-attrs/parallel_tensor_shape.h" #include "op-attrs/tensor_dims.dtg.h" -#include "pcg/device_id_t.h" -#include "pcg/device_type.dtg.h" +#include "realm-execution/processor_kind.h" #include "realm-execution/realm_allocator.h" #include "realm-execution/tasks/task_id_t.dtg.h" #include "realm-execution/tasks/task_id_t.h" @@ -14,6 +13,7 @@ #include "utils/exception.h" #include "utils/nonnegative_int/nonnegative_int.h" #include "utils/one_to_many/one_to_many.h" +#include "utils/optional.h" #include "utils/positive_int/positive_int.h" namespace FlexFlow { @@ -21,7 +21,11 @@ namespace FlexFlow { RealmContext::RealmContext(Realm::Processor processor) : processor(processor), allocator(get_realm_allocator( - processor, RealmContext::get_nearest_memory(processor))) {} + processor, RealmContext::get_nearest_memory(processor))) { + if (processor != Realm::Processor::NO_PROC) { + this->discover_machine_topology(); + } +} RealmContext::~RealmContext() { if (!this->outstanding_events.empty()) { @@ -31,31 +35,22 @@ RealmContext::~RealmContext() { } static std::tuple - convert_machine_space_coordinate( - MachineSpaceCoordinate const &device_coord) { + convert_machine_space_coordinate(MachineSpaceCoordinate const &device_coord, + DeviceType device_type) { Realm::AddressSpace as = int{device_coord.node_idx}; - Realm::Processor::Kind kind; - switch (device_coord.device_type) { - case DeviceType::CPU: - kind = Realm::Processor::Kind::LOC_PROC; - break; - case DeviceType::GPU: - kind = Realm::Processor::Kind::TOC_PROC; - break; - default: - PANIC("Unhandled DeviceType", fmt::to_string(device_coord.device_type)); - break; - } + Realm::Processor::Kind kind = processor_kind_from_device_type(device_type); nonnegative_int proc_in_node = device_coord.device_idx; return std::tuple{as, kind, proc_in_node}; } Realm::Processor RealmContext::map_device_coord_to_processor( - MachineSpaceCoordinate const &device_coord) { - this->discover_machine_topology(); - auto [as, kind, proc_in_node] = - convert_machine_space_coordinate(device_coord); - return this->processors.at(std::pair{as, kind}).at(int{proc_in_node}); + device_id_t const &device_id) const { + return assert_unwrap(this->processors).at_r(device_id); +} + +device_id_t + RealmContext::map_processor_to_device_coord(Realm::Processor p) const { + return assert_unwrap(this->processors).at_l(p); } Realm::Memory RealmContext::get_nearest_memory(Realm::Processor proc) { @@ -81,27 +76,7 @@ Allocator &RealmContext::get_current_device_allocator() { device_id_t RealmContext::get_current_device_idx() const { Realm::Processor proc = this->get_current_processor(); - - // FIXME: find a more efficient way to implement this than scanning the - // machine every time - Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); - pq.same_address_space_as(proc); - nonnegative_int idx{0}; - for (Realm::Processor p : pq) { - if (p == proc) { - break; - } - idx++; - } - - switch (proc.kind()) { - case Realm::Processor::LOC_PROC: - return make_device_id_t_from_idx(idx, DeviceType::CPU); - case Realm::Processor::TOC_PROC: - return make_device_id_t_from_idx(idx, DeviceType::GPU); - default: - PANIC("Unhandled Realm::ProcessorKind", fmt::to_string(int{proc.kind()})); - } + return this->map_processor_to_device_coord(proc); } Realm::Event @@ -316,16 +291,49 @@ Realm::Event RealmContext::merge_outstanding_events() { } void RealmContext::discover_machine_topology() { - if (!this->processors.empty()) { + if (this->processors.has_value()) { return; } + std::unordered_map, nonnegative_int> + next_device_idx; + + auto fresh_device_id = [&](nonnegative_int node_idx, + DeviceType device_type) -> device_id_t { + std::pair key = + std::pair{node_idx, device_type}; + if (!contains_key(next_device_idx, key)) { + next_device_idx.insert({key, 0_n}); + } + + nonnegative_int device_idx = next_device_idx.at(key); + next_device_idx.at(key)++; + + return device_id_t{ + MachineSpaceCoordinate{node_idx, device_idx}, + device_type, + }; + }; + + bidict procs; Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); for (Realm::Processor proc : pq) { Realm::AddressSpace as = proc.address_space(); Realm::Processor::Kind kind = proc.kind(); - this->processors[std::pair{as, kind}].push_back(proc); + + nonnegative_int node_idx = nonnegative_int{static_cast(as)}; + + if (kind != Realm::Processor::LOC_PROC && + kind != Realm::Processor::TOC_PROC) { + continue; + } + + DeviceType device_type = device_type_from_processor_kind(kind); + device_id_t coord = fresh_device_id(node_idx, device_type); + procs.equate_strict(proc, coord); } + + this->processors = procs; } Realm::Runtime RealmContext::get_runtime() { diff --git a/lib/realm-execution/src/realm-execution/realm_manager.cc b/lib/realm-execution/src/realm-execution/realm_manager.cc index e76be7054b..cbc6cc2367 100644 --- a/lib/realm-execution/src/realm-execution/realm_manager.cc +++ b/lib/realm-execution/src/realm-execution/realm_manager.cc @@ -22,6 +22,7 @@ RealmManager::~RealmManager() { ControllerTaskResult RealmManager::start_controller(std::function thunk, Realm::Event wait_on) { + Realm::Processor target_proc = Realm::Machine::ProcessorQuery(Realm::Machine::get_machine()) .only_kind(Realm::Processor::LOC_PROC) diff --git a/lib/realm-execution/src/realm-execution/tasks/impl/controller_task.cc b/lib/realm-execution/src/realm-execution/tasks/impl/controller_task.cc index 925402bcb4..21253604b8 100644 --- a/lib/realm-execution/src/realm-execution/tasks/impl/controller_task.cc +++ b/lib/realm-execution/src/realm-execution/tasks/impl/controller_task.cc @@ -48,8 +48,10 @@ ControllerTaskResult sizeof(raw_ptr), precondition); - return ControllerTaskResult{std::unique_ptr(raw_ptr), - event}; + return ControllerTaskResult{ + std::unique_ptr(raw_ptr), + event, + }; } } // namespace FlexFlow 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 39eae4e1cb..f698d2a07f 100644 --- a/lib/realm-execution/test/src/realm-execution/test_e2e.cc +++ b/lib/realm-execution/test/src/realm-execution/test_e2e.cc @@ -36,6 +36,184 @@ static bool did_loss_decrease(GenericTensorAccessorR const &first_epoch, compare_tensor_accessors_le(last_epoch, first_epoch, allocator)); } +struct E2ETrainingConfig { + MappedParallelComputationGraph mapped_pcg; + LossAttrs loss_attrs; + MappedOperatorTaskGroup loss_mapping; + OptimizerAttrs optimizer_attrs; + parallel_tensor_guid_t logit_tensor; + TensorShape input_shape; + TensorShape logit_shape; + TensorShape label_shape; + TensorShape loss_shape; +}; + +static E2ETrainingConfig create_e2e_test_case() { + 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 input_tensor_shape = + TensorShape{TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; + + TensorShape label_tensor_shape = TensorShape{ + TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; + + TensorShape loss_tensor_shape = TensorShape{ + TensorDims{FFOrdered{output_dim, hidden_dim}}, DataType::FLOAT}; + + TensorShape weight_shape_1 = + TensorShape{TensorDims{FFOrdered{hidden_dim, data_dim}}, DataType::FLOAT}; + + TensorShape weight_shape_2 = TensorShape{ + TensorDims{FFOrdered{output_dim, hidden_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_1 = add_parallel_layer( + pcg, + ParallelLayerAttrs{ + PCGOperatorAttrs{WeightAttrs{weight_shape_1, + InitializerAttrs{GlorotNormalAttrs{0}}}}, + std::nullopt}, + {}, + {}); + parallel_tensor_guid_t t_weights_1 = + require_only_key(weights_layer_1.outputs, TensorSlotName::OUTPUT); + + ParallelLayerAddedResult weights_layer_2 = add_parallel_layer( + pcg, + ParallelLayerAttrs{ + PCGOperatorAttrs{WeightAttrs{weight_shape_2, + InitializerAttrs{GlorotNormalAttrs{0}}}}, + std::nullopt}, + {}, + {}); + parallel_tensor_guid_t t_weights_2 = + require_only_key(weights_layer_2.outputs, TensorSlotName::OUTPUT); + + ParallelLayerAddedResult linear_operator_1 = add_parallel_layer( + pcg, + ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{hidden_dim, + /*use_bias=*/false, + DataType::FLOAT, + Activation::RELU, + std::nullopt}}, + std::nullopt}, + { + { + TensorSlotName::INPUT, + t_input, + }, + }, + { + { + TensorSlotName::WEIGHT, + t_weights_1, + }, + }); + parallel_tensor_guid_t t_linear_1 = + require_only_key(linear_operator_1.outputs, TensorSlotName::OUTPUT); + + ParallelLayerAddedResult linear_operator_2 = add_parallel_layer( + pcg, + ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{output_dim, + /*use_bias=*/false, + DataType::FLOAT, + Activation::RELU, + std::nullopt}}, + std::nullopt}, + { + { + TensorSlotName::INPUT, + t_linear_1, + }, + }, + { + { + TensorSlotName::WEIGHT, + t_weights_2, + }, + }); + parallel_tensor_guid_t t_linear_2 = + require_only_key(linear_operator_2.outputs, TensorSlotName::OUTPUT); + + MachineSpaceCoordinate cpu0{0_n, 0_n}; + MachineSpaceCoordinate cpu1{0_n, 1_n}; + ParallelTensorSpaceCoordinate tensor_coord0{0_n, 0_n, FFOrdered{0_n}}; + + std::unordered_map mapping = { + {inputs_layer.parallel_layer, + MappedOperatorTaskGroup{ + {{cpu0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, + {weights_layer_1.parallel_layer, + MappedOperatorTaskGroup{ + {{cpu0, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, + {weights_layer_2.parallel_layer, + MappedOperatorTaskGroup{ + {{cpu1, + OperatorAtomicTaskShardBinding{ + {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, + {linear_operator_1.parallel_layer, + MappedOperatorTaskGroup{{{cpu0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord0}, + {TensorSlotName::WEIGHT, tensor_coord0}, + {TensorSlotName::OUTPUT, tensor_coord0}, + }}}}}}, + {linear_operator_2.parallel_layer, + MappedOperatorTaskGroup{{{cpu1, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord0}, + {TensorSlotName::WEIGHT, tensor_coord0}, + {TensorSlotName::OUTPUT, tensor_coord0}, + }}}}}}, + }; + + TensorShape output_tensor_shape = TensorShape{ + TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; + + MappedParallelComputationGraph mpcg = + mapped_pcg_from_pcg_and_mapped_op_task_groups(pcg, mapping); + + MappedOperatorTaskGroup loss_mapping{ + {{cpu0, + OperatorAtomicTaskShardBinding{{ + {TensorSlotName::INPUT, tensor_coord0}, + {TensorSlotName::LOGIT, tensor_coord0}, + }}}}}; + + LossAttrs loss_attrs = LossAttrs{ + NonconfigurableLossAttrs{LossFunction::CATEGORICAL_CROSSENTROPY}}; + OptimizerAttrs optimizer_attrs = + OptimizerAttrs{SGDOptimizerAttrs{/*lr=*/0.001, + /*momentum=*/0.9, + /*nesterov=*/false, + /*weight_decay=*/0.001}}; + + return E2ETrainingConfig{ + /*mapped_pcg=*/mpcg, + /*loss_attrs=*/loss_attrs, + /*loss_mapping=*/loss_mapping, + /*optimizer_attrs=*/optimizer_attrs, + /*logit_tensor=*/t_linear_2, + /*input_shape=*/input_tensor_shape, + /*logit_shape=*/output_tensor_shape, + /*label_shape=*/label_tensor_shape, + /*loss_shape=*/loss_tensor_shape, + }; +} + TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("RealmBackend e2e Training (CPU Model Parallelism)") { std::vector fake_args = @@ -46,163 +224,16 @@ TEST_SUITE(FF_TEST_SUITE) { RealmManager manager = RealmManager{&fake_argc, &fake_argv}; (void)manager.start_controller([](RealmContext &ctx) { - Allocator allocator = ctx.get_current_device_allocator(); - - positive_int batch_size = 10_p; - positive_int data_dim = 16_p; - positive_int hidden_dim = 32_p; - positive_int output_dim = 1_p; + ASSERT(ctx.processors.has_value()); - TensorShape output_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; + E2ETrainingConfig cfg = create_e2e_test_case(); - GenericTensorAccessorW label_tensor_backing = - allocator.allocate_tensor(output_tensor_shape); - - // construct computation graph - ParallelComputationGraph pcg = empty_parallel_computation_graph(); - - TensorShape input_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; + Allocator allocator = ctx.get_current_device_allocator(); - TensorShape label_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; + GenericTensorAccessorW output_tensor = + allocator.allocate_tensor(cfg.logit_shape); GenericTensorAccessorW label_tensor = - allocator.allocate_tensor(label_tensor_shape); - - TensorShape weight_shape_1 = TensorShape{ - TensorDims{FFOrdered{hidden_dim, data_dim}}, DataType::FLOAT}; - TensorShape weight_shape_2 = TensorShape{ - TensorDims{FFOrdered{output_dim, hidden_dim}}, DataType::FLOAT}; - - 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_1 = add_parallel_layer( - pcg, - ParallelLayerAttrs{ - PCGOperatorAttrs{WeightAttrs{ - weight_shape_1, InitializerAttrs{GlorotNormalAttrs{0}}}}, - std::nullopt}, - {}, - {}); - parallel_tensor_guid_t t_weights_1 = - require_only_key(weights_layer_1.outputs, TensorSlotName::OUTPUT); - - ParallelLayerAddedResult weights_layer_2 = add_parallel_layer( - pcg, - ParallelLayerAttrs{ - PCGOperatorAttrs{WeightAttrs{ - weight_shape_2, InitializerAttrs{GlorotNormalAttrs{0}}}}, - std::nullopt}, - {}, - {}); - parallel_tensor_guid_t t_weights_2 = - require_only_key(weights_layer_2.outputs, TensorSlotName::OUTPUT); - - ParallelLayerAddedResult linear_operator_1 = add_parallel_layer( - pcg, - ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{hidden_dim, - /*use_bias=*/false, - DataType::FLOAT, - Activation::RELU, - std::nullopt}}, - std::nullopt}, - { - { - TensorSlotName::INPUT, - t_input, - }, - }, - { - { - TensorSlotName::WEIGHT, - t_weights_1, - }, - }); - parallel_tensor_guid_t t_linear_1 = - require_only_key(linear_operator_1.outputs, TensorSlotName::OUTPUT); - - ParallelLayerAddedResult linear_operator_2 = add_parallel_layer( - pcg, - ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{output_dim, - /*use_bias=*/false, - DataType::FLOAT, - Activation::RELU, - std::nullopt}}, - std::nullopt}, - { - { - TensorSlotName::INPUT, - t_linear_1, - }, - }, - { - { - TensorSlotName::WEIGHT, - t_weights_2, - }, - }); - parallel_tensor_guid_t t_linear_2 = - require_only_key(linear_operator_2.outputs, TensorSlotName::OUTPUT); - - MachineSpaceCoordinate cpu0{0_n, 0_n, DeviceType::CPU}; - MachineSpaceCoordinate cpu1{0_n, 1_n, DeviceType::CPU}; - ParallelTensorSpaceCoordinate tensor_coord0{0_n, 0_n, 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{ - {{cpu0, - OperatorAtomicTaskShardBinding{ - {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, - {weights_layer_1.parallel_layer, - MappedOperatorTaskGroup{ - {{cpu0, - OperatorAtomicTaskShardBinding{ - {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, - {weights_layer_2.parallel_layer, - MappedOperatorTaskGroup{ - {{cpu1, - OperatorAtomicTaskShardBinding{ - {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, - {linear_operator_1.parallel_layer, - MappedOperatorTaskGroup{ - {{cpu0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::WEIGHT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}}}}}, - {linear_operator_2.parallel_layer, - MappedOperatorTaskGroup{ - {{cpu1, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::WEIGHT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}}}}}, - }); - - MappedOperatorTaskGroup loss_mapping{ - {{cpu0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::LOGIT, tensor_coord0}, - }}}}}; - - // instantiate computation graph - LossAttrs loss_attrs = LossAttrs{ - NonconfigurableLossAttrs{LossFunction::CATEGORICAL_CROSSENTROPY}}; - OptimizerAttrs optimizer_attrs = - OptimizerAttrs{SGDOptimizerAttrs{/*lr=*/0.001, - /*momentum=*/0.9, - /*nesterov=*/false, - /*weight_decay=*/0.001}}; + allocator.allocate_tensor(cfg.label_shape); std::unordered_map input_tensors; @@ -214,18 +245,19 @@ TEST_SUITE(FF_TEST_SUITE) { PCGInstance pcg_instance = create_pcg_instance( /*ctx=*/ctx, - /*mpcg=*/mpcg, - /*optimizer=*/optimizer_attrs, + /*mpcg=*/cfg.mapped_pcg, + /*optimizer=*/cfg.optimizer_attrs, /*loss=*/ ParallelLossConfig{ - /*loss_attrs=*/loss_attrs, + /*loss_attrs=*/cfg.loss_attrs, /*label_tensor=*/label_tensor, - /*logit_tensor=*/t_linear_2, - /*loss_mapping=*/loss_mapping, + /*logit_tensor=*/cfg.logit_tensor, + /*loss_mapping=*/cfg.loss_mapping, }, /*input_tensors=*/input_tensors, /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle); + /*device_handle=*/device_handle, + /*device_type=*/DeviceType::CPU); // begin training loop int num_epochs = 5; @@ -240,9 +272,7 @@ TEST_SUITE(FF_TEST_SUITE) { dynamic_tensor_accessor_from_instance( pcg_instance.get_loss_tensor_instance().value(), Realm::Event::NO_EVENT, - lift_to_parallel( - TensorShape{TensorDims{FFOrdered{output_dim, hidden_dim}}, - DataType::FLOAT}), + lift_to_parallel(cfg.loss_shape), Permissions::RO, ctx.get_current_processor()) .require_read(), @@ -265,155 +295,7 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_SUITE(FF_CUDA_TEST_SUITE) { TEST_CASE("RealmBackend e2e Training (GPU Model Parallelism)") { - 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}; - - // construct computation graph - ParallelComputationGraph pcg = empty_parallel_computation_graph(); - - TensorShape input_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT}; - - TensorShape label_tensor_shape = TensorShape{ - TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT}; - - TensorShape weight_shape_1 = TensorShape{ - TensorDims{FFOrdered{hidden_dim, data_dim}}, DataType::FLOAT}; - TensorShape weight_shape_2 = TensorShape{ - TensorDims{FFOrdered{output_dim, hidden_dim}}, DataType::FLOAT}; - - 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_1 = add_parallel_layer( - pcg, - ParallelLayerAttrs{ - PCGOperatorAttrs{WeightAttrs{ - weight_shape_1, InitializerAttrs{GlorotNormalAttrs{0}}}}, - std::nullopt}, - {}, - {}); - parallel_tensor_guid_t t_weights_1 = - require_only_key(weights_layer_1.outputs, TensorSlotName::OUTPUT); - - ParallelLayerAddedResult weights_layer_2 = add_parallel_layer( - pcg, - ParallelLayerAttrs{ - PCGOperatorAttrs{WeightAttrs{ - weight_shape_2, InitializerAttrs{GlorotNormalAttrs{0}}}}, - std::nullopt}, - {}, - {}); - parallel_tensor_guid_t t_weights_2 = - require_only_key(weights_layer_2.outputs, TensorSlotName::OUTPUT); - - ParallelLayerAddedResult linear_operator_1 = add_parallel_layer( - pcg, - ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{hidden_dim, - /*use_bias=*/false, - DataType::FLOAT, - Activation::RELU, - std::nullopt}}, - std::nullopt}, - { - { - TensorSlotName::INPUT, - t_input, - }, - }, - { - { - TensorSlotName::WEIGHT, - t_weights_1, - }, - }); - parallel_tensor_guid_t t_linear_1 = - require_only_key(linear_operator_1.outputs, TensorSlotName::OUTPUT); - - ParallelLayerAddedResult linear_operator_2 = add_parallel_layer( - pcg, - ParallelLayerAttrs{PCGOperatorAttrs{LinearAttrs{output_dim, - /*use_bias=*/false, - DataType::FLOAT, - Activation::RELU, - std::nullopt}}, - std::nullopt}, - { - { - TensorSlotName::INPUT, - t_linear_1, - }, - }, - { - { - TensorSlotName::WEIGHT, - t_weights_2, - }, - }); - parallel_tensor_guid_t t_linear_2 = - require_only_key(linear_operator_2.outputs, TensorSlotName::OUTPUT); - - MachineSpaceCoordinate gpu0{0_n, 0_n, DeviceType::GPU}; - ParallelTensorSpaceCoordinate tensor_coord0{0_n, 0_n, 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{ - {{gpu0, - OperatorAtomicTaskShardBinding{ - {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, - {weights_layer_1.parallel_layer, - MappedOperatorTaskGroup{ - {{gpu0, - OperatorAtomicTaskShardBinding{ - {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, - {weights_layer_2.parallel_layer, - MappedOperatorTaskGroup{ - {{gpu0, - OperatorAtomicTaskShardBinding{ - {{TensorSlotName::OUTPUT, tensor_coord0}}}}}}}, - {linear_operator_1.parallel_layer, - MappedOperatorTaskGroup{ - {{gpu0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::WEIGHT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}}}}}, - {linear_operator_2.parallel_layer, - MappedOperatorTaskGroup{ - {{gpu0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::WEIGHT, tensor_coord0}, - {TensorSlotName::OUTPUT, tensor_coord0}, - }}}}}}, - }); - - MappedOperatorTaskGroup loss_mapping{ - {{gpu0, - OperatorAtomicTaskShardBinding{{ - {TensorSlotName::INPUT, tensor_coord0}, - {TensorSlotName::LOGIT, tensor_coord0}, - }}}}}; - - // instantiate computation graph - LossAttrs loss_attrs = LossAttrs{ - NonconfigurableLossAttrs{LossFunction::CATEGORICAL_CROSSENTROPY}}; - OptimizerAttrs optimizer_attrs = - OptimizerAttrs{SGDOptimizerAttrs{/*lr=*/0.001, - /*momentum=*/0.9, - /*nesterov=*/false, - /*weight_decay=*/0.001}}; + E2ETrainingConfig cfg = create_e2e_test_case(); //! [realm-execution example] std::vector fake_args = @@ -427,11 +309,10 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { manager.start_controller([&](RealmContext &ctx) { Allocator allocator = ctx.get_current_device_allocator(); - GenericTensorAccessorW label_tensor_backing = - allocator.allocate_tensor(output_tensor_shape); - + GenericTensorAccessorW logit_tensor = + allocator.allocate_tensor(cfg.logit_shape); GenericTensorAccessorW label_tensor = - allocator.allocate_tensor(label_tensor_shape); + allocator.allocate_tensor(cfg.label_shape); std::unordered_map input_tensors; @@ -443,18 +324,19 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { PCGInstance pcg_instance = create_pcg_instance( /*ctx=*/ctx, - /*mpcg=*/mpcg, - /*optimizer=*/optimizer_attrs, + /*mpcg=*/cfg.mapped_pcg, + /*optimizer=*/cfg.optimizer_attrs, /*loss=*/ ParallelLossConfig{ - /*loss_attrs=*/loss_attrs, + /*loss_attrs=*/cfg.loss_attrs, /*label_tensor=*/label_tensor, - /*logit_tensor=*/t_linear_2, - /*loss_mapping=*/loss_mapping, + /*logit_tensor=*/cfg.logit_tensor, + /*loss_mapping=*/cfg.loss_mapping, }, /*input_tensors=*/input_tensors, /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle); + /*device_handle=*/device_handle, + /*device_type=*/DeviceType::GPU); // begin training loop int num_epochs = 5; @@ -465,13 +347,12 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*instance=*/pcg_instance, /*profiling_settings=*/ProfilingSettings{0, 0}, /*device_handle=*/device_handle); + loss_values.push_back(copy_tensor_accessor_r( dynamic_tensor_accessor_from_instance( pcg_instance.get_loss_tensor_instance().value(), Realm::Event::NO_EVENT, - lift_to_parallel(TensorShape{ - TensorDims{FFOrdered{output_dim, hidden_dim}}, - DataType::FLOAT}), + lift_to_parallel(cfg.loss_shape), Permissions::RO, ctx.get_current_processor()) .require_read(), diff --git a/lib/task-spec/include/task-spec/concrete_arg_spec.h b/lib/task-spec/include/task-spec/concrete_arg_spec.h index 45bbd6ba6b..32df07d325 100644 --- a/lib/task-spec/include/task-spec/concrete_arg_spec.h +++ b/lib/task-spec/include/task-spec/concrete_arg_spec.h @@ -1,7 +1,6 @@ #ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_CONCRETE_ARG_SPEC_H #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_CONCRETE_ARG_SPEC_H -#include "task-spec/serialization.h" #include "utils/hash-utils.h" #include "utils/type_index.h" #include diff --git a/lib/task-spec/include/task-spec/device_id_t.dtg.toml b/lib/task-spec/include/task-spec/device_id_t.dtg.toml new file mode 100644 index 0000000000..f6d29ee1a7 --- /dev/null +++ b/lib/task-spec/include/task-spec/device_id_t.dtg.toml @@ -0,0 +1,23 @@ +namespace = "FlexFlow" +name = "device_id_t" +type = "struct" +features = [ + "eq", + "ord", + "hash", + "json", + "fmt", +] + +includes = [ + "pcg/machine_space_coordinate.dtg.h", + "pcg/device_type.dtg.h", +] + +[[fields]] +name = "coord" +type = "::FlexFlow::MachineSpaceCoordinate" + +[[fields]] +name = "device_type" +type = "::FlexFlow::DeviceType" diff --git a/lib/task-spec/include/task-spec/device_specific.h b/lib/task-spec/include/task-spec/device_specific.h index 2055888b1b..49a2555411 100644 --- a/lib/task-spec/include/task-spec/device_specific.h +++ b/lib/task-spec/include/task-spec/device_specific.h @@ -1,9 +1,9 @@ #ifndef _FLEXFLOW_LOCAL_EXECUTION_DEVICE_SPECIFIC_H #define _FLEXFLOW_LOCAL_EXECUTION_DEVICE_SPECIFIC_H -#include "pcg/device_id_t.dtg.h" -#include "task-spec/serialization.h" +#include "task-spec/device_id_t.dtg.h" #include "utils/hash/tuple.h" +#include namespace FlexFlow { @@ -12,7 +12,8 @@ struct DeviceSpecific { DeviceSpecific() = delete; template - static DeviceSpecific create(device_id_t device_idx, Args &&...args) { + static DeviceSpecific create(device_id_t const &device_idx, + Args &&...args) { return DeviceSpecific(std::make_shared(std::forward(args)...), device_idx); } @@ -25,13 +26,13 @@ struct DeviceSpecific { return this->tie() != other.tie(); } - T const *get(device_id_t curr_device_idx) const { + T const *get(device_id_t const &curr_device_idx) const { ASSERT(curr_device_idx == this->device_idx); return (T const *)this->ptr.get(); } private: - DeviceSpecific(std::shared_ptr ptr, device_id_t device_idx) + DeviceSpecific(std::shared_ptr ptr, device_id_t const &device_idx) : ptr(ptr), device_idx(device_idx) {} private: @@ -57,11 +58,6 @@ std::ostream &operator<<(std::ostream &s, DeviceSpecific const &d) { return (s << fmt::to_string(d)); } -// manually force serialization to make DeviceSpecific trivially -// serializable -// template -// struct is_trivially_serializable> : std::true_type {}; - } // namespace FlexFlow namespace std { diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.dtg.toml new file mode 100644 index 0000000000..578bfe7585 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.dtg.toml @@ -0,0 +1,29 @@ +namespace = "FlexFlow" +name = "DynamicGraphEdge" +type = "struct" +features = [ + "eq", + "hash", + "fmt", +] + +includes = [ + "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h", + "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h", + "task-spec/dynamic_graph/dynamic_slot_site.h", +] + +src_includes = [ +] + +[[fields]] +name = "src" +type = "::FlexFlow::DynamicSlotSite" + +[[fields]] +name = "dst_node" +type = "::FlexFlow::DynamicNodeInvocation" + +[[fields]] +name = "dst_slot" +type = "::FlexFlow::DynamicTensorSlot" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.h new file mode 100644 index 0000000000..b561dba8c4 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_graph_edge.h @@ -0,0 +1,15 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_GRAPH_EDGE_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_GRAPH_EDGE_H + +#include "task-spec/dynamic_graph/dynamic_graph_edge.dtg.h" +#include "task-spec/dynamic_graph/dynamic_slot_site.dtg.h" + +namespace FlexFlow { + +DynamicGraphEdge + dynamic_graph_edge_from_slot_sites(DynamicSlotSite const &src, + InternalDynamicSlotSite const &dst); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml index 73c023fd40..7655bd24e1 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml @@ -10,11 +10,11 @@ features = [ includes = [ "", "task-spec/dynamic_graph/dynamic_task_type.dtg.h", - "pcg/machine_space_coordinate.dtg.h", - "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h", "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h", "task-spec/dynamic_graph/training_operation_attrs.dtg.h", "task-spec/device_specific_per_device_op_state.dtg.h", + "task-spec/device_id_t.dtg.h", + "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h", ] src_includes = [ @@ -27,8 +27,13 @@ type = "std::optional<::FlexFlow::DynamicTaskType>" [[fields]] name = "device_coord" -type = "std::optional<::FlexFlow::MachineSpaceCoordinate>" +type = "std::optional<::FlexFlow::device_id_t>" docstring = ''' +\brief The device on which this task should execute. + +For a standard \ref ParallelComputationGraph, filled in by +\ref perform_shard_expansion. + \note Right now the \c device_coord for a copy node is sort of meaningless because we have one controller issuing all copies for the entire graph, no matter where they are. However the intention is this to be the "owner" or @@ -41,7 +46,7 @@ choices. [[fields]] name = "mapping" -type = "std::optional<::FlexFlow::MappedOperatorTaskGroup>" +type = "std::optional<::FlexFlow::DynamicNodeMapping>" [[fields]] name = "op_attrs" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_invocation.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_invocation.h new file mode 100644 index 0000000000..6f27b43ad3 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_invocation.h @@ -0,0 +1,22 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_INVOCATION_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_INVOCATION_H + +#include "pcg/tensor_direction.dtg.h" +#include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h" +#include "task-spec/dynamic_graph/dynamic_slot_site.dtg.h" +#include "task-spec/dynamic_graph/training_op_type.dtg.h" + +namespace FlexFlow { + +std::unordered_map + get_slot_map_for_direction(DynamicNodeInvocation const &, TensorDirection); + +TrainingOpType + dynamic_node_invocation_get_op_type(DynamicNodeInvocation const &); + +std::unordered_set + get_dynamic_slot_sites_for_invocation(DynamicNodeInvocation const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.dtg.toml new file mode 100644 index 0000000000..859902e2bf --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.dtg.toml @@ -0,0 +1,26 @@ +namespace = "FlexFlow" +name = "DynamicNodeMapping" +type = "struct" +features = [ + "eq", + "ord", + "hash", + "json", + "fmt", +] + +includes = [ + "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h", + "pcg/device_type.dtg.h", +] + +src_includes = [ +] + +[[fields]] +name = "op_task_group" +type = "::FlexFlow::MappedOperatorTaskGroup" + +[[fields]] +name = "device_type" +type = "::FlexFlow::DeviceType" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h new file mode 100644 index 0000000000..6a52773a71 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h @@ -0,0 +1,18 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H + +#include "task-spec/device_id_t.dtg.h" +#include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" + +namespace FlexFlow { + +bidict + dynamic_node_mapping_bindings_for_slot_name(DynamicNodeMapping const &, + TensorSlotName const &); + +std::unordered_set + target_devices_of_dynamic_node_mapping(DynamicNodeMapping const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_open_dataflow_graph.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_open_dataflow_graph.h index 4ca62db5b1..6141c70e54 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_open_dataflow_graph.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_open_dataflow_graph.h @@ -1,14 +1,19 @@ #ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_OPEN_DATAFLOW_GRAPH_H #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_OPEN_DATAFLOW_GRAPH_H +#include "task-spec/dynamic_graph/dynamic_graph_edge.dtg.h" #include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" +#include "task-spec/dynamic_graph/dynamic_slot_site.dtg.h" #include "utils/graph/labelled_open_kwarg_dataflow_graph/labelled_open_kwarg_dataflow_graph.h" namespace FlexFlow { DynamicOpenDataflowGraph make_empty_dynamic_open_dataflow_graph(); +void check_dynamic_open_dataflow_graph_is_valid( + DynamicOpenDataflowGraph const &); + nonnegative_int dynamic_graph_num_nodes(DynamicOpenDataflowGraph const &); bool full_dynamic_graph_satisfies( @@ -32,6 +37,28 @@ std::unordered_multiset std::unordered_set get_dynamic_invocation_set(DynamicOpenDataflowGraph const &); +std::unordered_set + get_dynamic_graph_edges(DynamicOpenDataflowGraph const &); +std::unordered_set + get_dynamic_graph_edges_incoming_to_invocation( + DynamicOpenDataflowGraph const &, DynamicNodeInvocation const &); +std::unordered_set + get_dynamic_graph_edges_outgoing_from_invocation( + DynamicOpenDataflowGraph const &, DynamicNodeInvocation const &); + +std::unordered_set + get_internal_dynamic_slot_sites(DynamicOpenDataflowGraph const &); + +std::unordered_set + get_dynamic_slot_sites(DynamicOpenDataflowGraph const &); + +DynamicSlotSite + dynamic_graph_find_source_of_value(DynamicOpenDataflowGraph const &, + DynamicValueAttrs const &); +std::unordered_set + dynamic_graph_find_sinks_of_value(DynamicOpenDataflowGraph const &, + DynamicValueAttrs const &); + std::optional find_output_value_attrs(DynamicOpenDataflowGraph const &, dynamic_tensor_guid_t, diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.dtg.toml new file mode 100644 index 0000000000..fa56b0f105 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.dtg.toml @@ -0,0 +1,21 @@ +namespace = "FlexFlow" +name = "DynamicSlotSite" +type = "variant" +features = [ + "eq", + "hash", + "fmt", +] + +includes = [ + "task-spec/dynamic_graph/external_dynamic_slot_site.dtg.h", + "task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.h", +] + +[[values]] +type = "::FlexFlow::InternalDynamicSlotSite" +key = "internal" + +[[values]] +type = "::FlexFlow::ExternalDynamicSlotSite" +key = "external" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.h new file mode 100644 index 0000000000..4340cdcf39 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_slot_site.h @@ -0,0 +1,13 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_SLOT_SITE_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_SLOT_SITE_H + +#include "task-spec/dynamic_graph/dynamic_slot_site.dtg.h" +#include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" + +namespace FlexFlow { + +DynamicValueAttrs dynamic_value_attrs_for_slot_site(DynamicSlotSite const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml index 490a51f88d..926cc9893c 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml @@ -16,6 +16,7 @@ includes = [ "utils/bidict/bidict.h", "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h", "task-spec/dynamic_graph/dynamic_tensor_role.dtg.h", + "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h", ] src_includes = [ @@ -36,7 +37,13 @@ type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>" [[fields]] name = "mapping" -type = "std::optional<::FlexFlow::bidict<::FlexFlow::ParallelTensorSpaceCoordinate, ::FlexFlow::MachineSpaceCoordinate>>" +type = "std::optional<::FlexFlow::ParallelTensorMapping>" +docstring = ''' +\brief Which devices the shards of the parallel tensor are stored on. + +For computations originating from a \ref ParallelComputationGraph, +this field is usually resolved/assigned by \ref perform_copy_insertion. +''' [[fields]] name = "accessor" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h index 9cccc565cc..d6090e5f14 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h @@ -2,12 +2,17 @@ #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_VALUE_ATTRS_H #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" +#include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" namespace FlexFlow { DynamicValueAttrs decide_dynamic_value_attrs_role(DynamicValueAttrs const &, DynamicTensorRole); +DynamicValueAttrs + dynamic_value_attrs_with_mapping(DynamicValueAttrs const &, + ParallelTensorMapping const &); + } // namespace FlexFlow #endif diff --git a/lib/task-spec/include/task-spec/dynamic_graph/external_dynamic_slot_site.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/external_dynamic_slot_site.dtg.toml new file mode 100644 index 0000000000..bdcc93251c --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/external_dynamic_slot_site.dtg.toml @@ -0,0 +1,19 @@ +namespace = "FlexFlow" +name = "ExternalDynamicSlotSite" +type = "struct" +features = [ + "eq", + "hash", + "fmt", +] + +includes = [ + "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h", +] + +src_includes = [ +] + +[[fields]] +name = "value" +type = "::FlexFlow::DynamicValueAttrs" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml new file mode 100644 index 0000000000..bf29bf2b2f --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml @@ -0,0 +1,29 @@ +namespace = "FlexFlow" +name = "InternalDynamicSlotSite" +type = "struct" +features = [ + "eq", + "hash", + "fmt", +] + +includes = [ + "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h", + "pcg/tensor_direction.dtg.h", + "task-spec/dynamic_graph/dynamic_tensor_slot.h", +] + +src_includes = [ +] + +[[fields]] +name = "invocation" +type = "::FlexFlow::DynamicNodeInvocation" + +[[fields]] +name = "direction" +type = "::FlexFlow::TensorDirection" + +[[fields]] +name = "slot_name" +type = "::FlexFlow::DynamicTensorSlot" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/loss_insertion.h b/lib/task-spec/include/task-spec/dynamic_graph/loss_insertion.h index b3b2a465f8..8584b4bed8 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/loss_insertion.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/loss_insertion.h @@ -14,7 +14,7 @@ LossInsertionResult perform_loss_insertion( DynamicOpenDataflowGraph const &dg, LossAttrs const &loss_attrs, dynamic_tensor_guid_t logit_tensor, - std::optional const &loss_mapping); + std::optional const &loss_mapping); } // namespace FlexFlow diff --git a/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h b/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h index 823f962c25..b40fabe7bf 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h @@ -7,11 +7,11 @@ namespace FlexFlow { std::unordered_set perform_machine_slicing_for_invocation(DynamicNodeInvocation const &, - MachineSpaceCoordinate const &); + device_id_t const &); DynamicOpenDataflowGraph perform_machine_slicing(DynamicOpenDataflowGraph const &, - MachineSpaceCoordinate const &); + device_id_t const &); } // namespace FlexFlow diff --git a/lib/task-spec/include/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h b/lib/task-spec/include/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h index 6a269ec3c9..c5db832a6a 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h @@ -7,7 +7,7 @@ namespace FlexFlow { DynamicOpenDataflowGraph make_dynamic_open_dataflow_graph_from_mapped_pcg( - MappedParallelComputationGraph const &); + MappedParallelComputationGraph const &, DeviceType device_type); } // namespace FlexFlow diff --git a/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml new file mode 100644 index 0000000000..ba9785bc92 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml @@ -0,0 +1,22 @@ +namespace = "FlexFlow" +name = "ParallelTensorMapping" +type = "struct" +features = [ + "eq", + "hash", + "json", + "fmt", +] + +includes = [ + "utils/bidict/bidict.h", + "op-attrs/parallel_tensor_space_coordinate.dtg.h", + "task-spec/device_id_t.dtg.h", +] + +src_includes = [ +] + +[[fields]] +name = "raw" +type = "::FlexFlow::bidict<::FlexFlow::ParallelTensorSpaceCoordinate, ::FlexFlow::device_id_t>" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml index 3c43e1d637..43fb8490ec 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml @@ -11,8 +11,8 @@ features = [ includes = [ "", "task-spec/dynamic_graph/dynamic_task_type.dtg.h", - "pcg/machine_space_coordinate.dtg.h", - "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h", + "task-spec/device_id_t.dtg.h", + "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h", "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h", "task-spec/dynamic_graph/training_operation_attrs.dtg.h", ] @@ -28,11 +28,11 @@ type = "std::optional<::FlexFlow::DynamicTaskType>" [[fields]] name = "device_coord" -type = "std::optional<::FlexFlow::MachineSpaceCoordinate>" +type = "std::optional<::FlexFlow::device_id_t>" [[fields]] name = "mapping" -type = "std::optional<::FlexFlow::MappedOperatorTaskGroup>" +type = "std::optional<::FlexFlow::DynamicNodeMapping>" [[fields]] name = "op_attrs" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml index 454f1b7e8c..c884bd8abe 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml @@ -16,6 +16,7 @@ includes = [ "pcg/machine_space_coordinate.dtg.h", "utils/bidict/bidict.h", "task-spec/dynamic_graph/dynamic_tensor_role.dtg.h", + "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h", ] src_includes = [ @@ -37,7 +38,7 @@ type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>" [[fields]] name = "mapping" -type = "std::optional<::FlexFlow::bidict<::FlexFlow::ParallelTensorSpaceCoordinate, ::FlexFlow::MachineSpaceCoordinate>>" +type = "std::optional<::FlexFlow::ParallelTensorMapping>" [[fields]] name = "role" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml new file mode 100644 index 0000000000..f1013771d1 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml @@ -0,0 +1,15 @@ +namespace = "FlexFlow" +name = "TrainingOnlyOpType" +type = "enum" +features = [ + "hash", + "json", + "rapidcheck", + "fmt", +] + +[[values]] +name = "COPY" + +[[values]] +name = "LOSS" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/training_op_type.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/training_op_type.dtg.toml new file mode 100644 index 0000000000..7744f3775f --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/training_op_type.dtg.toml @@ -0,0 +1,22 @@ +namespace = "FlexFlow" +name = "TrainingOpType" +type = "variant" +features = [ + "eq", + "hash", + "fmt", + "json", +] + +includes = [ + "task-spec/dynamic_graph/training_only_op_type.dtg.h", + "op-attrs/operator_type.dtg.h", +] + +[[values]] +type = "::FlexFlow::OperatorType" +key = "pcg_op" + +[[values]] +type = "::FlexFlow::TrainingOnlyOpType" +key = "training_only_op" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.h b/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.h new file mode 100644 index 0000000000..92446cb1b7 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.h @@ -0,0 +1,13 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_TRAINING_OPERATION_ATTRS_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_TRAINING_OPERATION_ATTRS_H + +#include "task-spec/dynamic_graph/training_op_type.dtg.h" +#include "task-spec/dynamic_graph/training_operation_attrs.dtg.h" + +namespace FlexFlow { + +TrainingOpType training_op_attrs_get_op_type(TrainingOperationAttrs const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/serialization.h b/lib/task-spec/include/task-spec/serialization.h deleted file mode 100644 index 29f9144a3b..0000000000 --- a/lib/task-spec/include/task-spec/serialization.h +++ /dev/null @@ -1,113 +0,0 @@ -#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_SERIALIZATION_H -#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_SERIALIZATION_H - -#include "kernels/device.h" -#include "kernels/nccl.h" -#include "op-attrs/ff_ordered/ff_ordered.h" -#include "utils/required.h" -#include "utils/type_traits.h" -#include "utils/variant.h" - -namespace FlexFlow { - -template -struct needs_serialization {}; - -template -struct visit_trivially_serializable; - -template -struct is_trivially_serializable : std::false_type {}; - -template -struct visit_trivially_serializable { - static constexpr bool value = is_trivially_serializable::value && - visit_trivially_serializable::value; -}; - -template -struct visit_trivially_serializable> { - static constexpr bool value = visit_trivially_serializable::value; -}; - -template <> -struct visit_trivially_serializable<> : std::true_type {}; - -template -struct is_trivially_serializable< - T, - typename std::enable_if::value>::type> - : std::true_type {}; - -template <> -struct is_trivially_serializable : std::true_type {}; -template <> -struct is_trivially_serializable : std::true_type {}; - -template -struct is_trivially_serializable< - T, - typename std::enable_if::value>::type> : std::true_type {}; - -template -struct is_trivially_serializable< - T, - typename std::enable_if::value>::type> - : std::true_type {}; - -template -struct is_trivially_serializable> - : is_trivially_serializable {}; - -template -struct is_trivially_serializable> : is_trivially_serializable { -}; - -template -struct is_trivially_serializable> - : elements_satisfy> {}; - -template -struct is_trivially_serializable> - : is_trivially_serializable {}; - -template -struct std_array_size_helper; - -template -struct std_array_size_helper> { - static const std::size_t value = N; -}; - -template -using std_array_size = std_array_size_helper; - -template -struct is_trivially_serializable< - T, - std::enable_if::value>>::value>> - : std::true_type {}; - -template -struct is_serializable : std::false_type {}; - -template -struct is_serializable< - T, - typename std::enable_if::value>::type> - : std::true_type {}; - -static_assert(is_trivially_serializable::value, ""); -static_assert(is_trivially_serializable::value, ""); -static_assert(is_trivially_serializable::value, ""); -static_assert(is_trivially_serializable::value, ""); -static_assert(is_trivially_serializable::value, ""); -static_assert(is_trivially_serializable::value, ""); -static_assert(is_trivially_serializable>::value, - ""); - -} // namespace FlexFlow - -#endif diff --git a/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h b/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h index b4c8dcdf36..165631889f 100644 --- a/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h +++ b/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h @@ -7,9 +7,9 @@ #include "op-attrs/ops/loss_functions/loss_attrs.dtg.h" #include "op-attrs/pcg_operator_attrs.dtg.h" #include "op-attrs/tensor_slot_name.dtg.h" -#include "pcg/device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/concrete_arg_spec.h" +#include "task-spec/device_id_t.dtg.h" #include "task-spec/ops/arg_slot_id_t.dtg.h" #include "task-spec/per_device_op_state.dtg.h" #include "task-spec/privilege_tensor_accessor.h" 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 79bd091c0a..f611735e26 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,21 +3,35 @@ #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 "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" +#include "task-spec/dynamic_graph/dynamic_node_invocation.h" +#include "task-spec/dynamic_graph/dynamic_node_mapping.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" +#include "task-spec/dynamic_graph/dynamic_slot_site.h" #include "task-spec/dynamic_graph/dynamic_task_type.h" #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_tensor_mapping.dtg.h" #include "utils/bidict/algorithms/bidict_from_pairs.h" #include "utils/bidict/algorithms/unordered_set_of.h" #include "utils/containers/contains_key.h" +#include "utils/containers/count.h" +#include "utils/containers/filter_values.h" +#include "utils/containers/filtermap_keys.h" +#include "utils/containers/filtrans.h" #include "utils/containers/flatmap.h" #include "utils/containers/map_values2.h" +#include "utils/containers/merge_disjoint_maps.h" #include "utils/containers/set_difference.h" #include "utils/containers/set_intersection.h" #include "utils/containers/transform.h" +#include "utils/containers/values.h" +#include "utils/containers/zip_values_strict_with.h" #include "utils/optional.h" +#include "utils/overload.h" namespace FlexFlow { @@ -44,107 +58,171 @@ bool graph_is_fully_copy_inserted(DynamicOpenDataflowGraph const &g) { g, node_is_any, value_is_mapped, slot_is_mapped); } -static DynamicValueAttrs map_dynamic_value_attrs_for_task_group( - DynamicTensorSlot const &slot, - DynamicValueAttrs const &value, - MappedOperatorTaskGroup const &mapping) { - DynamicValueAttrs result = value; - result.mapping = get_tensor_bindings_for_slot_name(mapping, slot.slot_name); - return result; -} - static std::pair filter_mapping_to_avoid_degenerate_copies(DynamicValueAttrs const &input, DynamicValueAttrs const &output) { - std::unordered_set< - std::pair> - input_mapping = unordered_set_of(assert_unwrap(input.mapping)); - std::unordered_set< - std::pair> - output_mapping = unordered_set_of(assert_unwrap(output.mapping)); + std::unordered_set> + input_mapping = unordered_set_of(assert_unwrap(input.mapping).raw); + + std::unordered_set> + output_mapping = unordered_set_of(assert_unwrap(output.mapping).raw); // Exclude the point shared between the input and output mappings, because // those will not result in actual copies once shard expansion is performed std::unordered_set< - std::pair> + std::pair> remove = set_intersection(input_mapping, output_mapping); DynamicValueAttrs filtered_input = input; - filtered_input.mapping = - bidict_from_pairs(set_difference(input_mapping, remove)); + filtered_input.mapping = ParallelTensorMapping{ + bidict_from_pairs(set_difference(input_mapping, remove)), + }; DynamicValueAttrs filtered_output = output; - filtered_output.mapping = - bidict_from_pairs(set_difference(output_mapping, remove)); + filtered_output.mapping = ParallelTensorMapping{ + bidict_from_pairs(set_difference(output_mapping, remove)), + }; return std::pair{filtered_input, filtered_output}; } -std::unordered_set perform_copy_insertion_for_invocation( +std::unordered_map + get_mappings_for_invocation( + DynamicNodeInvocation const &i, + std::unordered_map const + &mappings) { + return filtermap_keys(mappings, + [&](InternalDynamicSlotSite const &s) + -> std::optional { + if (s.invocation == i) { + return s.slot_name; + } else { + return std::nullopt; + } + }); +} + +DynamicNodeInvocation apply_mappings_for_invocation( DynamicNodeInvocation const &i, - std::unordered_map const - &unmapped_value_to_mapped_source_value) { + std::unordered_map const + &all_mappings) { + std::unordered_map i_mappings = + get_mappings_for_invocation(i, all_mappings); + + std::unordered_map + i_input_mappings = restrict_keys(i_mappings, keys(i.inputs)); - MappedOperatorTaskGroup mapping = assert_unwrap(i.node_attrs.mapping); + std::unordered_map + i_output_mappings = restrict_keys(i_mappings, keys(i.outputs)); - auto map_tensor = [&](DynamicTensorSlot const &slot, - DynamicValueAttrs const &value) { - return map_dynamic_value_attrs_for_task_group(slot, value, mapping); + auto apply_mapping = + [&](DynamicValueAttrs const &v, + ParallelTensorMapping const &mapping) -> DynamicValueAttrs { + return dynamic_value_attrs_with_mapping(v, mapping); }; - std::unordered_map mapped_inputs = - map_values2(i.inputs, map_tensor); - std::unordered_map mapped_outputs = - map_values2(i.outputs, map_tensor); + return DynamicNodeInvocation{ + /*inputs=*/ + zip_values_strict_with(i.inputs, i_input_mappings, apply_mapping), + /*node_attrs=*/ + i.node_attrs, + /*outputs=*/ + zip_values_strict_with(i.outputs, i_output_mappings, apply_mapping), + }; +} - std::unordered_set result{DynamicNodeInvocation{ - /*inputs=*/mapped_inputs, - /*node_attrs=*/i.node_attrs, - /*outputs=*/mapped_outputs, - }}; +std::unordered_set copies_for_value( + DynamicOpenDataflowGraph const &g, + DynamicValueAttrs const &v, + std::unordered_map const + &mappings) { + InternalDynamicSlotSite src = ({ + DynamicSlotSite found = dynamic_graph_find_source_of_value(g, v); - for (auto const &[slot, input] : i.inputs) { - if (!contains_key(unmapped_value_to_mapped_source_value, input)) { - continue; + if (found.is_external()) { + return {}; } - DynamicValueAttrs source_value = - unmapped_value_to_mapped_source_value.at(input); - DynamicValueAttrs use_value = mapped_inputs.at(slot); - if (source_value != use_value) { - auto const &[filtered_source, filtered_use] = - filter_mapping_to_avoid_degenerate_copies(source_value, use_value); - DynamicNodeInvocation copy{ - /*inputs=*/{ - { - DynamicTensorSlot{TensorSlotName::INPUT, - slot.slot_tensor_role}, - filtered_source, - }, - }, - /*node_attrs=*/ - DynamicNodeAttrs{ - /*task_type=*/transform( - slot.slot_tensor_role, - dynamic_task_type_from_tensor_role_for_copy), - /*device_coord=*/std::nullopt, - /*mapping=*/std::nullopt, - /*op_attrs*/ TrainingOperationAttrs{CopyAttrs{}}, - /*layer_guid=*/dynamic_layer_guid_t{dynamic_copy_layer_guid_t{}}, - /*per_device_op_state=*/std::nullopt, - }, - /*outputs=*/ - { - { - DynamicTensorSlot{TensorSlotName::OUTPUT, - slot.slot_tensor_role}, - filtered_use, - }, - }, - }; - result.insert(copy); - } - } + found.require_internal(); + }); + + std::unordered_set sinks = + dynamic_graph_find_sinks_of_value(g, DynamicValueAttrs{v}); + + ParallelTensorMapping src_mapping = mappings.at(src); + + std::unordered_map + mappings_for_sinks = generate_map( + sinks, + [&](InternalDynamicSlotSite const &s) -> ParallelTensorMapping { + return mappings.at(s); + }); + + std::unordered_set sink_mapping_set = + unordered_set_of(values(mappings_for_sinks)); + + std::unordered_set required_copies = + set_difference(sink_mapping_set, std::unordered_set{src_mapping}); + + auto make_copy_to = + [&](ParallelTensorMapping const &sink_mapping) -> DynamicNodeInvocation { + return DynamicNodeInvocation{ + /*inputs=*/{ + { + DynamicTensorSlot{ + TensorSlotName::INPUT, + src.slot_name.slot_tensor_role, + }, + dynamic_value_attrs_with_mapping(v, src_mapping), + }, + }, + /*node_attrs=*/ + DynamicNodeAttrs{ + /*task_type=*/transform( + src.slot_name.slot_tensor_role, + dynamic_task_type_from_tensor_role_for_copy), + /*device_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*op_attrs*/ TrainingOperationAttrs{CopyAttrs{}}, + /*layer_guid=*/dynamic_layer_guid_t{dynamic_copy_layer_guid_t{}}, + /*per_device_op_state=*/std::nullopt, + }, + /*outputs=*/ + { + { + DynamicTensorSlot{ + TensorSlotName::OUTPUT, + src.slot_name.slot_tensor_role, + }, + dynamic_value_attrs_with_mapping(v, sink_mapping), + }, + }, + }; + }; + + return transform(required_copies, make_copy_to); +} + +std::unordered_map + resolve_tensor_mappings_from_node_mappings( + DynamicOpenDataflowGraph const &g) { + + auto get_mappings_for_invocation = [&](DynamicNodeInvocation const &i) + -> std::unordered_map { + return generate_map( + get_dynamic_slot_sites_for_invocation(i), + [&](InternalDynamicSlotSite const &s) -> ParallelTensorMapping { + return ParallelTensorMapping{ + dynamic_node_mapping_bindings_for_slot_name( + assert_unwrap(i.node_attrs.mapping), + s.slot_name.slot_name), + }; + }); + }; + + std::unordered_map result = + merge_disjoint_maps(transform(get_dynamic_invocation_set(g), + get_mappings_for_invocation)); return result; } @@ -154,25 +232,26 @@ DynamicOpenDataflowGraph ASSERT(no_part_of_graph_is_copy_inserted(g)); - std::unordered_map - unmapped_value_to_mapped_source_value; - for (DynamicNodeInvocation const &i : g.invocations) { - for (auto const &[slot, value] : i.outputs) { - unmapped_value_to_mapped_source_value.insert( - std::pair{value, - map_dynamic_value_attrs_for_task_group( - slot, value, assert_unwrap(i.node_attrs.mapping))}); - } - } + std::unordered_map + fully_resolved_tensor_mappings = + resolve_tensor_mappings_from_node_mappings(g); + + std::unordered_set all_copies = + flatmap(unordered_set_of(get_dynamic_values(g)), + [&](DynamicValueAttrs const &v) + -> std::unordered_set { + return copies_for_value(g, v, fully_resolved_tensor_mappings); + }); + + std::unordered_set mapped_invocations = transform( + get_dynamic_invocation_set(g), + [&](DynamicNodeInvocation const &i) -> DynamicNodeInvocation { + return apply_mappings_for_invocation(i, fully_resolved_tensor_mappings); + }); - // Use regular flatmap here to remove duplicates (we don't want to copy the - // same tensor to the same place multiple times) DynamicOpenDataflowGraph result = dynamic_open_dataflow_graph_from_invocation_set( - flatmap(g.invocations, [&](DynamicNodeInvocation const &i) { - return perform_copy_insertion_for_invocation( - i, unmapped_value_to_mapped_source_value); - })); + set_union(all_copies, mapped_invocations)); ASSERT(graph_is_fully_copy_inserted(result)); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_graph_edge.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_graph_edge.cc new file mode 100644 index 0000000000..788c010084 --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_graph_edge.cc @@ -0,0 +1,20 @@ +#include "task-spec/dynamic_graph/dynamic_graph_edge.h" + +namespace FlexFlow { + +DynamicGraphEdge + dynamic_graph_edge_from_slot_sites(DynamicSlotSite const &src, + InternalDynamicSlotSite const &dst) { + if (src.is_internal()) { + ASSERT(src.require_internal().direction == TensorDirection::OUTPUT); + } + ASSERT(dst.direction == TensorDirection::INCOMING); + + return DynamicGraphEdge{ + /*src=*/src, + /*dst_node=*/dst.invocation, + /*dst_slot=*/dst.slot_name, + }; +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_invocation.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_invocation.cc new file mode 100644 index 0000000000..93ad9269ce --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_invocation.cc @@ -0,0 +1,62 @@ +#include "task-spec/dynamic_graph/dynamic_node_invocation.h" +#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" +#include "task-spec/dynamic_graph/training_operation_attrs.h" +#include "utils/containers/are_disjoint.h" +#include "utils/containers/set_union.h" +#include "utils/containers/unordered_set_of.h" +#include "utils/optional.h" + +namespace FlexFlow { + +std::unordered_map + get_slot_map_for_direction(DynamicNodeInvocation const &invocation, + TensorDirection direction) { + switch (direction) { + case TensorDirection::INCOMING: + return invocation.inputs; + case TensorDirection::OUTPUT: + return invocation.outputs; + default: + PANIC("Unexpected direction {}", direction); + } +} + +TrainingOpType + dynamic_node_invocation_get_op_type(DynamicNodeInvocation const &i) { + TrainingOperationAttrs training_op_attrs = + assert_unwrap(i.node_attrs.op_attrs); + + return training_op_attrs_get_op_type(training_op_attrs); +} + +std::unordered_set + get_dynamic_slot_sites_for_invocation(DynamicNodeInvocation const &i) { + + std::unordered_set input_slots = + transform(unordered_set_of(i.inputs), + [&](std::pair const &p) + -> InternalDynamicSlotSite { + return InternalDynamicSlotSite{ + /*invocation=*/i, + /*direction=*/TensorDirection::INCOMING, + /*slot_name=*/p.first, + }; + }); + + std::unordered_set output_slots = + transform(unordered_set_of(i.outputs), + [&](std::pair const &p) + -> InternalDynamicSlotSite { + return InternalDynamicSlotSite{ + /*invocation=*/i, + /*direction=*/TensorDirection::OUTPUT, + /*slot_name=*/p.first, + }; + }); + + ASSERT(are_disjoint(input_slots, output_slots)); + + return set_union(input_slots, output_slots); +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc new file mode 100644 index 0000000000..9e56dddcf9 --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc @@ -0,0 +1,31 @@ +#include "task-spec/dynamic_graph/dynamic_node_mapping.h" +#include "utils/bidict/algorithms/transform_values.h" +#include "utils/containers/transform.h" + +namespace FlexFlow { + +bidict + dynamic_node_mapping_bindings_for_slot_name( + DynamicNodeMapping const &mapping, TensorSlotName const &slot_name) { + bidict coord_bindings = + get_tensor_bindings_for_slot_name(mapping.op_task_group, slot_name); + + return transform_values( + coord_bindings, [&](MachineSpaceCoordinate const &coord) -> device_id_t { + return device_id_t{coord, mapping.device_type}; + }); +} + +std::unordered_set + target_devices_of_dynamic_node_mapping(DynamicNodeMapping const &mapping) { + + return transform(mapping.op_task_group.get_shard_bindings().left_values(), + [&](MachineSpaceCoordinate const &c) -> device_id_t { + return device_id_t{ + /*coord=*/c, + /*device_type=*/mapping.device_type, + }; + }); +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc index d2a5b653e5..07c49664a5 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc @@ -1,19 +1,29 @@ #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" +#include "op-attrs/ff_ordered/filtrans.h" +#include "op-attrs/parallel_tensor_shape.h" +#include "op-attrs/pcg_operator_attrs.h" +#include "task-spec/dynamic_graph/dynamic_graph_edge.h" +#include "task-spec/dynamic_graph/dynamic_node_invocation.h" +#include "task-spec/dynamic_graph/dynamic_slot_site.h" #include "task-spec/dynamic_graph/serializable_dynamic_node_attrs.h" #include "task-spec/dynamic_graph/serializable_dynamic_value_attrs.h" #include "utils/containers/all_of.h" #include "utils/containers/concat_vectors.h" #include "utils/containers/contains_duplicates.h" +#include "utils/containers/contains_value.h" +#include "utils/containers/filter_values.h" #include "utils/containers/flatmap.h" +#include "utils/containers/get_only.h" #include "utils/containers/multiset_union.h" +#include "utils/containers/transform.h" #include "utils/containers/zip_strict.h" #include "utils/containers/zip_values_strict.h" #include "utils/graph/dataflow_graph/algorithms.h" #include "utils/graph/instances/unordered_set_labelled_open_dataflow_graph.h" #include "utils/graph/instances/unordered_set_labelled_open_kwarg_dataflow_graph.h" +#include "utils/graph/labelled_kwarg_dataflow_graph/algorithms/labelled_kwarg_dataflow_graph_view_as_dot.h" #include "utils/graph/labelled_open_dataflow_graph/algorithms/find_isomorphism.h" #include "utils/graph/labelled_open_kwarg_dataflow_graph/algorithms/find_isomorphism_between_labelled_open_kwarg_dataflow_graphs.h" -#include "utils/graph/labelled_open_kwarg_dataflow_graph/algorithms/labelled_open_kwarg_dataflow_graph_view_as_dot.h" #include "utils/graph/labelled_open_kwarg_dataflow_graph/labelled_open_kwarg_dataflow_graph.h" #include "utils/graph/node/algorithms.h" #include "utils/graph/open_dataflow_graph/algorithms/get_inputs.h" @@ -28,6 +38,30 @@ DynamicOpenDataflowGraph make_empty_dynamic_open_dataflow_graph() { }; } +void check_dynamic_open_dataflow_graph_is_valid( + DynamicOpenDataflowGraph const &g) { + std::unordered_map> + invocations_by_value_produced; + + for (DynamicNodeInvocation const &i : g.invocations) { + for (DynamicValueAttrs const &v : values(i.outputs)) { + invocations_by_value_produced[v].push_back(i); + } + } + + std::unordered_map> + values_produced_multiple_times = filter_values( + invocations_by_value_produced, + [](std::vector const &producers) -> bool { + return producers.size() > 1; + }); + + ASSERT(values_produced_multiple_times.size() == 0, + keys(values_produced_multiple_times)); + + labelled_open_kwarg_dataflow_graph_from_dynamic_open_dataflow_graph(g); +} + nonnegative_int dynamic_graph_num_nodes(DynamicOpenDataflowGraph const &g) { return num_elements(get_dynamic_nodes(g)); } @@ -88,6 +122,140 @@ std::unordered_set return g.invocations; } +std::unordered_set + get_dynamic_graph_edges(DynamicOpenDataflowGraph const &g) { + return flatmap(get_dynamic_invocation_set(g), + [&](DynamicNodeInvocation const &i) + -> std::unordered_set { + return get_dynamic_graph_edges_incoming_to_invocation(g, i); + }); +} + +std::unordered_set + get_dynamic_graph_edges_incoming_to_invocation( + DynamicOpenDataflowGraph const &g, DynamicNodeInvocation const &i) { + return transform(unordered_set_of(i.inputs), + [&](std::pair const &p) + -> DynamicGraphEdge { + DynamicSlotSite src = + dynamic_graph_find_source_of_value(g, p.second); + + InternalDynamicSlotSite dst = InternalDynamicSlotSite{ + /*invocation=*/i, + /*direction=*/TensorDirection::INCOMING, + /*slot_name=*/p.first, + }; + + return dynamic_graph_edge_from_slot_sites(src, dst); + }); +} + +std::unordered_set + get_dynamic_graph_edges_outgoing_from_invocation( + DynamicOpenDataflowGraph const &g, DynamicNodeInvocation const &i) { + return flatmap( + unordered_set_of(i.outputs), + [&](std::pair const &p) + -> std::unordered_set { + DynamicSlotSite src = DynamicSlotSite{ + InternalDynamicSlotSite{ + /*invocation=*/i, + /*direction=*/TensorDirection::OUTPUT, + /*slot_name=*/p.first, + }, + }; + + return transform( + dynamic_graph_find_sinks_of_value(g, p.second), + [&](InternalDynamicSlotSite const &sink) -> DynamicGraphEdge { + return dynamic_graph_edge_from_slot_sites(src, sink); + }); + }); +} + +std::unordered_set + get_internal_dynamic_slot_sites(DynamicOpenDataflowGraph const &g) { + return flatmap(get_dynamic_invocation_set(g), + [](DynamicNodeInvocation const &i) + -> std::unordered_set { + return get_dynamic_slot_sites_for_invocation(i); + }); +} + +std::unordered_set + get_dynamic_slot_sites(DynamicOpenDataflowGraph const &g) { + std::unordered_set internal_slot_sites = + get_internal_dynamic_slot_sites(g); + + std::unordered_set internal_values = + filtrans(internal_slot_sites, + [&](InternalDynamicSlotSite const &s) + -> std::optional { + if (s.direction == TensorDirection::OUTPUT) { + return dynamic_value_attrs_for_slot_site(DynamicSlotSite{s}); + } else { + return std::nullopt; + } + }); + + std::unordered_set all_values = + unordered_set_of(get_dynamic_values(g)); + + std::unordered_set external_values = + set_minus(all_values, internal_values); + + std::unordered_set external_slot_sites = transform( + external_values, + [](DynamicValueAttrs const &external_value) -> ExternalDynamicSlotSite { + return ExternalDynamicSlotSite{external_value}; + }); + + return set_union( + transform(internal_slot_sites, + [](InternalDynamicSlotSite const &s) -> DynamicSlotSite { + return DynamicSlotSite{s}; + }), + transform(external_slot_sites, + [](ExternalDynamicSlotSite const &s) -> DynamicSlotSite { + return DynamicSlotSite{s}; + })); +} + +std::unordered_set + dynamic_graph_find_sinks_of_value(DynamicOpenDataflowGraph const &g, + DynamicValueAttrs const &v) { + std::unordered_set found = filter( + get_internal_dynamic_slot_sites(g), + [&](InternalDynamicSlotSite const &s) -> bool { + return dynamic_value_attrs_for_slot_site(DynamicSlotSite{s}) == v && + s.direction == TensorDirection::INCOMING; + }); + + return found; +} + +DynamicSlotSite + dynamic_graph_find_source_of_value(DynamicOpenDataflowGraph const &g, + DynamicValueAttrs const &v) { + + auto is_source_of_value = [&](DynamicSlotSite const &s) -> bool { + return s.visit(overload{ + [&](InternalDynamicSlotSite const &internal_slot_site) -> bool { + return dynamic_value_attrs_for_slot_site(s) == v && + internal_slot_site.direction == TensorDirection::OUTPUT; + }, + [&](ExternalDynamicSlotSite const &external_slot_site) -> bool { + return external_slot_site.value == v; + }, + }); + }; + + std::unordered_set found = + filter(get_dynamic_slot_sites(g), is_source_of_value); + + return get_only(found); +} + std::optional find_output_value_attrs(DynamicOpenDataflowGraph const &dg, dynamic_tensor_guid_t tensor_guid, @@ -133,9 +301,13 @@ DynamicOpenDataflowGraph flatmap_dynamic_invocation_set( DynamicOpenDataflowGraph dynamic_open_dataflow_graph_from_invocation_set( std::unordered_set const &invocation_set) { - return DynamicOpenDataflowGraph{ + DynamicOpenDataflowGraph result = DynamicOpenDataflowGraph{ invocation_set, }; + + check_dynamic_open_dataflow_graph_is_valid(result); + + return result; } std::pair RecordFormatter { + return mk_record_for_map(mapping.raw.as_unordered_map()); + }; + std::function render_value_label = [&](DynamicValueAttrs const &a) -> nlohmann::json { nlohmann::json result = dynamic_value_attrs_to_serializable(a); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_slot_site.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_slot_site.cc new file mode 100644 index 0000000000..af13345abd --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_slot_site.cc @@ -0,0 +1,26 @@ +#include "task-spec/dynamic_graph/dynamic_slot_site.h" +#include "utils/overload.h" + +namespace FlexFlow { + +DynamicValueAttrs + dynamic_value_attrs_for_slot_site(DynamicSlotSite const &slot) { + return slot.visit(overload{ + + [](ExternalDynamicSlotSite const &external_slot) -> DynamicValueAttrs { + return external_slot.value; + }, + + [](InternalDynamicSlotSite const &internal_slot) -> DynamicValueAttrs { + switch (internal_slot.direction) { + case TensorDirection::INCOMING: + return internal_slot.invocation.inputs.at(internal_slot.slot_name); + case TensorDirection::OUTPUT: + return internal_slot.invocation.outputs.at(internal_slot.slot_name); + default: + PANIC("Unexpected direction {}", internal_slot.direction); + } + }}); +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc index 282279edbe..a05330b1b8 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc @@ -13,4 +13,13 @@ DynamicValueAttrs return result; } +DynamicValueAttrs + dynamic_value_attrs_with_mapping(DynamicValueAttrs const &v, + ParallelTensorMapping const &m) { + ASSERT(v.mapping == std::nullopt); + DynamicValueAttrs result = v; + result.mapping = m; + return result; +} + } // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc index 8066926262..9bcdc351fc 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc @@ -17,7 +17,7 @@ LossInsertionResult perform_loss_insertion( DynamicOpenDataflowGraph const &dg, LossAttrs const &loss_attrs, dynamic_tensor_guid_t logit_tensor, - std::optional const &loss_mapping) { + std::optional const &loss_mapping) { DynamicValueAttrs logit_value = assert_unwrap( find_output_value_attrs(dg, logit_tensor, mk_dynamic_tensor_role_fwd())); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc b/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc index 0a22015ddf..e46f8b1510 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc @@ -6,7 +6,7 @@ namespace FlexFlow { std::unordered_set perform_machine_slicing_for_invocation( DynamicNodeInvocation const &invocation, - MachineSpaceCoordinate const &device_coord) { + device_id_t const &device_coord) { ASSERT(invocation.node_attrs.device_coord.has_value()); @@ -19,7 +19,7 @@ std::unordered_set DynamicOpenDataflowGraph perform_machine_slicing(DynamicOpenDataflowGraph const &g, - MachineSpaceCoordinate const &device_coord) { + device_id_t const &device_coord) { DynamicOpenDataflowGraph result = flatmap_dynamic_invocation_set( g, [&](DynamicNodeInvocation const &invocation) diff --git a/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc b/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc index 380c2d17a1..57efdd150b 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc @@ -15,7 +15,7 @@ namespace FlexFlow { DynamicOpenDataflowGraph make_dynamic_open_dataflow_graph_from_mapped_pcg( - MappedParallelComputationGraph const &mpcg) { + MappedParallelComputationGraph const &mpcg, DeviceType device_type) { DynamicOpenDataflowGraph result = make_empty_dynamic_open_dataflow_graph(); ParallelComputationGraph pcg = pcg_from_mpcg(mpcg); @@ -24,7 +24,11 @@ DynamicOpenDataflowGraph make_dynamic_open_dataflow_graph_from_mapped_pcg( DynamicNodeAttrs result_attrs{ /*task_type=*/std::nullopt, /*device_coord=*/std::nullopt, - /*mapping=*/mpcg_get_mapping_for_layer(mpcg, layer), + /*mapping=*/ + DynamicNodeMapping{ + /*op_task_group=*/mpcg_get_mapping_for_layer(mpcg, layer), + /*device_type=*/device_type, + }, /*op_attrs=*/TrainingOperationAttrs{attrs.op_attrs}, /*pcg_layer_guid=*/dynamic_layer_guid_t{layer}, /*per_device_op_state=*/std::nullopt, 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 fb6efb96d0..117819a639 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 @@ -1,6 +1,8 @@ #include "task-spec/dynamic_graph/shard_expansion.h" +#include "task-spec/dynamic_graph/dynamic_node_mapping.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" +#include "task-spec/dynamic_graph/shard_expansion.h" #include "utils/bidict/algorithms/filter_keys.h" #include "utils/containers/get_only.h" #include "utils/containers/map_values2.h" @@ -40,10 +42,9 @@ bool graph_is_fully_shard_expanded(DynamicOpenDataflowGraph const &g) { slot_is_shard_expanded); } -static bidict +static bidict restrict_tensor_mapping_keys_to_coord( - bidict const - &mapping, + bidict const &mapping, ParallelTensorSpaceCoordinate const ¶llel_tensor_coord) { return filter_keys(mapping, [&](ParallelTensorSpaceCoordinate const &p) { return p == parallel_tensor_coord; @@ -52,7 +53,7 @@ static bidict static DynamicNodeInvocation shard_invocation_for_binding( DynamicNodeInvocation const &i, - MachineSpaceCoordinate const &machine_coord, + device_id_t const &device_coord, OperatorAtomicTaskShardBinding const &binding) { auto shard_expand_value_attrs = [&](DynamicTensorSlot const &s, @@ -64,17 +65,18 @@ static DynamicNodeInvocation shard_invocation_for_binding( result.shard_coord = parallel_tensor_coord; result.mapping = transform( v.mapping, - [&](bidict const - &mapping) { - return restrict_tensor_mapping_keys_to_coord(mapping, - parallel_tensor_coord); + [&](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_coord = machine_coord; + result.device_coord = device_coord; return result; }(); @@ -89,10 +91,10 @@ static std::unordered_set perform_shard_expansion_for_copy(DynamicNodeInvocation const &i) { auto [input_slot, input] = get_only(i.inputs); auto [output_slot, output] = get_only(i.outputs); - bidict input_mapping = - assert_unwrap(input.mapping); + bidict input_mapping = + assert_unwrap(input.mapping).raw; require_same(input_mapping.left_values(), - assert_unwrap(output.mapping).left_values()); + assert_unwrap(output.mapping).raw.left_values()); return transform( input_mapping.left_values(), [&](ParallelTensorSpaceCoordinate const &p) { @@ -103,7 +105,7 @@ static std::unordered_set // because we expect this to align with the most efficient way to issue // copies in Realm, although the current Realm backend uses a // centralized controller and thus issues copies all from a single node. - MachineSpaceCoordinate machine_coord = input_mapping.at_l(p); + device_id_t machine_coord = input_mapping.at_l(p); return shard_invocation_for_binding(i, machine_coord, @@ -121,16 +123,15 @@ std::unordered_set return perform_shard_expansion_for_copy(i); } - MappedOperatorTaskGroup mapping = assert_unwrap(i.node_attrs.mapping); + DynamicNodeMapping mapping = assert_unwrap(i.node_attrs.mapping); - std::unordered_set shard_machine_coords = - mapping.get_shard_bindings().left_values(); + std::unordered_set shard_machine_coords = + target_devices_of_dynamic_node_mapping(mapping); return transform( - shard_machine_coords, - [&](MachineSpaceCoordinate const &c) -> DynamicNodeInvocation { + shard_machine_coords, [&](device_id_t const &c) -> DynamicNodeInvocation { OperatorAtomicTaskShardBinding slot_bindings = - mapping.get_shard_bindings().at_l(c); + mapping.op_task_group.get_shard_bindings().at_l(c.coord); return shard_invocation_for_binding(i, c, slot_bindings); }); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc new file mode 100644 index 0000000000..90db73e617 --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc @@ -0,0 +1,24 @@ +#include "task-spec/dynamic_graph/training_operation_attrs.h" +#include "op-attrs/pcg_operator_attrs.h" +#include "utils/overload.h" + +namespace FlexFlow { + +TrainingOpType training_op_attrs_get_op_type( + TrainingOperationAttrs const &training_op_attrs) { + return training_op_attrs.visit(overload{ + [](PCGOperatorAttrs const &a) -> TrainingOpType { + return TrainingOpType{ + pcg_op_attrs_get_op_type(a), + }; + }, + [](LossAttrs const &) -> TrainingOpType { + return TrainingOpType{TrainingOnlyOpType::LOSS}; + }, + [](CopyAttrs const &) -> TrainingOpType { + return TrainingOpType{TrainingOnlyOpType::COPY}; + }, + }); +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/update_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/update_insertion.cc index 58a32db6c1..0c3769512b 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/update_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/update_insertion.cc @@ -58,8 +58,8 @@ static DynamicNodeInvocation get_update_invocation_for_invocation( auto create_binding_for_role = [&](DynamicTensorRole const &role) -> std::pair { DynamicTensorSlot binding_slot = tensor_slot_with_role(slot, role); - DynamicValueAttrs binding_attrs = dynamic_value_attrs_with_role( - value_attrs, mk_dynamic_tensor_role_fwd()); + DynamicValueAttrs binding_attrs = + dynamic_value_attrs_with_role(value_attrs, role); return std::pair{ binding_slot, diff --git a/lib/task-spec/src/task-spec/serialization.cc b/lib/task-spec/src/task-spec/serialization.cc deleted file mode 100644 index a2ad6eabfa..0000000000 --- a/lib/task-spec/src/task-spec/serialization.cc +++ /dev/null @@ -1 +0,0 @@ -#include "task-spec/serialization.h" diff --git a/lib/task-spec/test/CMakeLists.txt b/lib/task-spec/test/CMakeLists.txt index 9665dba88e..d417280c26 100644 --- a/lib/task-spec/test/CMakeLists.txt +++ b/lib/task-spec/test/CMakeLists.txt @@ -8,7 +8,6 @@ ff_add_test_executable( DEPS doctest utils-test-common - # local-execution kernels task-spec op-attrs diff --git a/lib/task-spec/test/src/task-spec/device_specific.cc b/lib/task-spec/test/src/task-spec/device_specific.cc index 34ef9b2bef..6a42a9b570 100644 --- a/lib/task-spec/test/src/task-spec/device_specific.cc +++ b/lib/task-spec/test/src/task-spec/device_specific.cc @@ -6,12 +6,14 @@ using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("DeviceSpecific") { DeviceSpecific device_specific1 = - DeviceSpecific::create(device_id_t{gpu_id_t{0_n}}, - "hello world"); + DeviceSpecific::create( + device_id_t{MachineSpaceCoordinate{0_n, 0_n}, DeviceType::GPU}, + "hello world"); DeviceSpecific device_specific2 = - DeviceSpecific::create(device_id_t{gpu_id_t{1_n}}, - "hello world"); + DeviceSpecific::create( + device_id_t{MachineSpaceCoordinate{0_n, 1_n}, DeviceType::GPU}, + "hello world"); std::string result1 = fmt::to_string(device_specific1); std::string result2 = fmt::to_string(device_specific2); 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 2160f6bf82..58a8a36fcd 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,433 +1,478 @@ #include "task-spec/dynamic_graph/copy_insertion.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" #include "task-spec/dynamic_graph/dynamic_task_type.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_role.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" +#include "test/utils/doctest/check_kv.h" #include "test/utils/doctest/fmt/unordered_set.h" #include using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { - TEST_CASE("perform_copy_insertion_for_invocation") { + TEST_CASE("perform_copy_insertion") { + + auto mk_slot = [](TensorSlotName tensor_slot_name) -> DynamicTensorSlot { + return DynamicTensorSlot{ + tensor_slot_name, + std::nullopt, + }; + }; + + auto mk_value_attrs = + [](size_t src_layer_guid, + TensorSlotName src_slot, + std::optional const &mapping) + -> DynamicValueAttrs { + return DynamicValueAttrs{ + /*tensor_guid=*/dynamic_tensor_guid_t{ + parallel_tensor_guid_t{ + KwargDataflowOutput{ + Node{ + src_layer_guid, + }, + src_slot, + }, + }, + }, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/mapping, + /*accessor=*/std::nullopt, + /*role=*/std::nullopt, + }; + }; + + auto mk_ptensor_coord = + [](nonnegative_int shard_idx) -> ParallelTensorSpaceCoordinate { + return ParallelTensorSpaceCoordinate{ + /*sum_component=*/0_n, + /*discard_copy_component=*/0_n, + /*shard_components=*/ + FFOrdered{ + shard_idx, + }, + }; + }; + auto mk_machine_coord = - [](nonnegative_int node_idx, - nonnegative_int device_idx) -> MachineSpaceCoordinate { + [](nonnegative_int device_idx) -> MachineSpaceCoordinate { return MachineSpaceCoordinate{ - /*node_idx=*/node_idx, + /*node_idx=*/0_n, /*device_idx=*/device_idx, - /*device_type=*/DeviceType::GPU, }; }; - auto mk_pt_coord = - [](nonnegative_int idx1, - nonnegative_int idx2, - nonnegative_int idx3, - nonnegative_int idx4) -> ParallelTensorSpaceCoordinate { - return ParallelTensorSpaceCoordinate{ - /*sum_component=*/idx1, - /*discard_copy_component=*/idx2, - /*shard_components=*/ - FFOrdered{ - idx3, - idx4, + auto mk_device_id = [&](nonnegative_int device_idx) -> device_id_t { + return device_id_t{ + mk_machine_coord(device_idx), + DeviceType::GPU, + }; + }; + + auto mk_pcg_layer_guid = [](size_t pcg_layer_guid) -> dynamic_layer_guid_t { + return dynamic_layer_guid_t{ + parallel_layer_guid_t{ + Node{pcg_layer_guid}, }, }; }; - auto mk_input_shard_binding = [&](ParallelTensorSpaceCoordinate const &c) + auto mk_node_attrs = + [](dynamic_layer_guid_t layer_guid, + std::optional const &mapping, + std::optional const &op_attrs) + -> DynamicNodeAttrs { + return DynamicNodeAttrs{ + /*task_type=*/std::nullopt, + /*device_coord=*/std::nullopt, + /*mapping=*/mapping, + /*op_attrs=*/op_attrs, + /*layer_guid=*/layer_guid, + /*per_device_op_state=*/std::nullopt, + }; + }; + + auto mk_binding = [&](nonnegative_int input_shard_idx, + nonnegative_int output_shard_idx) -> OperatorAtomicTaskShardBinding { return OperatorAtomicTaskShardBinding{ /*tensor_coords=*/{ + { + TensorSlotName::INPUT, + mk_ptensor_coord(input_shard_idx), + }, { TensorSlotName::OUTPUT, - c, + mk_ptensor_coord(output_shard_idx), }, }, }; }; - auto mk_shard_binding = [&](ParallelTensorSpaceCoordinate const &c1, - ParallelTensorSpaceCoordinate const &c2, - ParallelTensorSpaceCoordinate const &c3, - ParallelTensorSpaceCoordinate const &c4) - -> OperatorAtomicTaskShardBinding { - return OperatorAtomicTaskShardBinding{ - /*tensor_coords=*/{ + DynamicValueAttrs v1 = mk_value_attrs( + /*src_layer_guid=*/0, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/std::nullopt); + + DynamicValueAttrs v2 = mk_value_attrs( + /*src_layer_guid=*/1, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/std::nullopt); + + DynamicValueAttrs v3 = mk_value_attrs( + /*src_layer_guid=*/2, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/std::nullopt); + + SUBCASE("inserts copy when necessary") { + DynamicNodeMapping mapping1 = DynamicNodeMapping{ + MappedOperatorTaskGroup{ + bidict{ + { + mk_machine_coord(0_n), + mk_binding(0_n, 0_n), + }, + { + mk_machine_coord(1_n), + mk_binding(1_n, 1_n), + }, + }, + }, + DeviceType::GPU, + }; + + DynamicNodeMapping mapping2 = DynamicNodeMapping{ + MappedOperatorTaskGroup{ + bidict{ + { + mk_machine_coord(0_n), + mk_binding(0_n, 0_n), + }, + { + mk_machine_coord(2_n), + mk_binding(1_n, 1_n), + }, + }, + }, + DeviceType::GPU, + }; + + DynamicNodeInvocation inv1 = DynamicNodeInvocation{ + /*inputs=*/{ { - TensorSlotName::INPUT, - c1, + mk_slot(TensorSlotName::INPUT), + v1, }, + }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(1), mapping1, /*op_attrs=*/std::nullopt), + /*outputs=*/ + { { - TensorSlotName::WEIGHT, - c2, + mk_slot(TensorSlotName::OUTPUT), + v2, }, + }, + }; + + DynamicNodeInvocation inv2 = DynamicNodeInvocation{ + /*inputs=*/{ { - TensorSlotName::OUTPUT_1, - c3, + mk_slot(TensorSlotName::INPUT), + v2, }, + }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(2), mapping2, /*op_attrs=*/std::nullopt), + /*outputs=*/ + { { - TensorSlotName::OUTPUT_2, - c4, + mk_slot(TensorSlotName::OUTPUT), + v3, }, }, }; - }; - MachineSpaceCoordinate mc1 = mk_machine_coord(0_n, 0_n); - MachineSpaceCoordinate mc2 = mk_machine_coord(1_n, 0_n); - MachineSpaceCoordinate mc3 = mk_machine_coord(2_n, 0_n); - MachineSpaceCoordinate mc4 = mk_machine_coord(3_n, 0_n); - - ParallelTensorSpaceCoordinate mc1_input_coord = - mk_pt_coord(0_n, 0_n, 0_n, 0_n); - ParallelTensorSpaceCoordinate mc1_weight_coord = - mk_pt_coord(0_n, 1_n, 2_n, 0_n); - ParallelTensorSpaceCoordinate mc1_output_1_coord = - mk_pt_coord(1_n, 0_n, 0_n, 1_n); - ParallelTensorSpaceCoordinate mc1_output_2_coord = - mk_pt_coord(3_n, 0_n, 0_n, 0_n); - - ParallelTensorSpaceCoordinate mc2_input_coord = - mk_pt_coord(0_n, 1_n, 0_n, 0_n); - ParallelTensorSpaceCoordinate mc2_weight_coord = - mk_pt_coord(0_n, 4_n, 2_n, 0_n); - ParallelTensorSpaceCoordinate mc2_output_1_coord = - mk_pt_coord(1_n, 2_n, 0_n, 1_n); - ParallelTensorSpaceCoordinate mc2_output_2_coord = - mk_pt_coord(0_n, 0_n, 0_n, 0_n); - - MappedOperatorTaskGroup input_mapping_same = MappedOperatorTaskGroup{ - bidict{ - { - mc1, - mk_input_shard_binding(mc1_input_coord), - }, - { - mc2, - mk_input_shard_binding(mc2_input_coord), - }, - }, - }; + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set({inv1, inv2}); - MappedOperatorTaskGroup weight_mapping_same = MappedOperatorTaskGroup{ - bidict{ - { - mc1, - mk_input_shard_binding(mc1_weight_coord), - }, - { - mc2, - mk_input_shard_binding(mc2_weight_coord), - }, - }, - }; + DynamicOpenDataflowGraph result = perform_copy_insertion(g); - MappedOperatorTaskGroup invocation_mapping = MappedOperatorTaskGroup{ - bidict{ - { - mc1, - mk_shard_binding(mc1_input_coord, - mc1_weight_coord, - mc1_output_1_coord, - mc1_output_2_coord), + DynamicOpenDataflowGraph correct = [&] { + DynamicValueAttrs mapped_v1 = mk_value_attrs( + /*src_layer_guid=*/0, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(1_n)}, + }, + }); + + DynamicValueAttrs mapped_v2_placement1 = mk_value_attrs( + /*src_layer_guid=*/1, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(1_n)}, + }, + }); + + DynamicValueAttrs mapped_v2_placement2 = mk_value_attrs( + /*src_layer_guid=*/1, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(2_n)}, + }, + }); + + DynamicValueAttrs mapped_v3 = mk_value_attrs( + /*src_layer_guid=*/2, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(2_n)}, + }, + }); + + DynamicNodeInvocation mapped_inv1 = DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mapped_v1, + }, }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(1), mapping1, /*op_attrs=*/std::nullopt), + /*outputs=*/ { - mc2, - mk_shard_binding(mc2_input_coord, - mc2_weight_coord, - mc2_output_1_coord, - mc2_output_2_coord), - }, - }, - }; - - MappedOperatorTaskGroup invocation_mapping_diff_vs_copy1 = - MappedOperatorTaskGroup{ - bidict{ { - mc2, - mk_shard_binding(mc2_input_coord, - mc2_weight_coord, - mc2_output_1_coord, - mc2_output_2_coord), + mk_slot(TensorSlotName::OUTPUT), + mapped_v2_placement1, }, }, }; - auto mk_slot = [](TensorSlotName const &slot_name) -> DynamicTensorSlot { - return DynamicTensorSlot{ - /*slot_name=*/slot_name, - /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), - }; - }; - auto mk_value = [&](size_t src_node_id, - TensorSlotName src_slot_name, - MappedOperatorTaskGroup const &mapping, - std::optional const &use_slot_name) - -> DynamicValueAttrs { - return DynamicValueAttrs{ - /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ - KwargDataflowOutput{ - Node{src_node_id}, - src_slot_name, - }, - }}, - /*parallel_tensor_shape=*/std::nullopt, - /*shard_coord=*/std::nullopt, - /*mapping=*/ - transform(use_slot_name, - [&](TensorSlotName s) { - return get_tensor_bindings_for_slot_name(mapping, s); - }), - /*accessor=*/std::nullopt, - /*role=*/std::nullopt, - }; - }; - - size_t invocation1_id = 20; - - DynamicValueAttrs graph_input1 = - mk_value(0, TensorSlotName::OUTPUT, invocation_mapping, std::nullopt); - DynamicValueAttrs graph_input1_use = mk_value( - 0, TensorSlotName::OUTPUT, invocation_mapping, TensorSlotName::INPUT); - DynamicValueAttrs graph_input1_use_diff_vs_copy1 = - mk_value(0, - TensorSlotName::OUTPUT, - invocation_mapping_diff_vs_copy1, - TensorSlotName::INPUT); - DynamicValueAttrs graph_input2 = - mk_value(1, TensorSlotName::OUTPUT, invocation_mapping, std::nullopt); - DynamicValueAttrs graph_input2_use = mk_value( - 1, TensorSlotName::OUTPUT, invocation_mapping, TensorSlotName::WEIGHT); - DynamicValueAttrs invocation1_output1 = mk_value(invocation1_id, - TensorSlotName::OUTPUT_1, - invocation_mapping, - std::nullopt); - DynamicValueAttrs invocation1_output1_src = - mk_value(invocation1_id, - TensorSlotName::OUTPUT_1, - invocation_mapping, - TensorSlotName::OUTPUT_1); - DynamicValueAttrs invocation1_output2 = mk_value(invocation1_id, - TensorSlotName::OUTPUT_2, - invocation_mapping, - std::nullopt); - DynamicValueAttrs invocation1_output2_src = - mk_value(invocation1_id, - TensorSlotName::OUTPUT_2, - invocation_mapping, - TensorSlotName::OUTPUT_2); - - DynamicValueAttrs graph_input1_src_same = mk_value( - 0, TensorSlotName::OUTPUT, input_mapping_same, TensorSlotName::OUTPUT); - DynamicValueAttrs graph_input2_src_same = mk_value( - 1, TensorSlotName::OUTPUT, weight_mapping_same, TensorSlotName::OUTPUT); - - DynamicNodeInvocation input = DynamicNodeInvocation{ - /*inputs=*/{ - { - mk_slot(TensorSlotName::INPUT), - graph_input1, - }, - { - mk_slot(TensorSlotName::WEIGHT), - graph_input2, - }, - }, - /*node_attrs=*/ - DynamicNodeAttrs{ - /*task_type=*/DynamicTaskType::FWD, - /*device_coord=*/std::nullopt, - /*mapping=*/invocation_mapping, - /*op_attrs=*/std::nullopt, - /*layer_guid=*/ - dynamic_layer_guid_t{parallel_layer_guid_t{Node{20}}}, - /*per_device_op_state=*/std::nullopt, - }, - /*outputs=*/ - { - { - mk_slot(TensorSlotName::OUTPUT_1), - invocation1_output1, + DynamicNodeInvocation inserted_copy = DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mapped_v2_placement1, + }, }, + /*node_attrs=*/ + mk_node_attrs(dynamic_layer_guid_t{dynamic_copy_layer_guid_t{}}, + std::nullopt, + /*op_attrs=*/TrainingOperationAttrs{CopyAttrs{}}), + /*outputs=*/ { - mk_slot(TensorSlotName::OUTPUT_2), - invocation1_output2, + { + mk_slot(TensorSlotName::OUTPUT), + mapped_v2_placement2, + }, }, - }, - }; - DynamicNodeInvocation mapped = DynamicNodeInvocation{ - /*inputs=*/{ - { - mk_slot(TensorSlotName::INPUT), - graph_input1_use, - }, - { - mk_slot(TensorSlotName::WEIGHT), - graph_input2_use, - }, - }, - /*node_attrs=*/ - DynamicNodeAttrs{ - /*task_type=*/DynamicTaskType::FWD, - /*device_coord=*/std::nullopt, - /*mapping=*/invocation_mapping, - /*op_attrs=*/std::nullopt, - /*layer_guid=*/ - dynamic_layer_guid_t{parallel_layer_guid_t{Node{20}}}, - /*per_device_op_state=*/std::nullopt, - }, - /*outputs=*/ - { - { - mk_slot(TensorSlotName::OUTPUT_1), - invocation1_output1_src, + }; + + DynamicNodeInvocation mapped_inv2 = DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mapped_v2_placement2, + }, }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(2), mapping2, /*op_attrs=*/std::nullopt), + /*outputs=*/ { - mk_slot(TensorSlotName::OUTPUT_2), - invocation1_output2_src, + { + mk_slot(TensorSlotName::OUTPUT), + mapped_v3, + }, }, - }, - }; - - auto mk_copy = [&](DynamicValueAttrs const &src, - DynamicValueAttrs const &dst) { - return DynamicNodeInvocation{ - /*inputs=*/{{mk_slot(TensorSlotName::INPUT), src}}, - /*node_attrs=*/ - DynamicNodeAttrs{ - /*task_type=*/DynamicTaskType::FWD, - /*device_coord=*/std::nullopt, - /*mapping=*/std::nullopt, - /*op_attrs*/ TrainingOperationAttrs{CopyAttrs{}}, - /*layer_guid=*/dynamic_layer_guid_t{dynamic_copy_layer_guid_t{}}, - /*per_device_op_state=*/std::nullopt, - }, - /*outputs=*/{{mk_slot(TensorSlotName::OUTPUT), dst}}, - }; - }; - - SUBCASE("same mapping, no copies") { - std::unordered_map sources_same{ - {graph_input1, graph_input1_src_same}, - {graph_input2, graph_input2_src_same}}; - - std::unordered_set result = - perform_copy_insertion_for_invocation(input, sources_same); + }; - std::unordered_set correct = {mapped}; + return dynamic_open_dataflow_graph_from_invocation_set( + {mapped_inv1, mapped_inv2, inserted_copy}); + }(); - CHECK(result.size() == correct.size()); - CHECK(result == correct); + CHECK_MESSAGE( + result == correct, + check_kv("result\n", dynamic_open_dataflow_graph_as_dot(result)), + check_kv("correct\n", dynamic_open_dataflow_graph_as_dot(correct))); } - SUBCASE("copy one tensor, one point") { - MappedOperatorTaskGroup input_mapping_copy1 = MappedOperatorTaskGroup{ - bidict{ - { - mc1, - mk_input_shard_binding(mc1_input_coord), - }, - { - mc3, - mk_input_shard_binding(mc2_input_coord), + SUBCASE("does not insert a copy when not necessary") { + DynamicNodeMapping mapping1 = DynamicNodeMapping{ + MappedOperatorTaskGroup{ + bidict{ + { + mk_machine_coord(0_n), + mk_binding(0_n, 0_n), + }, + { + mk_machine_coord(1_n), + mk_binding(1_n, 1_n), + }, }, }, + DeviceType::GPU, }; - MappedOperatorTaskGroup input_mapping_copy1_diff_vs_use = + + DynamicNodeMapping mapping2 = DynamicNodeMapping{ MappedOperatorTaskGroup{ bidict{ { - mc3, - mk_input_shard_binding(mc2_input_coord), + mk_machine_coord(0_n), + mk_binding(0_n, 0_n), + }, + { + mk_machine_coord(1_n), + mk_binding(1_n, 1_n), }, }, - }; - - DynamicValueAttrs graph_input1_src_copy1 = - mk_value(0, - TensorSlotName::OUTPUT, - input_mapping_copy1, - TensorSlotName::OUTPUT); - DynamicValueAttrs graph_input1_src_copy1_diff_vs_use = - mk_value(0, - TensorSlotName::OUTPUT, - input_mapping_copy1_diff_vs_use, - TensorSlotName::OUTPUT); - - std::unordered_map sources_copy1{ - {graph_input1, graph_input1_src_copy1}, - {graph_input2, graph_input2_src_same}}; - - std::unordered_set result = - perform_copy_insertion_for_invocation(input, sources_copy1); - - std::unordered_set correct = { - mapped, - mk_copy(graph_input1_src_copy1_diff_vs_use, - graph_input1_use_diff_vs_copy1), + }, + DeviceType::GPU, }; - CHECK(result.size() == correct.size()); - CHECK(result == correct); - } - - SUBCASE("copy two tensors, two points") { - MappedOperatorTaskGroup input_mapping_copy2 = MappedOperatorTaskGroup{ - bidict{ + DynamicNodeInvocation inv1 = DynamicNodeInvocation{ + /*inputs=*/{ { - mc3, - mk_input_shard_binding(mc1_input_coord), + mk_slot(TensorSlotName::INPUT), + v1, }, + }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(1), mapping1, /*op_attrs=*/std::nullopt), + /*outputs=*/ + { { - mc4, - mk_input_shard_binding(mc2_input_coord), + mk_slot(TensorSlotName::OUTPUT), + v2, }, }, }; - MappedOperatorTaskGroup weight_mapping_copy2 = MappedOperatorTaskGroup{ - bidict{ + + DynamicNodeInvocation inv2 = DynamicNodeInvocation{ + /*inputs=*/{ { - mc4, - mk_input_shard_binding(mc1_weight_coord), + mk_slot(TensorSlotName::INPUT), + v2, }, + }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(2), mapping2, /*op_attrs=*/std::nullopt), + /*outputs=*/ + { { - mc3, - mk_input_shard_binding(mc2_weight_coord), + mk_slot(TensorSlotName::OUTPUT), + v3, }, }, }; - DynamicValueAttrs graph_input1_src_copy2 = - mk_value(0, - TensorSlotName::OUTPUT, - input_mapping_copy2, - TensorSlotName::OUTPUT); - DynamicValueAttrs graph_input2_src_copy2 = - mk_value(1, - TensorSlotName::OUTPUT, - weight_mapping_copy2, - TensorSlotName::OUTPUT); - - std::unordered_map sources_copy2{ - {graph_input1, graph_input1_src_copy2}, - {graph_input2, graph_input2_src_copy2}}; - - std::unordered_set result = - perform_copy_insertion_for_invocation(input, sources_copy2); - - std::unordered_set correct = { - mapped, - mk_copy(graph_input1_src_copy2, graph_input1_use), - mk_copy(graph_input2_src_copy2, graph_input2_use), - }; + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set({inv1, inv2}); + + DynamicOpenDataflowGraph result = perform_copy_insertion(g); + + DynamicOpenDataflowGraph correct = [&] { + DynamicValueAttrs mapped_v1 = mk_value_attrs( + /*src_layer_guid=*/0, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(1_n)}, + }, + }); + + DynamicValueAttrs mapped_v2 = mk_value_attrs( + /*src_layer_guid=*/1, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(1_n)}, + }, + }); + + DynamicValueAttrs mapped_v3 = mk_value_attrs( + /*src_layer_guid=*/2, + /*src_slot=*/TensorSlotName::OUTPUT, + /*mapping=*/ + ParallelTensorMapping{ + bidict{ + {mk_ptensor_coord(0_n), mk_device_id(0_n)}, + {mk_ptensor_coord(1_n), mk_device_id(1_n)}, + }, + }); + + DynamicNodeInvocation mapped_inv1 = DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mapped_v1, + }, + }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(1), mapping1, /*op_attrs=*/std::nullopt), + /*outputs=*/ + { + { + mk_slot(TensorSlotName::OUTPUT), + mapped_v2, + }, + }, + }; + + DynamicNodeInvocation mapped_inv2 = DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mapped_v2, + }, + }, + /*node_attrs=*/ + mk_node_attrs( + mk_pcg_layer_guid(2), mapping2, /*op_attrs=*/std::nullopt), + /*outputs=*/ + { + { + mk_slot(TensorSlotName::OUTPUT), + mapped_v3, + }, + }, + }; + + return dynamic_open_dataflow_graph_from_invocation_set( + {mapped_inv1, mapped_inv2}); + }(); - CHECK(result.size() == correct.size()); - CHECK(result == correct); + CHECK_MESSAGE( + result == correct, + check_kv("result\n", dynamic_open_dataflow_graph_as_dot(result)), + check_kv("correct\n", dynamic_open_dataflow_graph_as_dot(correct))); } } } diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc index 49b8d4a77a..9e529bba04 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc @@ -1,4 +1,8 @@ #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" +#include "op-attrs/initializer_attrs.h" +#include "task-spec/dynamic_graph/dynamic_tensor_role.h" +#include "task-spec/dynamic_graph/serializable_dynamic_value_attrs.h" +#include "utils/graph/instances/unordered_set_labelled_open_kwarg_dataflow_graph.h" #include "utils/graph/node/algorithms.h" #include @@ -6,48 +10,212 @@ using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("dynamic_op_dataflow_graph_from_invocation_set") { - DynamicValueAttrs value_1 = DynamicValueAttrs{ - /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ - KwargDataflowOutput{ - Node{1}, - TensorSlotName::OUTPUT, - }, - }}, - /*parallel_tensor_shape=*/std::nullopt, - /*shard_coord=*/std::nullopt, - /*mapping=*/std::nullopt, - /*accessor=*/std::nullopt, - /*tensor_type=*/std::nullopt, + + auto mk_dynamic_value = [](size_t node_id, + TensorSlotName slot_name) -> DynamicValueAttrs { + return DynamicValueAttrs{ + /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ + KwargDataflowOutput{ + Node{node_id}, + slot_name, + }, + }}, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*tensor_type=*/std::nullopt, + }; }; - DynamicValueAttrs value_2 = DynamicValueAttrs{ - /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ - KwargDataflowOutput{ - Node{2}, - TensorSlotName::OUTPUT, - }, - }}, - /*parallel_tensor_shape=*/std::nullopt, - /*shard_coord=*/std::nullopt, - /*mapping=*/std::nullopt, - /*accessor=*/std::nullopt, - /*tensor_type=*/std::nullopt, + auto mk_slot = [](TensorSlotName slot_name) { + return DynamicTensorSlot{ + /*slot_name=*/slot_name, + /*slot_tensor_role=*/std::nullopt, + }; }; - DynamicValueAttrs value_3 = DynamicValueAttrs{ - /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ - KwargDataflowOutput{ - Node{3}, - TensorSlotName::OUTPUT, - }, - }}, - /*parallel_tensor_shape=*/std::nullopt, - /*shard_coord=*/std::nullopt, + DynamicValueAttrs value_1 = mk_dynamic_value(1, TensorSlotName::OUTPUT); + DynamicValueAttrs value_2 = mk_dynamic_value(2, TensorSlotName::OUTPUT); + DynamicValueAttrs value_3 = mk_dynamic_value(3, TensorSlotName::OUTPUT); + + DynamicNodeAttrs node_attrs = DynamicNodeAttrs{ + /*task_type=*/std::nullopt, + /*device_coord=*/std::nullopt, /*mapping=*/std::nullopt, - /*accessor=*/std::nullopt, - /*tensor_type=*/std::nullopt, + /*op_attrs=*/std::nullopt, + /*layer_guid=*/dynamic_layer_guid_t{parallel_layer_guid_t{Node{4}}}, + /*per_device_op_state=*/std::nullopt, }; + SUBCASE("correct usage") { + DynamicNodeInvocation invocation_1 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{ + { + mk_slot(TensorSlotName::INPUT), + value_1, + }, + }, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{ + { + mk_slot(TensorSlotName::OUTPUT), + value_2, + }, + }, + }; + + DynamicNodeInvocation invocation_2 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{}, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{ + { + mk_slot(TensorSlotName::OUTPUT), + value_3, + }, + }, + }; + + DynamicNodeInvocation invocation_3 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{ + { + mk_slot(TensorSlotName::INPUT), + value_1, + }, + { + mk_slot(TensorSlotName::WEIGHT), + value_2, + }, + { + mk_slot(TensorSlotName::BIAS), + value_1, + }, + }, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{}, + }; + + std::unordered_set invocation_set = { + invocation_1, + invocation_2, + invocation_3, + }; + + DynamicOpenDataflowGraph result = + dynamic_open_dataflow_graph_from_invocation_set(invocation_set); + + CHECK(dynamic_graph_num_nodes(result) == 3); + } + + SUBCASE("throws if multiple invocations produce the same value") { + DynamicNodeInvocation invocation_1 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{ + { + mk_slot(TensorSlotName::INPUT), + value_1, + }, + }, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{ + { + mk_slot(TensorSlotName::OUTPUT), + value_2, + }, + }, + }; + + DynamicNodeInvocation invocation_2 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{}, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{ + { + mk_slot(TensorSlotName::OUTPUT), + value_2, + }, + }, + }; + + std::unordered_set invocation_set = { + invocation_1, + invocation_2, + }; + + CHECK_THROWS( + dynamic_open_dataflow_graph_from_invocation_set(invocation_set)); + } + + SUBCASE("throws if invocations contain/create cycle") { + DynamicNodeInvocation invocation_1 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{ + { + mk_slot(TensorSlotName::INPUT), + value_1, + }, + }, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{ + { + mk_slot(TensorSlotName::OUTPUT), + value_2, + }, + }, + }; + + DynamicNodeInvocation invocation_2 = DynamicNodeInvocation{ + /*inputs=*/std::unordered_map{ + { + mk_slot(TensorSlotName::INPUT), + value_2, + }, + }, + /*node_attrs=*/node_attrs, + /*outputs=*/ + std::unordered_map{ + { + mk_slot(TensorSlotName::OUTPUT), + value_1, + }, + }, + }; + + std::unordered_set invocation_set = { + invocation_1, + invocation_2, + }; + + CHECK_THROWS( + dynamic_open_dataflow_graph_from_invocation_set(invocation_set)); + } + } + + TEST_CASE("get_dynamic_slot_sites") { + auto mk_dynamic_value = [](int node_id, + TensorSlotName slot_name) -> DynamicValueAttrs { + return DynamicValueAttrs{ + /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ + KwargDataflowOutput{ + Node{static_cast(node_id)}, + slot_name, + }, + }}, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*tensor_type=*/std::nullopt, + }; + }; + + DynamicValueAttrs value_1 = mk_dynamic_value(1, TensorSlotName::OUTPUT); + DynamicValueAttrs value_2 = mk_dynamic_value(2, TensorSlotName::OUTPUT); + DynamicValueAttrs value_3 = mk_dynamic_value(3, TensorSlotName::OUTPUT); + DynamicNodeAttrs node_attrs = DynamicNodeAttrs{ /*task_type=*/std::nullopt, /*device_coord=*/std::nullopt, @@ -59,25 +227,14 @@ TEST_SUITE(FF_TEST_SUITE) { DynamicNodeInvocation invocation_1 = DynamicNodeInvocation{ /*inputs=*/std::unordered_map{ - {DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::INPUT, - /*slot_tensor_role=*/std::nullopt, - }, - value_1}, - }, - /*node_attrs=*/node_attrs, - /*outputs=*/ - std::unordered_map{ - {DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/std::nullopt, - }, - value_2}, + { + DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::INPUT, + /*slot_tensor_role=*/std::nullopt, + }, + value_1, + }, }, - }; - - DynamicNodeInvocation invocation_2 = DynamicNodeInvocation{ - /*inputs=*/std::unordered_map{}, /*node_attrs=*/node_attrs, /*outputs=*/ std::unordered_map{ @@ -86,48 +243,250 @@ TEST_SUITE(FF_TEST_SUITE) { /*slot_name=*/TensorSlotName::OUTPUT, /*slot_tensor_role=*/std::nullopt, }, - value_3, + value_2, }, }, }; - DynamicNodeInvocation invocation_3 = DynamicNodeInvocation{ + DynamicNodeInvocation invocation_2 = DynamicNodeInvocation{ /*inputs=*/std::unordered_map{ { DynamicTensorSlot{ /*slot_name=*/TensorSlotName::INPUT, /*slot_tensor_role=*/std::nullopt, }, - value_1, + value_2, }, { DynamicTensorSlot{ /*slot_name=*/TensorSlotName::WEIGHT, /*slot_tensor_role=*/std::nullopt, }, - value_2, + value_3, }, - { - DynamicTensorSlot{ - /*slot_name=*/TensorSlotName::BIAS, - /*slot_tensor_role=*/std::nullopt, - }, + }, + /*node_attrs=*/node_attrs, + /*outputs=*/{}, + }; + + DynamicOpenDataflowGraph g = + dynamic_open_dataflow_graph_from_invocation_set( + std::unordered_set{invocation_1, invocation_2}); + + std::unordered_set result = get_dynamic_slot_sites(g); + + auto mk_internal_slot_site = [](DynamicNodeInvocation const &invocation, + TensorDirection direction, + TensorSlotName slot_name) { + return DynamicSlotSite{ + InternalDynamicSlotSite{ + /*invocation=*/invocation, + /*direction=*/direction, + /*slot_name=*/ + DynamicTensorSlot{ + /*slot_name=*/slot_name, + /*slot_tensor_role=*/std::nullopt, + }, + }, + }; + }; + + std::unordered_set correct = { + DynamicSlotSite{ + ExternalDynamicSlotSite{ value_1, }, }, - /*node_attrs=*/node_attrs, - /*outputs=*/std::unordered_map{}, + DynamicSlotSite{ + ExternalDynamicSlotSite{ + value_3, + }, + }, + mk_internal_slot_site( + invocation_1, TensorDirection::INCOMING, TensorSlotName::INPUT), + mk_internal_slot_site( + invocation_1, TensorDirection::OUTPUT, TensorSlotName::OUTPUT), + mk_internal_slot_site( + invocation_2, TensorDirection::INCOMING, TensorSlotName::INPUT), + mk_internal_slot_site( + invocation_2, TensorDirection::INCOMING, TensorSlotName::WEIGHT), + }; + + CHECK(result == correct); + } + + TEST_CASE( + "labelled_open_kwarg_dataflow_graph_from_dynamic_open_dataflow_graph") { + dynamic_layer_guid_t layer_guid = dynamic_layer_guid_t{ + parallel_layer_guid_t{ + Node{0}, + }, + }; + + dynamic_tensor_guid_t tensor_guid = dynamic_tensor_guid_t{ + parallel_tensor_guid_t{ + KwargDataflowOutput{ + /*node=*/Node{0}, + /*slot_name=*/TensorSlotName::OUTPUT, + }, + }, + }; + + TrainingOperationAttrs weight_attrs = TrainingOperationAttrs{ + PCGOperatorAttrs{ + WeightAttrs{ + /*tensor_shape=*/TensorShape{ + /*dims=*/TensorDims{ + FFOrdered{ + 4_p, + 3_p, + }, + }, + /*data_type=*/DataType::FLOAT, + }, + /*initializer=*/make_zero_initializer(), + }, + }, + }; + + DynamicNodeAttrs fwd_weight_node_attrs = DynamicNodeAttrs{ + /*task_type=*/DynamicTaskType::FWD, + /*device_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*op_attrs=*/weight_attrs, + /*layer_guid=*/layer_guid, + /*per_device_op_state=*/std::nullopt, + }; + + DynamicTensorSlot fwd_weight_output_slot1 = DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), + }; + + DynamicValueAttrs fwd_weight_output_attrs1 = DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/mk_dynamic_tensor_role_fwd(), + }; + + DynamicNodeInvocation weight_invocation = DynamicNodeInvocation{ + /*inputs=*/{}, + /*node_attrs=*/fwd_weight_node_attrs, + /*outputs=*/ + std::unordered_map{ + { + fwd_weight_output_slot1, + fwd_weight_output_attrs1, + }, + }, + }; + + DynamicNodeAttrs upd_weight_node_attrs = DynamicNodeAttrs{ + /*task_type=*/DynamicTaskType::UPD, + /*device_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*op_attrs=*/weight_attrs, + /*layer_guid=*/layer_guid, + /*per_device_op_state=*/std::nullopt, + }; + + DynamicTensorSlot upd_weight_input_slot2 = DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/mk_dynamic_tensor_role_bwd(), + }; + + DynamicValueAttrs upd_weight_input_attrs2 = DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/mk_dynamic_tensor_role_bwd(), }; - std::unordered_set invocation_set = { - invocation_1, - invocation_2, - invocation_3, + DynamicTensorSlot upd_weight_input_slot3 = DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/ + mk_dynamic_tensor_role_opt(OptimizerSlotName::SGD_V), }; - DynamicOpenDataflowGraph result = - dynamic_open_dataflow_graph_from_invocation_set(invocation_set); + DynamicValueAttrs upd_weight_input_attrs3 = DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/mk_dynamic_tensor_role_opt(OptimizerSlotName::SGD_V), + }; - ASSERT(dynamic_graph_num_nodes(result) == 3); + DynamicOpenDataflowGraph input = + dynamic_open_dataflow_graph_from_invocation_set( + /*invocations=*/{ + weight_invocation, + DynamicNodeInvocation{/*inputs=*/{ + { + fwd_weight_output_slot1, + fwd_weight_output_attrs1, + }, + { + upd_weight_input_slot2, + upd_weight_input_attrs2, + }, + { + upd_weight_input_slot3, + upd_weight_input_attrs3, + }, + }, + /*node_attrs=*/upd_weight_node_attrs, + /*outputs=*/{}}, + }); + + std::pair, + bidict> + result = + labelled_open_kwarg_dataflow_graph_from_dynamic_open_dataflow_graph( + input); + + LabelledOpenKwargDataflowGraph + correct = LabelledOpenKwargDataflowGraph:: + create>(); + + KwargNodeAddedResult fwd_weight_added = correct.add_node( + /*node_label=*/fwd_weight_node_attrs, + /*inputs=*/{}, + /*output_labels=*/ + { + { + fwd_weight_output_slot1, + fwd_weight_output_attrs1, + }, + }); + + KwargNodeAddedResult upd_weight_added = correct.add_node( + /*node_label=*/fwd_weight_node_attrs, + /*inputs=*/{}, + /*output_labels=*/ + { + { + fwd_weight_output_slot1, + fwd_weight_output_attrs1, + }, + }); } } diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc index 40b3460ee5..6bc31999a6 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc @@ -6,13 +6,14 @@ using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("perform_machine_slicing_for_invocation") { - auto mk_machine_coord = - [](nonnegative_int node_idx, - nonnegative_int device_idx) -> MachineSpaceCoordinate { - return MachineSpaceCoordinate{ - /*node_idx=*/node_idx, - /*device_idx=*/device_idx, - /*device_type=*/DeviceType::GPU, + auto mk_device_id = [](nonnegative_int node_idx, + nonnegative_int device_idx) -> device_id_t { + return device_id_t{ + MachineSpaceCoordinate{ + /*node_idx=*/node_idx, + /*device_idx=*/device_idx, + }, + DeviceType::GPU, }; }; @@ -32,9 +33,9 @@ TEST_SUITE(FF_TEST_SUITE) { }; }; - MachineSpaceCoordinate mc1 = mk_machine_coord(0_n, 0_n); - MachineSpaceCoordinate mc2 = mk_machine_coord(2_n, 0_n); - MachineSpaceCoordinate mc3 = mk_machine_coord(4_n, 0_n); + device_id_t mc1 = mk_device_id(0_n, 0_n); + device_id_t mc2 = mk_device_id(2_n, 0_n); + device_id_t mc3 = mk_device_id(4_n, 0_n); ParallelTensorSpaceCoordinate mc1_input_coord = mk_pt_coord(0_n, 0_n, 0_n, 0_n); diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc @@ -0,0 +1 @@ + diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc index efe21146db..65870192d7 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc @@ -2,6 +2,7 @@ #include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h" #include "task-spec/dynamic_graph/copy_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_copy_layer_guid_t.dtg.h" +#include "task-spec/dynamic_graph/dynamic_node_mapping.h" #include "task-spec/dynamic_graph/training_operation_attrs.dtg.h" #include "test/utils/doctest/fmt/unordered_set.h" #include "utils/bidict/algorithms/filter_keys.h" @@ -14,7 +15,6 @@ static MachineSpaceCoordinate mk_machine_coord(nonnegative_int node_idx, return MachineSpaceCoordinate{ /*node_idx=*/node_idx, /*device_idx=*/device_idx, - /*device_type=*/DeviceType::GPU, }; }; @@ -33,200 +33,157 @@ static ParallelTensorSpaceCoordinate mk_pt_coord(nonnegative_int idx1, }; }; -DynamicTensorSlot mk_slot(TensorSlotName const &slot_name) { - return DynamicTensorSlot{ - /*slot_name=*/slot_name, - /*slot_tensor_role=*/std::nullopt, - }; -}; - -DynamicValueAttrs - mk_value(size_t src_node_id, - TensorSlotName src_slot_name, - bidict - tensor_binding, - std::optional const &shard_coord) { - if (shard_coord.has_value()) { - tensor_binding = filter_keys(tensor_binding, - [&](ParallelTensorSpaceCoordinate const &p) { - return p == shard_coord.value(); - }); - } - return DynamicValueAttrs{ - /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ - KwargDataflowOutput{ - Node{src_node_id}, - src_slot_name, - }, - }}, - /*parallel_tensor_shape=*/std::nullopt, - /*shard_coord=*/shard_coord, - /*mapping=*/ - tensor_binding, - /*accessor=*/std::nullopt, - /*role=*/std::nullopt, - }; -}; - TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("perform_shard_expansion_for_invocation") { - auto mk_shard_binding = [&](ParallelTensorSpaceCoordinate const &c1, - ParallelTensorSpaceCoordinate const &c2, - ParallelTensorSpaceCoordinate const &c3, - ParallelTensorSpaceCoordinate const &c4) - -> OperatorAtomicTaskShardBinding { - return OperatorAtomicTaskShardBinding{ - /*tensor_coords=*/{ - { - TensorSlotName::INPUT, - c1, - }, - { - TensorSlotName::WEIGHT, - c2, - }, - { - TensorSlotName::OUTPUT_1, - c3, - }, - { - TensorSlotName::OUTPUT_2, - c4, - }, - }, + + auto mk_slot = [](TensorSlotName const &slot_name) { + return DynamicTensorSlot{ + /*slot_name=*/slot_name, + /*slot_tensor_role=*/std::nullopt, }; }; - MachineSpaceCoordinate mc1 = mk_machine_coord(0_n, 0_n); - MachineSpaceCoordinate mc2 = mk_machine_coord(2_n, 0_n); - - ParallelTensorSpaceCoordinate mc1_input_coord = - mk_pt_coord(0_n, 0_n, 0_n, 0_n); - ParallelTensorSpaceCoordinate mc1_weight_coord = - mk_pt_coord(0_n, 1_n, 2_n, 0_n); - ParallelTensorSpaceCoordinate mc1_output_1_coord = - mk_pt_coord(1_n, 0_n, 0_n, 1_n); - ParallelTensorSpaceCoordinate mc1_output_2_coord = - mk_pt_coord(3_n, 0_n, 0_n, 0_n); - - ParallelTensorSpaceCoordinate mc2_input_coord = - mk_pt_coord(0_n, 1_n, 0_n, 0_n); - ParallelTensorSpaceCoordinate mc2_weight_coord = - mk_pt_coord(0_n, 4_n, 2_n, 0_n); - ParallelTensorSpaceCoordinate mc2_output_1_coord = - mk_pt_coord(1_n, 2_n, 0_n, 1_n); - ParallelTensorSpaceCoordinate mc2_output_2_coord = - mk_pt_coord(0_n, 0_n, 0_n, 0_n); - - MappedOperatorTaskGroup mapped_task_group = MappedOperatorTaskGroup{ - bidict{ - { - mc1, - mk_shard_binding(mc1_input_coord, - mc1_weight_coord, - mc1_output_1_coord, - mc1_output_2_coord), - }, - { - mc2, - mk_shard_binding(mc2_input_coord, - mc2_weight_coord, - mc2_output_1_coord, - mc2_output_2_coord), - }, - }, + DeviceType device_type = DeviceType::GPU; + auto mk_device_id = [&](MachineSpaceCoordinate const &c) -> device_id_t { + return device_id_t{c, device_type}; }; - auto mk_op_value = + auto mk_value = [&](size_t src_node_id, TensorSlotName src_slot_name, - TensorSlotName use_slot_name, + bidict tensor_binding, std::optional const &shard_coord) -> DynamicValueAttrs { - bidict - tensor_binding = get_tensor_bindings_for_slot_name(mapped_task_group, - use_slot_name); - return mk_value(src_node_id, src_slot_name, tensor_binding, shard_coord); + if (shard_coord.has_value()) { + tensor_binding = + filter_keys(tensor_binding, + [&](ParallelTensorSpaceCoordinate const &p) -> bool { + return p == shard_coord.value(); + }); + } + + return DynamicValueAttrs{ + /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{ + KwargDataflowOutput{ + Node{src_node_id}, + src_slot_name, + }, + }}, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/shard_coord, + /*mapping=*/ + ParallelTensorMapping{tensor_binding}, + /*accessor=*/std::nullopt, + /*role=*/std::nullopt, + }; }; - DynamicNodeInvocation input = DynamicNodeInvocation{ - /*inputs=*/{ - { - mk_slot(TensorSlotName::INPUT), - mk_op_value(0, - TensorSlotName::OUTPUT, - TensorSlotName::INPUT, - std::nullopt), - }, - { - mk_slot(TensorSlotName::WEIGHT), - mk_op_value(1, - TensorSlotName::OUTPUT, - TensorSlotName::WEIGHT, - std::nullopt), - }, - }, - /*node_attrs=*/ - DynamicNodeAttrs{ - /*task_type=*/std::nullopt, - /*device_coord=*/std::nullopt, - /*mapping=*/mapped_task_group, - /*op_attrs=*/std::nullopt, - /*layer_guid=*/ - dynamic_layer_guid_t{parallel_layer_guid_t{Node{20}}}, - /*per_device_op_state=*/std::nullopt, - }, - /*outputs=*/ - { - { - mk_slot(TensorSlotName::OUTPUT_1), - mk_op_value(20, - TensorSlotName::OUTPUT_1, - TensorSlotName::OUTPUT_1, - std::nullopt), + SUBCASE("standard operators") { + auto mk_shard_binding = [&](ParallelTensorSpaceCoordinate const &c1, + ParallelTensorSpaceCoordinate const &c2, + ParallelTensorSpaceCoordinate const &c3, + ParallelTensorSpaceCoordinate const &c4) + -> OperatorAtomicTaskShardBinding { + return OperatorAtomicTaskShardBinding{ + /*tensor_coords=*/{ + { + TensorSlotName::INPUT, + c1, + }, + { + TensorSlotName::WEIGHT, + c2, + }, + { + TensorSlotName::OUTPUT_1, + c3, + }, + { + TensorSlotName::OUTPUT_2, + c4, + }, }, - { - mk_slot(TensorSlotName::OUTPUT_2), - mk_op_value(20, - TensorSlotName::OUTPUT_2, - TensorSlotName::OUTPUT_2, - std::nullopt), - }, - }, - }; + }; + }; + + MachineSpaceCoordinate mc1 = mk_machine_coord(0_n, 0_n); + MachineSpaceCoordinate mc2 = mk_machine_coord(2_n, 0_n); + + ParallelTensorSpaceCoordinate mc1_input_coord = + mk_pt_coord(0_n, 0_n, 0_n, 0_n); + ParallelTensorSpaceCoordinate mc1_weight_coord = + mk_pt_coord(0_n, 1_n, 2_n, 0_n); + ParallelTensorSpaceCoordinate mc1_output_1_coord = + mk_pt_coord(1_n, 0_n, 0_n, 1_n); + ParallelTensorSpaceCoordinate mc1_output_2_coord = + mk_pt_coord(3_n, 0_n, 0_n, 0_n); + + ParallelTensorSpaceCoordinate mc2_input_coord = + mk_pt_coord(0_n, 1_n, 0_n, 0_n); + ParallelTensorSpaceCoordinate mc2_weight_coord = + mk_pt_coord(0_n, 4_n, 2_n, 0_n); + ParallelTensorSpaceCoordinate mc2_output_1_coord = + mk_pt_coord(1_n, 2_n, 0_n, 1_n); + ParallelTensorSpaceCoordinate mc2_output_2_coord = + mk_pt_coord(0_n, 0_n, 0_n, 0_n); + + DynamicNodeMapping node_mapping = DynamicNodeMapping{ + MappedOperatorTaskGroup{ + bidict{ + { + mc1, + mk_shard_binding(mc1_input_coord, + mc1_weight_coord, + mc1_output_1_coord, + mc1_output_2_coord), + }, + { + mc2, + mk_shard_binding(mc2_input_coord, + mc2_weight_coord, + mc2_output_1_coord, + mc2_output_2_coord), + }, + }, + }, + device_type, + }; - std::unordered_set result = - perform_shard_expansion_for_invocation(input); + auto mk_op_value = + [&](size_t src_node_id, + TensorSlotName src_slot_name, + TensorSlotName use_slot_name, + std::optional const &shard_coord) + -> DynamicValueAttrs { + bidict tensor_binding = + dynamic_node_mapping_bindings_for_slot_name(node_mapping, + use_slot_name); + return mk_value( + src_node_id, src_slot_name, tensor_binding, shard_coord); + }; - auto mk_invocation_shard = - [&](MachineSpaceCoordinate const &device_coord, - ParallelTensorSpaceCoordinate const &input_shard_coord, - ParallelTensorSpaceCoordinate const &weight_shard_coord, - ParallelTensorSpaceCoordinate const &output_1_shard_coord, - ParallelTensorSpaceCoordinate const &output_2_shard_coord) - -> DynamicNodeInvocation { - return DynamicNodeInvocation{ + DynamicNodeInvocation input = DynamicNodeInvocation{ /*inputs=*/{ { mk_slot(TensorSlotName::INPUT), mk_op_value(0, TensorSlotName::OUTPUT, TensorSlotName::INPUT, - input_shard_coord), + std::nullopt), }, { mk_slot(TensorSlotName::WEIGHT), mk_op_value(1, TensorSlotName::OUTPUT, TensorSlotName::WEIGHT, - weight_shard_coord), + std::nullopt), }, }, /*node_attrs=*/ DynamicNodeAttrs{ /*task_type=*/std::nullopt, - /*device_coord=*/device_coord, - /*mapping=*/mapped_task_group, + /*device_coord=*/std::nullopt, + /*mapping=*/node_mapping, /*op_attrs=*/std::nullopt, /*layer_guid=*/ dynamic_layer_guid_t{parallel_layer_guid_t{Node{20}}}, @@ -239,112 +196,173 @@ TEST_SUITE(FF_TEST_SUITE) { mk_op_value(20, TensorSlotName::OUTPUT_1, TensorSlotName::OUTPUT_1, - output_1_shard_coord), + std::nullopt), }, { mk_slot(TensorSlotName::OUTPUT_2), mk_op_value(20, TensorSlotName::OUTPUT_2, TensorSlotName::OUTPUT_2, - output_2_shard_coord), + std::nullopt), }, }, }; - }; - std::unordered_set correct = { - mk_invocation_shard(mc1, - mc1_input_coord, - mc1_weight_coord, - mc1_output_1_coord, - mc1_output_2_coord), - mk_invocation_shard(mc2, - mc2_input_coord, - mc2_weight_coord, - mc2_output_1_coord, - mc2_output_2_coord), - }; + std::unordered_set result = + perform_shard_expansion_for_invocation(input); - CHECK(result.size() == correct.size()); - CHECK(result == correct); - } + auto mk_invocation_shard = + [&](device_id_t const &device_coord, + ParallelTensorSpaceCoordinate const &input_shard_coord, + ParallelTensorSpaceCoordinate const &weight_shard_coord, + ParallelTensorSpaceCoordinate const &output_1_shard_coord, + ParallelTensorSpaceCoordinate const &output_2_shard_coord) + -> DynamicNodeInvocation { + return DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mk_op_value(0, + TensorSlotName::OUTPUT, + TensorSlotName::INPUT, + input_shard_coord), + }, + { + mk_slot(TensorSlotName::WEIGHT), + mk_op_value(1, + TensorSlotName::OUTPUT, + TensorSlotName::WEIGHT, + weight_shard_coord), + }, + }, + /*node_attrs=*/ + DynamicNodeAttrs{ + /*task_type=*/std::nullopt, + /*device_coord=*/device_coord, + /*mapping=*/node_mapping, + /*op_attrs=*/std::nullopt, + /*layer_guid=*/ + dynamic_layer_guid_t{parallel_layer_guid_t{Node{20}}}, + /*per_device_op_state=*/std::nullopt, + }, + /*outputs=*/ + { + { + mk_slot(TensorSlotName::OUTPUT_1), + mk_op_value(20, + TensorSlotName::OUTPUT_1, + TensorSlotName::OUTPUT_1, + output_1_shard_coord), + }, + { + mk_slot(TensorSlotName::OUTPUT_2), + mk_op_value(20, + TensorSlotName::OUTPUT_2, + TensorSlotName::OUTPUT_2, + output_2_shard_coord), + }, + }, + }; + }; - TEST_CASE("perform_shard_expansion_for_invocation (copy)") { - MachineSpaceCoordinate mc1 = mk_machine_coord(0_n, 0_n); - MachineSpaceCoordinate mc2 = mk_machine_coord(1_n, 0_n); - MachineSpaceCoordinate mc3 = mk_machine_coord(2_n, 0_n); - MachineSpaceCoordinate mc4 = mk_machine_coord(3_n, 0_n); + std::unordered_set correct = { + mk_invocation_shard(mk_device_id(mc1), + mc1_input_coord, + mc1_weight_coord, + mc1_output_1_coord, + mc1_output_2_coord), + mk_invocation_shard(mk_device_id(mc2), + mc2_input_coord, + mc2_weight_coord, + mc2_output_1_coord, + mc2_output_2_coord), + }; - ParallelTensorSpaceCoordinate pt1 = mk_pt_coord(0_n, 0_n, 0_n, 0_n); - ParallelTensorSpaceCoordinate pt2 = mk_pt_coord(0_n, 1_n, 0_n, 0_n); + CHECK(result.size() == correct.size()); + CHECK(result == correct); + } - bidict src_binding{ - {pt1, mc1}, - {pt2, mc2}, - }; - bidict dst_binding{ - {pt1, mc3}, - {pt2, mc4}, - }; + SUBCASE("for copy operator") { + device_id_t dev1 = mk_device_id(mk_machine_coord(0_n, 0_n)); + device_id_t dev2 = mk_device_id(mk_machine_coord(1_n, 0_n)); + device_id_t dev3 = mk_device_id(mk_machine_coord(2_n, 0_n)); + device_id_t dev4 = mk_device_id(mk_machine_coord(3_n, 0_n)); - DynamicNodeInvocation input = DynamicNodeInvocation{ - /*inputs=*/{ + ParallelTensorSpaceCoordinate pt1 = mk_pt_coord(0_n, 0_n, 0_n, 0_n); + ParallelTensorSpaceCoordinate pt2 = mk_pt_coord(0_n, 1_n, 0_n, 0_n); + + bidict src_binding{ + {pt1, dev1}, + {pt2, dev2}, + }; + bidict dst_binding{ + {pt1, dev3}, + {pt2, dev4}, + }; + + DynamicNodeInvocation input = DynamicNodeInvocation{ + /*inputs=*/{ + { + mk_slot(TensorSlotName::INPUT), + mk_value( + 0, TensorSlotName::OUTPUT, src_binding, std::nullopt), + }, + }, + /*node_attrs=*/ + DynamicNodeAttrs{ + /*task_type=*/std::nullopt, + /*device_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*op_attrs=*/TrainingOperationAttrs{CopyAttrs{}}, + /*layer_guid=*/dynamic_layer_guid_t{dynamic_copy_layer_guid_t{}}, + /*per_device_op_state=*/std::nullopt, + }, + /*outputs=*/ + { + { + mk_slot(TensorSlotName::OUTPUT), + mk_value( + 20, TensorSlotName::OUTPUT, dst_binding, std::nullopt), + }, + }, + }; + + std::unordered_set result = + perform_shard_expansion_for_invocation(input); + + auto mk_invocation_shard = + [&](device_id_t const &device_coord, + ParallelTensorSpaceCoordinate const &tensor_shard_coord) + -> DynamicNodeInvocation { + DynamicNodeInvocation result = input; + result.inputs = { { mk_slot(TensorSlotName::INPUT), - mk_value(0, TensorSlotName::OUTPUT, src_binding, std::nullopt), + mk_value( + 0, TensorSlotName::OUTPUT, src_binding, tensor_shard_coord), }, - }, - /*node_attrs=*/ - DynamicNodeAttrs{ - /*task_type=*/std::nullopt, - /*device_coord=*/std::nullopt, - /*mapping=*/std::nullopt, - /*op_attrs=*/TrainingOperationAttrs{CopyAttrs{}}, - /*layer_guid=*/dynamic_layer_guid_t{dynamic_copy_layer_guid_t{}}, - /*per_device_op_state=*/std::nullopt, - }, - /*outputs=*/ - { + }; + // See perform_shard_expansion_for_copy in shard_expansion.cc for explanation of the choice of device placement. + result.node_attrs.device_coord = device_coord; + result.outputs = { { mk_slot(TensorSlotName::OUTPUT), - mk_value(20, TensorSlotName::OUTPUT, dst_binding, std::nullopt), + mk_value(20, + TensorSlotName::OUTPUT, + dst_binding, + tensor_shard_coord), }, - }, - }; - - std::unordered_set result = - perform_shard_expansion_for_invocation(input); - - auto mk_invocation_shard = - [&](MachineSpaceCoordinate const &device_coord, - ParallelTensorSpaceCoordinate const &tensor_shard_coord) - -> DynamicNodeInvocation { - DynamicNodeInvocation result = input; - result.inputs = { - { - mk_slot(TensorSlotName::INPUT), - mk_value( - 0, TensorSlotName::OUTPUT, src_binding, tensor_shard_coord), - }, + }; + return result; }; - // See perform_shard_expansion_for_copy in shard_expansion.cc for explanation of the choice of device placement. - result.node_attrs.device_coord = device_coord; - result.outputs = { - { - mk_slot(TensorSlotName::OUTPUT), - mk_value( - 20, TensorSlotName::OUTPUT, dst_binding, tensor_shard_coord), - }, - }; - return result; - }; - std::unordered_set correct = { - mk_invocation_shard(mc1, pt1), - mk_invocation_shard(mc2, pt2), - }; + std::unordered_set correct = { + mk_invocation_shard(dev1, pt1), + mk_invocation_shard(dev2, pt2), + }; - CHECK(result.size() == correct.size()); - CHECK(result == correct); + CHECK(result.size() == correct.size()); + CHECK(result == correct); + } } } diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc new file mode 100644 index 0000000000..3fb94bb98a --- /dev/null +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc @@ -0,0 +1,157 @@ +#include "task-spec/dynamic_graph/update_insertion.h" +#include "op-attrs/initializer_attrs.h" +#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" +#include "task-spec/dynamic_graph/dynamic_tensor_role.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("perform_update_insertion") { + dynamic_layer_guid_t layer_guid = dynamic_layer_guid_t{ + parallel_layer_guid_t{ + Node{0}, + }, + }; + + dynamic_tensor_guid_t tensor_guid = dynamic_tensor_guid_t{ + parallel_tensor_guid_t{ + KwargDataflowOutput{ + /*node=*/Node{0}, + /*slot_name=*/TensorSlotName::OUTPUT, + }, + }, + }; + + TrainingOperationAttrs weight_attrs = TrainingOperationAttrs{ + PCGOperatorAttrs{ + WeightAttrs{ + /*tensor_shape=*/TensorShape{ + /*dims=*/TensorDims{ + FFOrdered{ + 4_p, + 3_p, + }, + }, + /*data_type=*/DataType::FLOAT, + }, + /*initializer=*/make_zero_initializer(), + }, + }, + }; + + DynamicNodeInvocation weight_invocation = DynamicNodeInvocation{ + /*inputs=*/{}, + /*node_attrs=*/ + DynamicNodeAttrs{ + /*task_type=*/DynamicTaskType::FWD, + /*device_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*op_attrs=*/weight_attrs, + /*layer_guid=*/layer_guid, + /*per_device_op_state=*/std::nullopt, + }, + /*outputs=*/ + std::unordered_map{ + { + DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(), + }, + DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/mk_dynamic_tensor_role_fwd(), + }, + }, + }, + }; + + DynamicOpenDataflowGraph input = + dynamic_open_dataflow_graph_from_invocation_set({weight_invocation}); + + OptimizerAttrs optimizer_attrs = OptimizerAttrs{ + SGDOptimizerAttrs{ + /*lr=*/0.001, + /*momentum=*/0.9, + /*nesterov=*/false, + /*weight_decay=*/0.001, + }, + }; + + DynamicOpenDataflowGraph result = + perform_update_insertion(input, optimizer_attrs); + + DynamicOpenDataflowGraph correct = + dynamic_open_dataflow_graph_from_invocation_set( + /*invocations=*/{ + weight_invocation, + DynamicNodeInvocation{ + /*inputs=*/{ + { + DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/ + mk_dynamic_tensor_role_fwd(), + }, + DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/mk_dynamic_tensor_role_fwd(), + }, + }, + { + DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/ + mk_dynamic_tensor_role_bwd(), + }, + DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/mk_dynamic_tensor_role_bwd(), + }, + }, + { + DynamicTensorSlot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/ + mk_dynamic_tensor_role_opt( + OptimizerSlotName::SGD_V), + }, + DynamicValueAttrs{ + /*tensor_guid=*/tensor_guid, + /*parallel_tensor_shape=*/std::nullopt, + /*shard_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*accessor=*/std::nullopt, + /*role=*/ + mk_dynamic_tensor_role_opt( + OptimizerSlotName::SGD_V), + }, + }, + }, + /*node_attrs=*/ + DynamicNodeAttrs{ + /*task_type=*/DynamicTaskType::UPD, + /*device_coord=*/std::nullopt, + /*mapping=*/std::nullopt, + /*op_attrs=*/weight_attrs, + /*layer_guid=*/layer_guid, + /*per_device_op_state=*/std::nullopt, + }, + /*outputs=*/{}}, + }); + + CHECK(result == correct); + } +} From e4fbc04644c0bb38522b1abe8f5161f097d9e6e7 Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Fri, 12 Jun 2026 20:27:38 -0700 Subject: [PATCH 3/5] Address PR comments --- flake.lock | 27 ++- .../computation_graph_instance.h | 12 +- .../cost_estimator/local_cost_estimator.h | 8 +- .../local_task_argument_accessor.h | 8 +- .../per_device_op_state_initialization.h | 6 +- .../include/local-execution/task_execution.h | 4 +- .../computation_graph_instance.cc | 12 +- .../cost_estimator/local_cost_estimator.cc | 4 +- .../local_task_argument_accessor.cc | 4 +- .../per_device_op_state_initialization.cc | 4 +- .../src/local-execution/task_execution.cc | 4 +- .../computation_graph_instance.cc | 10 +- .../cost_estimator/local_cost_estimator.cc | 4 +- .../local_task_argument_accessor.cc | 2 +- ..._space_to_parallel_tensor_space_mapping.cc | 21 -- .../include/pcg/device_in_node_idx_t.dtg.toml | 22 ++ lib/pcg/include/pcg/node_idx_t.dtg.toml | 21 ++ .../v1_mapped_parallel_computation_graph.cc | 1 - .../include/realm-execution/address_space.h | 14 ++ ...ce_specific_managed_per_device_ff_handle.h | 6 +- .../realm-execution/device_specific_ptr.h | 12 +- ...buted_per_device_op_state_initialization.h | 3 +- .../realm-execution/fmt/realm_processor.h | 34 +++ .../fmt/realm_processor_kind.h | 83 ++++++++ .../realm-execution/instance_allocation.h | 5 +- .../include/realm-execution/pcg_instance.h | 2 +- .../include/realm-execution/processor_query.h | 12 ++ .../include/realm-execution/realm_context.h | 20 +- .../include/realm-execution/realm_manager.h | 2 +- .../serializable_device_specific_ptr.dtg.toml | 4 +- .../src/realm-execution/address_space.cc | 17 ++ ...e_specific_managed_per_device_ff_handle.cc | 4 +- .../realm-execution/distributed_ff_handle.cc | 2 +- ...uted_per_device_op_state_initialization.cc | 9 +- .../realm-execution/fmt/realm_processor.cc | 9 + .../fmt/realm_processor_kind.cc | 9 + .../realm-execution/instance_allocation.cc | 9 +- .../src/realm-execution/pcg_instance.cc | 10 +- .../src/realm-execution/processor_kind.cc | 3 +- .../src/realm-execution/processor_query.cc | 13 ++ .../src/realm-execution/realm_context.cc | 193 ++++++++++++------ .../tasks/impl/ff_handle_init_task.cc | 2 +- .../src/realm-execution/tasks/impl/op_task.cc | 6 +- .../impl/per_device_op_state_init_task.cc | 8 +- .../fmt/realm_processor_kind.cc | 13 ++ .../test/src/realm-execution/test_e2e.cc | 2 - .../output_expr_to_result_sub_pcg_mapping.cc | 4 +- .../src/substitutions/pcg_pattern_match.cc | 4 +- .../include/task-spec/device_specific.h | 10 +- .../dynamic_graph/dynamic_node_attrs.dtg.toml | 6 +- .../dynamic_graph/dynamic_node_mapping.h | 6 +- .../task-spec/dynamic_graph/machine_slicing.h | 4 +- .../parallel_tensor_mapping.dtg.toml | 4 +- .../serializable_dynamic_node_attrs.dtg.toml | 6 +- ...t.dtg.toml => global_device_id_t.dtg.toml} | 2 +- .../include/task-spec/global_device_id_t.h | 15 ++ .../task-spec/local_device_id_t.dtg.toml | 23 +++ .../include/task-spec/per_device_op_state.h | 2 +- .../itask_argument_accessor.h | 4 +- .../task_argument_accessor.h | 2 +- .../task-spec/dynamic_graph/copy_insertion.cc | 6 +- .../dynamic_graph/dynamic_node_mapping.cc | 12 +- .../dynamic_open_dataflow_graph.cc | 9 + .../dynamic_graph/machine_slicing.cc | 10 +- .../serializable_dynamic_node_attrs.cc | 4 +- .../dynamic_graph/shard_expansion.cc | 18 +- .../src/task-spec/global_device_id_t.cc | 25 +++ .../src/task-spec/per_device_op_state.cc | 2 +- .../test/src/task-spec/device_specific.cc | 4 +- .../task-spec/dynamic_graph/copy_insertion.cc | 18 +- .../dynamic_graph/machine_slicing.cc | 10 +- .../dynamic_graph/shard_expansion.cc | 26 +-- .../binary_merge_disjoint_bidicts.h | 38 ++++ .../algorithms/merge_disjoint_bidicts.h | 33 +-- .../binary_merge_disjoint_bidicts.cc | 13 ++ .../algorithms/merge_disjoint_bidicts.cc | 10 + ...ts.cc => binary_merge_disjoint_bidicts.cc} | 12 +- 77 files changed, 724 insertions(+), 298 deletions(-) create mode 100644 lib/pcg/include/pcg/device_in_node_idx_t.dtg.toml create mode 100644 lib/pcg/include/pcg/node_idx_t.dtg.toml create mode 100644 lib/realm-execution/include/realm-execution/address_space.h create mode 100644 lib/realm-execution/include/realm-execution/fmt/realm_processor.h create mode 100644 lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h create mode 100644 lib/realm-execution/include/realm-execution/processor_query.h create mode 100644 lib/realm-execution/src/realm-execution/address_space.cc create mode 100644 lib/realm-execution/src/realm-execution/fmt/realm_processor.cc create mode 100644 lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc create mode 100644 lib/realm-execution/src/realm-execution/processor_query.cc create mode 100644 lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc rename lib/task-spec/include/task-spec/{device_id_t.dtg.toml => global_device_id_t.dtg.toml} (91%) create mode 100644 lib/task-spec/include/task-spec/global_device_id_t.h create mode 100644 lib/task-spec/include/task-spec/local_device_id_t.dtg.toml create mode 100644 lib/task-spec/src/task-spec/global_device_id_t.cc create mode 100644 lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h create mode 100644 lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc rename lib/utils/test/src/utils/bidict/algorithms/{merge_disjoint_bidicts.cc => binary_merge_disjoint_bidicts.cc} (72%) diff --git a/flake.lock b/flake.lock index e52833e4ed..71bb407e27 100644 --- a/flake.lock +++ b/flake.lock @@ -63,14 +63,15 @@ ], "nixpkgs": [ "nixpkgs" - ] + ], + "python38-nixpkgs": "python38-nixpkgs" }, "locked": { - "lastModified": 1778104328, - "narHash": "sha256-bn0G8xDqBrVjp5htw1i3u8fPdPMVtoZXFzX7hJ6m9YY=", + "lastModified": 1780709897, + "narHash": "sha256-55VJWpNnt/tMEP8tadRQSI66cX1BgJS9UNy/EkaQt3A=", "ref": "refs/heads/master", - "rev": "535ac756b2674dc10051b37be978b9d2cb9f817d", - "revCount": 156, + "rev": "18d04112d94ee0a1c173a30923aac9fa2b505539", + "revCount": 160, "type": "git", "url": "https://git.sr.ht/~lockshaw/proj" }, @@ -79,6 +80,22 @@ "url": "https://git.sr.ht/~lockshaw/proj" } }, + "python38-nixpkgs": { + "locked": { + "lastModified": 1645131486, + "narHash": "sha256-AuWJe0TiqHD6L7Vzcar4QOblMGubKHsuVs3YFMCRYl0=", + "owner": "nixos", + "repo": "nixpkgs", + "rev": "7592790b9e02f7f99ddcb1bd33fd44ff8df6a9a7", + "type": "github" + }, + "original": { + "owner": "nixos", + "repo": "nixpkgs", + "rev": "7592790b9e02f7f99ddcb1bd33fd44ff8df6a9a7", + "type": "github" + } + }, "root": { "inputs": { "flake-utils": "flake-utils", diff --git a/lib/local-execution/include/local-execution/computation_graph_instance.h b/lib/local-execution/include/local-execution/computation_graph_instance.h index aa05cbd582..e6c7ded9c1 100644 --- a/lib/local-execution/include/local-execution/computation_graph_instance.h +++ b/lib/local-execution/include/local-execution/computation_graph_instance.h @@ -7,7 +7,7 @@ #include "kernels/profiling_settings.dtg.h" #include "local-execution/loss_config.dtg.h" #include "pcg/computation_graph.dtg.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" @@ -49,31 +49,31 @@ ComputationGraphInstance create_computation_graph_instance( Allocator &allocator, ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, - device_id_t device_idx); + global_device_id_t global_device_id); std::unordered_map> perform_all_passes_for_computation_graph_instance( ComputationGraphInstance &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx); + global_device_id_t global_device_id); std::unordered_map> perform_forward_pass_for_computation_graph_instance( ComputationGraphInstance const &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx); + global_device_id_t global_device_id); std::unordered_map> perform_backward_pass_for_computation_graph_instance( ComputationGraphInstance const &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx); + global_device_id_t global_device_id); void perform_update_pass_for_computation_graph_instance( ComputationGraphInstance &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx); + global_device_id_t global_device_id); } // namespace FlexFlow diff --git a/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h b/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h index a978f3c996..15aa267eec 100644 --- a/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h +++ b/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h @@ -5,7 +5,7 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "pcg/machine_interconnect_specification.dtg.h" namespace FlexFlow { @@ -16,7 +16,7 @@ struct LocalCostEstimator : public ICostEstimator { Allocator &allocator, ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, - device_id_t device_idx); + global_device_id_t device_idx); LocalCostEstimator(LocalCostEstimator const &) = delete; LocalCostEstimator(LocalCostEstimator &&) = delete; @@ -31,7 +31,7 @@ struct LocalCostEstimator : public ICostEstimator { Allocator allocator; ProfilingSettings profiling_settings; device_handle_t device_handle; - device_id_t device_idx; + global_device_id_t device_idx; }; CHECK_RC_COPY_VIRTUAL_COMPLIANT(LocalCostEstimator); @@ -40,7 +40,7 @@ CostEstimator get_local_cost_estimator( Allocator &allocator, ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, - device_id_t device_idx); + global_device_id_t device_idx); } // namespace FlexFlow diff --git a/lib/local-execution/include/local-execution/local_task_argument_accessor.h b/lib/local-execution/include/local-execution/local_task_argument_accessor.h index a7df8814d0..38b1023cd2 100644 --- a/lib/local-execution/include/local-execution/local_task_argument_accessor.h +++ b/lib/local-execution/include/local-execution/local_task_argument_accessor.h @@ -2,7 +2,7 @@ #define _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_LOCAL_TASK_ARGUMENT_ACCESSOR_H #include "kernels/accessor.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" #include "task-spec/task_argument_accessor/itask_argument_accessor.h" #include "task-spec/task_argument_accessor/task_tensor_parameter.dtg.h" @@ -21,7 +21,7 @@ struct LocalTaskArgumentAccessor : public ITaskArgumentAccessor { std::optional const &loss_attrs, std::optional const &per_device_op_state, std::optional const &optimizer_attrs, - device_id_t device_idx); + global_device_id_t device_idx); LocalTaskArgumentAccessor(LocalTaskArgumentAccessor const &) = delete; LocalTaskArgumentAccessor(LocalTaskArgumentAccessor &&) = delete; @@ -41,7 +41,7 @@ struct LocalTaskArgumentAccessor : public ITaskArgumentAccessor { Allocator get_allocator() const override; - device_id_t get_device_idx() const override; + global_device_id_t get_device_idx() const override; private: Allocator allocator; @@ -56,7 +56,7 @@ struct LocalTaskArgumentAccessor : public ITaskArgumentAccessor { std::optional per_device_op_state; std::optional optimizer_attrs; - device_id_t device_idx; + global_device_id_t device_idx; }; CHECK_RC_COPY_VIRTUAL_COMPLIANT(LocalTaskArgumentAccessor); diff --git a/lib/local-execution/include/local-execution/per_device_op_state_initialization.h b/lib/local-execution/include/local-execution/per_device_op_state_initialization.h index 174e51f307..90a1ffc688 100644 --- a/lib/local-execution/include/local-execution/per_device_op_state_initialization.h +++ b/lib/local-execution/include/local-execution/per_device_op_state_initialization.h @@ -4,7 +4,7 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" @@ -18,7 +18,7 @@ DynamicNodeInvocation ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, OptimizerAttrs const &optimizer_attrs, - device_id_t device_idx); + global_device_id_t device_idx); /** * @brief Initialize all operators and save the per-device op state @@ -29,7 +29,7 @@ DynamicOpenDataflowGraph perform_per_device_op_state_initialization( ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, OptimizerAttrs const &optimizer_attrs, - device_id_t device_idx); + global_device_id_t device_idx); } // namespace FlexFlow diff --git a/lib/local-execution/include/local-execution/task_execution.h b/lib/local-execution/include/local-execution/task_execution.h index 3bb0c6b92b..83f69a96e3 100644 --- a/lib/local-execution/include/local-execution/task_execution.h +++ b/lib/local-execution/include/local-execution/task_execution.h @@ -16,7 +16,7 @@ TaskArgumentAccessor make_task_argument_accessor_for_invocation( device_handle_t const &ff_handle, std::optional const &per_device_op_state, std::optional const &optimizer_attrs, - device_id_t device_idx); + global_device_id_t device_idx); std::optional execute_dynamic_node_invocation( DynamicNodeInvocation const &invocation, @@ -25,7 +25,7 @@ std::optional execute_dynamic_node_invocation( device_handle_t const &ff_handle, std::optional const &per_device_op_state, std::optional const &optimizer_attrs, - device_id_t device_idx); + global_device_id_t device_idx); } // namespace FlexFlow diff --git a/lib/local-execution/src/local-execution/computation_graph_instance.cc b/lib/local-execution/src/local-execution/computation_graph_instance.cc index d2781472b0..2bc15d1187 100644 --- a/lib/local-execution/src/local-execution/computation_graph_instance.cc +++ b/lib/local-execution/src/local-execution/computation_graph_instance.cc @@ -67,7 +67,7 @@ ComputationGraphInstance create_computation_graph_instance( Allocator &allocator, ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { DynamicOpenDataflowGraph dg = make_dynamic_open_dataflow_graph_from_cg(cg); dg = perform_pass_expansion(dg); @@ -116,7 +116,7 @@ static std::unordered_map> OptimizerAttrs const &optimizer_attrs, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { return unordered_map_from_pairs( transform(invocations, [&](DynamicNodeInvocation const &invocation) { std::optional timing = execute_dynamic_node_invocation( @@ -141,7 +141,7 @@ std::unordered_map> ComputationGraphInstance &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { std::vector execution_order = instance.get_execution_order(); std::unordered_map> @@ -161,7 +161,7 @@ std::unordered_map> ComputationGraphInstance const &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { std::vector execution_order = filter(instance.get_execution_order(), [](DynamicNodeInvocation const &invocation) { @@ -184,7 +184,7 @@ std::unordered_map> ComputationGraphInstance const &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { std::vector execution_order = filter(instance.get_execution_order(), [](DynamicNodeInvocation const &invocation) { @@ -206,7 +206,7 @@ void perform_update_pass_for_computation_graph_instance( ComputationGraphInstance &instance, ProfilingSettings const &profiling_settings, device_handle_t const &ff_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { std::vector execution_order = filter(instance.get_execution_order(), [](DynamicNodeInvocation const &invocation) { diff --git a/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc b/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc index 27a1c39677..f151ab1d17 100644 --- a/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc +++ b/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc @@ -31,7 +31,7 @@ LocalCostEstimator::LocalCostEstimator( Allocator &allocator, ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, - device_id_t device_idx) + global_device_id_t device_idx) : interconnect_specification(interconnect_specification), allocator(allocator), profiling_settings(profiling_settings), device_handle(device_handle), device_idx(device_idx) {} @@ -181,7 +181,7 @@ CostEstimator get_local_cost_estimator( Allocator &allocator, ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, - device_id_t device_idx) { + global_device_id_t device_idx) { return CostEstimator::create(interconnect_specification, allocator, profiling_settings, diff --git a/lib/local-execution/src/local-execution/local_task_argument_accessor.cc b/lib/local-execution/src/local-execution/local_task_argument_accessor.cc index 5fbe207e6c..a5f553db0b 100644 --- a/lib/local-execution/src/local-execution/local_task_argument_accessor.cc +++ b/lib/local-execution/src/local-execution/local_task_argument_accessor.cc @@ -16,7 +16,7 @@ LocalTaskArgumentAccessor::LocalTaskArgumentAccessor( std::optional const &loss_attrs, std::optional const &per_device_op_state, std::optional const &optimizer_attrs, - device_id_t device_idx) + global_device_id_t device_idx) : allocator(allocator), tensor_slots_backing(tensor_slots_backing), profiling_settings(profiling_settings), ff_handle(ff_handle), op_attrs(op_attrs), loss_attrs(loss_attrs), @@ -105,7 +105,7 @@ Allocator LocalTaskArgumentAccessor::get_allocator() const { return this->allocator; } -device_id_t LocalTaskArgumentAccessor::get_device_idx() const { +global_device_id_t LocalTaskArgumentAccessor::get_device_idx() const { return this->device_idx; } diff --git a/lib/local-execution/src/local-execution/per_device_op_state_initialization.cc b/lib/local-execution/src/local-execution/per_device_op_state_initialization.cc index bf72843daf..571e3dd1c1 100644 --- a/lib/local-execution/src/local-execution/per_device_op_state_initialization.cc +++ b/lib/local-execution/src/local-execution/per_device_op_state_initialization.cc @@ -24,7 +24,7 @@ DynamicNodeInvocation ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, OptimizerAttrs const &optimizer_attrs, - device_id_t device_idx) { + global_device_id_t device_idx) { if (!i.node_attrs.op_attrs.has_value() || !i.node_attrs.op_attrs.value().is_pcg_op()) { return i; @@ -61,7 +61,7 @@ DynamicOpenDataflowGraph perform_per_device_op_state_initialization( ProfilingSettings const &profiling_settings, device_handle_t const &device_handle, OptimizerAttrs const &optimizer_attrs, - device_id_t device_idx) { + global_device_id_t device_idx) { ASSERT(no_nodes_are_initialized(dg)); DynamicOpenDataflowGraph result = transform_dynamic_invocation_set( diff --git a/lib/local-execution/src/local-execution/task_execution.cc b/lib/local-execution/src/local-execution/task_execution.cc index 0b1bc3513a..48b9b89d94 100644 --- a/lib/local-execution/src/local-execution/task_execution.cc +++ b/lib/local-execution/src/local-execution/task_execution.cc @@ -49,7 +49,7 @@ TaskArgumentAccessor make_task_argument_accessor_for_invocation( device_handle_t const &ff_handle, std::optional const &per_device_op_state, std::optional const &optimizer_attrs, - device_id_t device_idx) { + global_device_id_t device_idx) { auto make_param = [&](DynamicTensorSlot const &slot) { return make_task_tensor_parameter_from_dynamic_slot(slot, optimizer_attrs); }; @@ -88,7 +88,7 @@ std::optional execute_dynamic_node_invocation( device_handle_t const &ff_handle, std::optional const &per_device_op_state, std::optional const &optimizer_attrs, - device_id_t device_idx) { + global_device_id_t device_idx) { TaskArgumentAccessor arg_accessor = make_task_argument_accessor_for_invocation( /*invocation=*/invocation, diff --git a/lib/local-execution/test/src/local-execution/computation_graph_instance.cc b/lib/local-execution/test/src/local-execution/computation_graph_instance.cc index a4049a609f..69091b5d2e 100644 --- a/lib/local-execution/test/src/local-execution/computation_graph_instance.cc +++ b/lib/local-execution/test/src/local-execution/computation_graph_instance.cc @@ -139,7 +139,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*nesterov=*/false, /*weight_decay=*/0.001}}; device_handle_t ff_handle = cpu_make_device_handle_t(); - device_id_t device_idx = device_id_t{ + global_device_id_t global_device_id = global_device_id_t{ /*coord=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, @@ -164,7 +164,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*allocator=*/allocator, /*profiling_settings=*/ProfilingSettings{0, 0}, /*device_handle=*/ff_handle, - /*device_idx=*/device_idx); + /*global_device_id=*/global_device_id); // begin training loop int num_epochs = 5; @@ -175,7 +175,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*instance=*/computation_graph_instance, /*profiling_settings=*/ProfilingSettings{0, 0}, /*ff_handle=*/ff_handle, - /*device_idx=*/device_idx); + /*global_device_id=*/global_device_id); loss_values.push_back(copy_tensor_accessor_r( computation_graph_instance.get_loss_tensor_accessor().value(), allocator)); @@ -312,7 +312,7 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*weight_decay=*/0.001, }, }; - device_id_t device_idx = device_id_t{ + global_device_id_t device_idx = global_device_id_t{ /*coord=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, @@ -433,7 +433,7 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }, }; - device_id_t device_idx = device_id_t{ + global_device_id_t device_idx = global_device_id_t{ /*coord=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, diff --git a/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc b/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc index 34cfbb4e3d..8d1506b5fc 100644 --- a/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc +++ b/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc @@ -18,7 +18,7 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("LocalCostEstimator") { Allocator allocator = create_local_cpu_memory_allocator(); device_handle_t ff_handle = cpu_make_device_handle_t(); - device_id_t device_idx = device_id_t{ + global_device_id_t device_idx = global_device_id_t{ /*coord=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, @@ -93,7 +93,7 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { Allocator allocator = create_local_cuda_memory_allocator(); - device_id_t device_idx = device_id_t{ + global_device_id_t device_idx = global_device_id_t{ /*coord=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, diff --git a/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc b/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc index 9594e13336..dba4464611 100644 --- a/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc +++ b/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc @@ -51,7 +51,7 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; - device_id_t device_idx = device_id_t{ + global_device_id_t device_idx = global_device_id_t{ /*coord=*/MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, diff --git a/lib/op-attrs/src/op-attrs/parallel_tensor_space_to_parallel_tensor_space_mapping.cc b/lib/op-attrs/src/op-attrs/parallel_tensor_space_to_parallel_tensor_space_mapping.cc index 2a161838cd..b502c456e6 100644 --- a/lib/op-attrs/src/op-attrs/parallel_tensor_space_to_parallel_tensor_space_mapping.cc +++ b/lib/op-attrs/src/op-attrs/parallel_tensor_space_to_parallel_tensor_space_mapping.cc @@ -11,27 +11,6 @@ ParallelTensorSpaceToParallelTensorSpaceMapping ParallelTensorDimDegrees const &l_degrees, ParallelTensorDimDegrees const &r_degrees) { - // TODO(@lockshaw)(#pr): - // { - // std::unordered_set - // l_dims = - // unordered_set_of(get_nontrivial_parallel_tensor_dim_indices(l_degrees)); - // std::unordered_set - // projection_input_dims = input_dims_of_projection(projection); - // - // ASSERT(l_dims == projection_input_dims); - // } - // - // { - // std::unordered_set - // r_dims = - // unordered_set_of(get_nontrivial_parallel_tensor_dim_indices(r_degrees)); - // std::unordered_set - // projection_output_dims = output_dims_of_projection(projection); - // - // ASSERT(r_dims == projection_output_dims); - // } - return ParallelTensorSpaceToParallelTensorSpaceMapping{ dim_domain_mapping_from_projection( /*projection=*/projection, diff --git a/lib/pcg/include/pcg/device_in_node_idx_t.dtg.toml b/lib/pcg/include/pcg/device_in_node_idx_t.dtg.toml new file mode 100644 index 0000000000..c7b73e7724 --- /dev/null +++ b/lib/pcg/include/pcg/device_in_node_idx_t.dtg.toml @@ -0,0 +1,22 @@ +namespace = "FlexFlow" +name = "device_in_node_idx_t" +type = "struct" +features = [ + "eq", + "ord", + "hash", + "json", + "fmt", + "rapidcheck", +] + +includes = [ + "utils/nonnegative_int/nonnegative_int.h", +] + +src_includes = [] + +[[fields]] +name = "raw" +type = "::FlexFlow::nonnegative_int" + diff --git a/lib/pcg/include/pcg/node_idx_t.dtg.toml b/lib/pcg/include/pcg/node_idx_t.dtg.toml new file mode 100644 index 0000000000..df592635b0 --- /dev/null +++ b/lib/pcg/include/pcg/node_idx_t.dtg.toml @@ -0,0 +1,21 @@ +namespace = "FlexFlow" +name = "node_idx_t" +type = "struct" +features = [ + "eq", + "ord", + "hash", + "json", + "fmt", + "rapidcheck", +] + +includes = [ + "utils/nonnegative_int/nonnegative_int.h", +] + +src_includes = [] + +[[fields]] +name = "raw" +type = "::FlexFlow::nonnegative_int" diff --git a/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc b/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc index 0c43ad06a0..ac264191a8 100644 --- a/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc +++ b/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc @@ -36,7 +36,6 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpaceCoordinate coord = MachineSpaceCoordinate{ /*node_idx=*/0_n, /*device_idx=*/0_n, - /*device_type=*/DeviceType::GPU, }; OperatorAtomicTaskShardBinding binding = OperatorAtomicTaskShardBinding{ diff --git a/lib/realm-execution/include/realm-execution/address_space.h b/lib/realm-execution/include/realm-execution/address_space.h new file mode 100644 index 0000000000..3d00bc2bd4 --- /dev/null +++ b/lib/realm-execution/include/realm-execution/address_space.h @@ -0,0 +1,14 @@ +#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_ADDRESS_SPACE_H +#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_ADDRESS_SPACE_H + +#include "pcg/node_idx_t.dtg.h" +#include "realm-execution/realm.h" + +namespace FlexFlow { + +node_idx_t node_idx_from_realm_address_space(Realm::AddressSpace); +Realm::AddressSpace realm_address_space_from_node_idx(node_idx_t); + +} // namespace FlexFlow + +#endif diff --git a/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h b/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h index e5ab8b597e..a7e265f777 100644 --- a/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h +++ b/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h @@ -4,17 +4,17 @@ #include "kernels/device_handle_t.dtg.h" #include "kernels/managed_per_device_ff_handle.h" #include "realm-execution/device_specific_ptr.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include namespace FlexFlow { DeviceSpecificPtr make_device_specific_managed_ff_handle( - device_id_t const &, std::optional const &); + global_device_id_t const &, std::optional const &); device_handle_t device_handle_t_from_device_specific_managed_ff_handle( - DeviceSpecificPtr const &, device_id_t); + DeviceSpecificPtr const &, global_device_id_t); } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/device_specific_ptr.h b/lib/realm-execution/include/realm-execution/device_specific_ptr.h index e8de4cfdad..f1126f64c1 100644 --- a/lib/realm-execution/include/realm-execution/device_specific_ptr.h +++ b/lib/realm-execution/include/realm-execution/device_specific_ptr.h @@ -1,7 +1,7 @@ #ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_DEVICE_SPECIFIC_PTR_H #define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_DEVICE_SPECIFIC_PTR_H -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include #include @@ -18,7 +18,7 @@ namespace FlexFlow { * transfer the pointers back-and-forth between workers and the controller * task. To prevent accidentally accessing one of these pointers on the wrong * device (as the pointer is only valid in the memory where it was created), we - * wrap them with \ref DeviceSpecificPtr, which holds the \ref device_id_t + * wrap them with \ref DeviceSpecificPtr, which holds the \ref global_device_id_t * where the pointer was created, and any attempt to interact with the raw * pointer value (i.e., \ref DeviceSpecificPtr::get) checks that the current * device matches the original device, and throws a readable error message if @@ -32,15 +32,15 @@ template struct DeviceSpecificPtr { public: DeviceSpecificPtr() = delete; - explicit DeviceSpecificPtr(device_id_t device_idx, std::optional ptr) + explicit DeviceSpecificPtr(global_device_id_t device_idx, std::optional ptr) : device_idx(device_idx), ptr(ptr) {} - std::optional get(device_id_t device_idx) const { + std::optional get(global_device_id_t device_idx) const { ASSERT(this->device_idx == device_idx); return this->ptr; } - device_id_t get_device_idx() const { + global_device_id_t get_device_idx() const { return this->device_idx; } @@ -49,7 +49,7 @@ struct DeviceSpecificPtr { } private: - device_id_t device_idx; + global_device_id_t device_idx; std::optional ptr; }; diff --git a/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h b/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h index f160146b96..5d52f8caaf 100644 --- a/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h +++ b/lib/realm-execution/include/realm-execution/distributed_per_device_op_state_initialization.h @@ -25,8 +25,7 @@ PerDeviceOpStateBacking perform_distributed_per_device_op_state_initialization( ProfilingSettings const &profiling_settings, DistributedFfHandle const &device_handle, OptimizerAttrs const &optimizer_attrs, - Realm::Event precondition, - DeviceType device_type); + Realm::Event precondition); } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/fmt/realm_processor.h b/lib/realm-execution/include/realm-execution/fmt/realm_processor.h new file mode 100644 index 0000000000..e2fadb6b18 --- /dev/null +++ b/lib/realm-execution/include/realm-execution/fmt/realm_processor.h @@ -0,0 +1,34 @@ +#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_FMT_REALM_PROCESSOR_H +#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_FMT_REALM_PROCESSOR_H + +#include "realm-execution/realm.h" +#include +#include + +namespace fmt { + +template +struct formatter< + ::FlexFlow::Realm::Processor, + Char, + std::enable_if_t::value>> + : formatter<::std::string> { + template + auto format(::FlexFlow::Realm::Processor const &m, FormatContext &ctx) + -> decltype(ctx.out()) { + + std::string result = fmt::format("", m.id); + + return formatter::format(result, ctx); + } +}; + +} // namespace fmt + +namespace FlexFlow { + +std::ostream &operator<<(std::ostream &s, ::FlexFlow::Realm::Processor const &m); + +} // namespace FlexFlow + +#endif diff --git a/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h b/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h new file mode 100644 index 0000000000..d75ab067cc --- /dev/null +++ b/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h @@ -0,0 +1,83 @@ +#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_FMT_REALM_PROCESSOR_KIND_H +#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_FMT_REALM_PROCESSOR_KIND_H + +#include "realm-execution/realm.h" +#include +#include +#include + +namespace fmt { + +template +struct formatter< + ::FlexFlow::Realm::Processor::Kind, + Char + > + : formatter<::std::string> { + template + auto format(::FlexFlow::Realm::Processor::Kind const &m, FormatContext &ctx) + -> decltype(ctx.out()) { + + std::string result; + switch (m) { + case ::FlexFlow::Realm::Processor::Kind::NO_KIND: { + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::TOC_PROC: { + // Throughput core + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::LOC_PROC: { + // Latency core + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::UTIL_PROC: { + // Utility core + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::IO_PROC: { + // I/O core + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::PROC_GROUP: { + // Processor group + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::PROC_SET: { + // Set of Processors for OpenMP/Kokkos etc. + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::OMP_PROC: { + // OpenMP (or similar) thread pool + result = ""; + break; + } + case ::FlexFlow::Realm::Processor::Kind::PY_PROC: { + // Python interpreter + result = ""; + break; + } + default: + PANIC("Unknown Processor::Kind {}", static_cast(m)); + }; + + return formatter::format(result, ctx); + } +}; + +} // namespace fmt + +namespace FlexFlow { + +std::ostream &operator<<(std::ostream &, ::FlexFlow::Realm::Processor::Kind const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/realm-execution/include/realm-execution/instance_allocation.h b/lib/realm-execution/include/realm-execution/instance_allocation.h index 1f2a6d9134..cc4af624a6 100644 --- a/lib/realm-execution/include/realm-execution/instance_allocation.h +++ b/lib/realm-execution/include/realm-execution/instance_allocation.h @@ -12,7 +12,7 @@ namespace FlexFlow { * on the device represented by \p device_coord. */ std::pair - perform_instance_allocation_for_value(device_id_t const &device_id, + perform_instance_allocation_for_value(global_device_id_t const &device_id, DynamicValueAttrs const &value, RealmContext &ctx); @@ -27,8 +27,7 @@ TensorInstanceBacking perform_instance_allocation( DynamicOpenDataflowGraph const &g, std::unordered_map const &preallocated, - RealmContext &ctx, - DeviceType device_type); + RealmContext &ctx); /** * @brief Destroys all of the instances held in \p instances. diff --git a/lib/realm-execution/include/realm-execution/pcg_instance.h b/lib/realm-execution/include/realm-execution/pcg_instance.h index 1e17856999..6c82e4c61c 100644 --- a/lib/realm-execution/include/realm-execution/pcg_instance.h +++ b/lib/realm-execution/include/realm-execution/pcg_instance.h @@ -11,7 +11,7 @@ #include "realm-execution/per_device_op_state_backing.dtg.h" #include "realm-execution/realm_context.h" #include "realm-execution/tensor_instance_backing.dtg.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" diff --git a/lib/realm-execution/include/realm-execution/processor_query.h b/lib/realm-execution/include/realm-execution/processor_query.h new file mode 100644 index 0000000000..693d488390 --- /dev/null +++ b/lib/realm-execution/include/realm-execution/processor_query.h @@ -0,0 +1,12 @@ +#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_PROCESSOR_QUERY_H +#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_PROCESSOR_QUERY_H + +#include "realm-execution/realm.h" + +namespace FlexFlow { + +std::set processor_set_from_query(Realm::Machine::ProcessorQuery const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/realm-execution/include/realm-execution/realm_context.h b/lib/realm-execution/include/realm-execution/realm_context.h index 0f2bed7de4..e06ac522d7 100644 --- a/lib/realm-execution/include/realm-execution/realm_context.h +++ b/lib/realm-execution/include/realm-execution/realm_context.h @@ -9,9 +9,10 @@ #include "pcg/machine_space_coordinate.dtg.h" #include "realm-execution/realm.h" #include "realm-execution/tasks/task_id_t.dtg.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include #include +#include "task-spec/local_device_id_t.dtg.h" namespace FlexFlow { @@ -32,8 +33,11 @@ struct RealmContext { /** \name Device mapping */ ///\{ - Realm::Processor map_device_coord_to_processor(device_id_t const &) const; - device_id_t map_processor_to_device_coord(Realm::Processor) const; + Realm::Processor processor_from_global_device_id(global_device_id_t const &); + global_device_id_t global_device_id_from_processor(Realm::Processor); + Realm::Processor processor_from_local_device_id(local_device_id_t const &) const; + local_device_id_t local_device_id_from_processor(Realm::Processor) const; + static Realm::Memory get_nearest_memory(Realm::Processor); ///\} @@ -41,7 +45,7 @@ struct RealmContext { ///\{ Realm::Processor get_current_processor() const; Allocator &get_current_device_allocator(); - device_id_t get_current_device_idx() const; + global_device_id_t get_current_global_device_id() const; ///\} /** \name Task creation */ @@ -109,15 +113,15 @@ struct RealmContext { */ Realm::Runtime get_runtime(); - void discover_machine_topology(); + bidict const &get_global_machine_topology(); -public: +private: Realm::Runtime runtime; Realm::Processor processor; Allocator allocator; std::vector outstanding_events; - std::optional> processors = - std::nullopt; + bidict local_machine_topology; + std::optional> cached_global_machine_topology = std::nullopt; }; } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/realm_manager.h b/lib/realm-execution/include/realm-execution/realm_manager.h index ed01f4a022..0aa81c0e74 100644 --- a/lib/realm-execution/include/realm-execution/realm_manager.h +++ b/lib/realm-execution/include/realm-execution/realm_manager.h @@ -6,7 +6,7 @@ #include "realm-execution/realm.h" #include "realm-execution/realm_context.h" #include "realm-execution/tasks/impl/controller_task.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" namespace FlexFlow { diff --git a/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml b/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml index 6a5675772e..492993edd2 100644 --- a/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml +++ b/lib/realm-execution/include/realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.toml @@ -9,7 +9,7 @@ features = [ ] includes = [ - "pcg/device_id_t.dtg.h", + "task-spec/global_device_id_t.dtg.h", "", "", ] @@ -21,7 +21,7 @@ src_includes = [ [[fields]] name = "device_idx" -type = "::FlexFlow::device_id_t" +type = "::FlexFlow::global_device_id_t" [[fields]] name = "ptr" diff --git a/lib/realm-execution/src/realm-execution/address_space.cc b/lib/realm-execution/src/realm-execution/address_space.cc new file mode 100644 index 0000000000..2886f16312 --- /dev/null +++ b/lib/realm-execution/src/realm-execution/address_space.cc @@ -0,0 +1,17 @@ +#include "realm-execution/address_space.h" + +namespace FlexFlow { + +node_idx_t node_idx_from_realm_address_space(Realm::AddressSpace address_space) { + return node_idx_t{ + nonnegative_int{ + static_cast(address_space), + }, + }; +} + +Realm::AddressSpace realm_address_space_from_node_idx(node_idx_t node_idx) { + return static_cast(node_idx.raw.unwrap_nonnegative()); +}; + +} // namespace FlexFlow diff --git a/lib/realm-execution/src/realm-execution/device_specific_managed_per_device_ff_handle.cc b/lib/realm-execution/src/realm-execution/device_specific_managed_per_device_ff_handle.cc index c16f17d168..bfe845e6ab 100644 --- a/lib/realm-execution/src/realm-execution/device_specific_managed_per_device_ff_handle.cc +++ b/lib/realm-execution/src/realm-execution/device_specific_managed_per_device_ff_handle.cc @@ -8,14 +8,14 @@ namespace FlexFlow { DeviceSpecificPtr make_device_specific_managed_ff_handle( - device_id_t const &device_id, + global_device_id_t const &device_id, std::optional const &managed_handle) { return DeviceSpecificPtr{device_id, managed_handle}; } device_handle_t device_handle_t_from_device_specific_managed_ff_handle( DeviceSpecificPtr const &device_specific, - device_id_t device_idx) { + global_device_id_t device_idx) { return device_handle_t_from_managed_ff_handle_ptr( device_specific.get(device_idx)); } diff --git a/lib/realm-execution/src/realm-execution/distributed_ff_handle.cc b/lib/realm-execution/src/realm-execution/distributed_ff_handle.cc index 2fdc31d2c5..2e5b832f2f 100644 --- a/lib/realm-execution/src/realm-execution/distributed_ff_handle.cc +++ b/lib/realm-execution/src/realm-execution/distributed_ff_handle.cc @@ -32,7 +32,7 @@ DistributedFfHandle proc.kind() == Realm::Processor::TOC_PROC) { handles.insert({proc, make_device_specific_managed_ff_handle( - ctx.get_current_device_idx(), std::nullopt)}); + ctx.get_current_global_device_id(), std::nullopt)}); } } diff --git a/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc b/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc index 015af2b1ae..a4a4240ff3 100644 --- a/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc +++ b/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc @@ -22,8 +22,7 @@ PerDeviceOpStateBacking perform_distributed_per_device_op_state_initialization( ProfilingSettings const &profiling_settings, DistributedFfHandle const &device_handle, OptimizerAttrs const &optimizer_attrs, - Realm::Event precondition, - DeviceType device_type) { + Realm::Event precondition) { // Initialize all operators and save the per-device op state ASSERT(no_nodes_are_initialized(dg)); @@ -32,15 +31,15 @@ PerDeviceOpStateBacking perform_distributed_per_device_op_state_initialization( DeviceSpecificPtr *> device_state_map; for (DynamicNodeInvocation const &invocation : dg.invocations) { - Realm::Processor target_proc = ctx.map_device_coord_to_processor( - assert_unwrap(invocation.node_attrs.device_coord)); + Realm::Processor target_proc = ctx.processor_from_global_device_id( + assert_unwrap(invocation.node_attrs.device_id)); TensorInstanceBacking tensor_backing = subset_tensor_instance_backing_for_invocation(tensor_instance_backing, invocation); DeviceSpecificPtr *device_state_ptr = - new DeviceSpecificPtr{ctx.get_current_device_idx(), + new DeviceSpecificPtr{ctx.get_current_global_device_id(), std::nullopt}; std::optional completion_event = diff --git a/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc b/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc new file mode 100644 index 0000000000..a4afe3cd1b --- /dev/null +++ b/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc @@ -0,0 +1,9 @@ +#include "realm-execution/fmt/realm_processor.h" + +namespace FlexFlow { + +std::ostream &operator<<(std::ostream &s, ::FlexFlow::Realm::Processor const &m) { + return s << fmt::to_string(m); +} + +} // namespace FlexFlow diff --git a/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc b/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc new file mode 100644 index 0000000000..2552ec7ef5 --- /dev/null +++ b/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc @@ -0,0 +1,9 @@ +#include "realm-execution/fmt/realm_processor_kind.h" + +namespace FlexFlow { + +std::ostream &operator<<(std::ostream &s, ::FlexFlow::Realm::Processor::Kind const &k) { + return (s << fmt::to_string(k)); +} + +} // namespace FlexFlow diff --git a/lib/realm-execution/src/realm-execution/instance_allocation.cc b/lib/realm-execution/src/realm-execution/instance_allocation.cc index 37bf0ca03d..8829cda275 100644 --- a/lib/realm-execution/src/realm-execution/instance_allocation.cc +++ b/lib/realm-execution/src/realm-execution/instance_allocation.cc @@ -22,14 +22,14 @@ namespace FlexFlow { std::pair - perform_instance_allocation_for_value(device_id_t const &device_id, + perform_instance_allocation_for_value(global_device_id_t const &device_id, DynamicValueAttrs const &value, RealmContext &ctx) { ASSERT(value.accessor == std::nullopt); TensorShape shape = get_piece_shape(value.parallel_tensor_shape.value()); - Realm::Processor proc = ctx.map_device_coord_to_processor(device_id); + Realm::Processor proc = ctx.processor_from_global_device_id(device_id); Realm::Memory memory = ctx.get_nearest_memory(proc); return ctx.create_instance(memory, shape, Realm::ProfilingRequestSet()); } @@ -38,8 +38,7 @@ TensorInstanceBacking perform_instance_allocation( DynamicOpenDataflowGraph const &g, std::unordered_map const &preallocated, - RealmContext &ctx, - DeviceType device_type) { + RealmContext &ctx) { ASSERT(no_tensors_are_allocated(g)); ASSERT(tensors_are_ready_for_allocation(g)); for (DynamicValueAttrs const &v : keys(preallocated)) { @@ -53,7 +52,7 @@ TensorInstanceBacking perform_instance_allocation( NOT_IMPLEMENTED(); } else { if (!contains_key(result.backing, v)) { - device_id_t device = assert_unwrap(n.device_coord); + global_device_id_t device = assert_unwrap(n.device_id); result.backing.insert(std::pair{ v, perform_instance_allocation_for_value(device, v, ctx)}); } diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc index 62747fda7a..e8a8494a55 100644 --- a/lib/realm-execution/src/realm-execution/pcg_instance.cc +++ b/lib/realm-execution/src/realm-execution/pcg_instance.cc @@ -117,11 +117,10 @@ PCGInstance create_pcg_instance( dg = perform_update_insertion(dg, optimizer_attrs); dg = perform_copy_insertion(dg); - debug_print_dynamic_open_dataflow_graph_as_dot(dg); dg = perform_shard_expansion(dg); TensorInstanceBacking tensor_instance_backing = - perform_instance_allocation(dg, inputs, ctx, device_type); + perform_instance_allocation(dg, inputs, ctx); logit_grad_value = transform(logit_grad_value, [&](DynamicValueAttrs const &lgv) { @@ -154,8 +153,7 @@ PCGInstance create_pcg_instance( profiling_settings, device_handle, optimizer_attrs, - ctx.get_outstanding_events(), - device_type); + ctx.get_outstanding_events()); // Compute the topological ordering of the graph auto [kwarg_graph, node_map] = @@ -199,8 +197,8 @@ static Realm::Event spawn_dynamic_node_invocation( invocation); auto spawn_task = [&]() { - Realm::Processor target_proc = ctx.map_device_coord_to_processor( - assert_unwrap(invocation.node_attrs.device_coord)); + Realm::Processor target_proc = ctx.processor_from_global_device_id( + assert_unwrap(invocation.node_attrs.device_id)); return spawn_op_task(ctx, target_proc, invocation, diff --git a/lib/realm-execution/src/realm-execution/processor_kind.cc b/lib/realm-execution/src/realm-execution/processor_kind.cc index 3a40de0a42..29a9b05274 100644 --- a/lib/realm-execution/src/realm-execution/processor_kind.cc +++ b/lib/realm-execution/src/realm-execution/processor_kind.cc @@ -1,4 +1,5 @@ #include "realm-execution/processor_kind.h" +#include "realm-execution/fmt/realm_processor_kind.h" #include namespace FlexFlow { @@ -11,7 +12,7 @@ DeviceType case Realm::Processor::Kind::TOC_PROC: return DeviceType::GPU; default: - PANIC("Unhandled Realm::Processor::Kind", processor_kind); + PANIC("Unhandled Realm::Processor::Kind", fmt::to_string(processor_kind)); } } diff --git a/lib/realm-execution/src/realm-execution/processor_query.cc b/lib/realm-execution/src/realm-execution/processor_query.cc new file mode 100644 index 0000000000..b5dd3e12da --- /dev/null +++ b/lib/realm-execution/src/realm-execution/processor_query.cc @@ -0,0 +1,13 @@ +#include "realm-execution/processor_query.h" + +namespace FlexFlow { + +std::set processor_set_from_query(Realm::Machine::ProcessorQuery const &pq) { + std::set result; + for (Realm::Processor p : pq) { + result.insert(p); + } + return result; +} + +} // namespace FlexFlow diff --git a/lib/realm-execution/src/realm-execution/realm_context.cc b/lib/realm-execution/src/realm-execution/realm_context.cc index 4cd7bdc221..07d60aa73c 100644 --- a/lib/realm-execution/src/realm-execution/realm_context.cc +++ b/lib/realm-execution/src/realm-execution/realm_context.cc @@ -15,17 +15,116 @@ #include "utils/one_to_many/one_to_many.h" #include "utils/optional.h" #include "utils/positive_int/positive_int.h" +#include "utils/bidict/algorithms/merge_disjoint_bidicts.h" +#include "realm-execution/address_space.h" +#include "utils/bidict/algorithms/bidict_from_enumerating.h" +#include "utils/bidict/algorithms/transform_values.h" +#include "utils/containers/group_by.h" +#include "task-spec/global_device_id_t.h" +#include "utils/containers/are_all_same.h" +#include "utils/containers/set_of.h" +#include "realm-execution/processor_query.h" +#include "realm-execution/fmt/realm_processor.h" +#include "realm-execution/fmt/realm_processor_kind.h" namespace FlexFlow { + +bidict + build_local_machine_topology(std::set const &local_procs) +{ + { + bool procs_are_local = are_all_same( + transform(local_procs, + [&](Realm::Processor p) -> Realm::AddressSpace { + return p.address_space(); + })); + ASSERT(procs_are_local); + } + + OneToMany by_proc_kind = + group_by(local_procs, + [](Realm::Processor p) -> Realm::Processor::Kind { + return p.kind(); + }); + + auto local_machine_topology_for_proc_kind = [&](Realm::Processor::Kind k) + -> bidict + { + if (!contains(by_proc_kind.left_values(), k)) { + return {}; + } + + bidict enumerated = + bidict_from_enumerating(set_of(by_proc_kind.at_l(k).unwrap_as_unordered_set())).reversed(); + + bidict result = + transform_values( + enumerated, + [&](nonnegative_int idx) -> local_device_id_t { + return local_device_id_t{ + /*idx=*/device_in_node_idx_t{idx}, + /*device_type=*/device_type_from_processor_kind(k), + }; + }); + + return result; + }; + + return binary_merge_disjoint_bidicts( + local_machine_topology_for_proc_kind(Realm::Processor::Kind::LOC_PROC), + local_machine_topology_for_proc_kind(Realm::Processor::Kind::TOC_PROC)); +} + +static bidict + build_global_machine_topology(std::set const &global_procs) +{ + OneToMany by_node_idx = + group_by(global_procs, + [](Realm::Processor p) -> node_idx_t { + return node_idx_from_realm_address_space(p.address_space()); + }); + + auto build_global_machine_topology_for_node = [&](node_idx_t const &node_idx) + -> bidict + { + std::set procs_for_node = set_of(by_node_idx.at_l(node_idx).unwrap_as_unordered_set()); + + bidict + local_topology_for_node = build_local_machine_topology(procs_for_node); + + return transform_values( + local_topology_for_node, + [&](local_device_id_t const &local_device_id) -> global_device_id_t { + return global_device_id_from_local(local_device_id, node_idx); + }); + }; + + return merge_disjoint_bidicts( + transform( + set_of(by_node_idx.left_values()), + build_global_machine_topology_for_node)); +} + +static bidict discover_local_machine_topology(Realm::Processor local_processor) { + Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); + pq.same_address_space_as(local_processor); + + return build_local_machine_topology(processor_set_from_query(pq)); +} + +static bidict discover_global_machine_topology() { + Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); + + return build_global_machine_topology(processor_set_from_query(pq)); +} + RealmContext::RealmContext(Realm::Processor processor) : processor(processor), allocator(get_realm_allocator( - processor, RealmContext::get_nearest_memory(processor))) { - if (processor != Realm::Processor::NO_PROC) { - this->discover_machine_topology(); - } -} + processor, RealmContext::get_nearest_memory(processor))), + local_machine_topology(discover_local_machine_topology(processor)) +{ } RealmContext::~RealmContext() { if (!this->outstanding_events.empty()) { @@ -34,23 +133,28 @@ RealmContext::~RealmContext() { } } -static std::tuple - convert_machine_space_coordinate(MachineSpaceCoordinate const &device_coord, - DeviceType device_type) { - Realm::AddressSpace as = int{device_coord.node_idx}; - Realm::Processor::Kind kind = processor_kind_from_device_type(device_type); - nonnegative_int proc_in_node = device_coord.device_idx; - return std::tuple{as, kind, proc_in_node}; +Realm::Processor RealmContext::processor_from_global_device_id( + global_device_id_t const &global_device_id) { + + return this->get_global_machine_topology().at_r(global_device_id); } -Realm::Processor RealmContext::map_device_coord_to_processor( - device_id_t const &device_id) const { - return assert_unwrap(this->processors).at_r(device_id); +global_device_id_t RealmContext::global_device_id_from_processor( + Realm::Processor processor) { + + return this->get_global_machine_topology().at_l(processor); } -device_id_t - RealmContext::map_processor_to_device_coord(Realm::Processor p) const { - return assert_unwrap(this->processors).at_l(p); +Realm::Processor RealmContext::processor_from_local_device_id( + local_device_id_t const &local_device_id) const { + + return this->local_machine_topology.at_r(local_device_id); +} + +local_device_id_t RealmContext::local_device_id_from_processor( + Realm::Processor processor) const { + + return this->local_machine_topology.at_l(processor); } Realm::Memory RealmContext::get_nearest_memory(Realm::Processor proc) { @@ -74,9 +178,12 @@ Allocator &RealmContext::get_current_device_allocator() { return this->allocator; } -device_id_t RealmContext::get_current_device_idx() const { +global_device_id_t RealmContext::get_current_global_device_id() const { Realm::Processor proc = this->get_current_processor(); - return this->map_processor_to_device_coord(proc); + + return global_device_id_from_local( + this->local_device_id_from_processor(proc), + node_idx_from_realm_address_space(proc.address_space())); } Realm::Event @@ -290,50 +397,12 @@ Realm::Event RealmContext::merge_outstanding_events() { return result; } -void RealmContext::discover_machine_topology() { - if (this->processors.has_value()) { - return; - } - - std::unordered_map, nonnegative_int> - next_device_idx; - - auto fresh_device_id = [&](nonnegative_int node_idx, - DeviceType device_type) -> device_id_t { - std::pair key = - std::pair{node_idx, device_type}; - if (!contains_key(next_device_idx, key)) { - next_device_idx.insert({key, 0_n}); - } - - nonnegative_int device_idx = next_device_idx.at(key); - next_device_idx.at(key)++; - - return device_id_t{ - MachineSpaceCoordinate{node_idx, device_idx}, - device_type, - }; - }; - - bidict procs; - Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); - for (Realm::Processor proc : pq) { - Realm::AddressSpace as = proc.address_space(); - Realm::Processor::Kind kind = proc.kind(); - - nonnegative_int node_idx = nonnegative_int{static_cast(as)}; - - if (kind != Realm::Processor::LOC_PROC && - kind != Realm::Processor::TOC_PROC) { - continue; - } - - DeviceType device_type = device_type_from_processor_kind(kind); - device_id_t coord = fresh_device_id(node_idx, device_type); - procs.equate_strict(proc, coord); +bidict const &RealmContext::get_global_machine_topology() { + if (!this->cached_global_machine_topology.has_value()) { + this->cached_global_machine_topology = discover_global_machine_topology(); } - this->processors = procs; + return assert_unwrap(this->cached_global_machine_topology); } Realm::Runtime RealmContext::get_runtime() { diff --git a/lib/realm-execution/src/realm-execution/tasks/impl/ff_handle_init_task.cc b/lib/realm-execution/src/realm-execution/tasks/impl/ff_handle_init_task.cc index f47b957f32..42e69a967b 100644 --- a/lib/realm-execution/src/realm-execution/tasks/impl/ff_handle_init_task.cc +++ b/lib/realm-execution/src/realm-execution/tasks/impl/ff_handle_init_task.cc @@ -39,7 +39,7 @@ void ff_handle_init_task_body(void const *args, RealmContext ctx{proc}; DeviceSpecificPtr managed_handle = make_device_specific_managed_ff_handle( - ctx.get_current_device_idx(), + ctx.get_current_global_device_id(), make_ff_handle_for_processor(proc, task_args.workSpaceSize, task_args.allowTensorOpMathConversion)); diff --git a/lib/realm-execution/src/realm-execution/tasks/impl/op_task.cc b/lib/realm-execution/src/realm-execution/tasks/impl/op_task.cc index 02b0dd7f54..9026caba96 100644 --- a/lib/realm-execution/src/realm-execution/tasks/impl/op_task.cc +++ b/lib/realm-execution/src/realm-execution/tasks/impl/op_task.cc @@ -27,7 +27,7 @@ void op_task_body(void const *args, RealmContext ctx{proc}; device_handle_t device_handle = device_handle_t_from_device_specific_managed_ff_handle( - task_args.device_handle, ctx.get_current_device_idx()); + task_args.device_handle, ctx.get_current_global_device_id()); // Patch the invocation to include the provided instances auto map_instance_to_accessor = [&](DynamicValueAttrs const &value) { @@ -53,11 +53,11 @@ void op_task_body(void const *args, /*per_device_op_state=*/ transform(and_then(task_args.device_state, [&](DeviceSpecificPtr const &d) { - return d.get(ctx.get_current_device_idx()); + return d.get(ctx.get_current_global_device_id()); }), [](PerDeviceOpState *ptr) { return *ptr; }), /*optimizer_attrs=*/task_args.optimizer_attrs, - /*device_idx=*/ctx.get_current_device_idx()); + /*global_device_id=*/ctx.get_current_global_device_id()); } Realm::Event spawn_op_task( diff --git a/lib/realm-execution/src/realm-execution/tasks/impl/per_device_op_state_init_task.cc b/lib/realm-execution/src/realm-execution/tasks/impl/per_device_op_state_init_task.cc index c05932ff61..3466890d26 100644 --- a/lib/realm-execution/src/realm-execution/tasks/impl/per_device_op_state_init_task.cc +++ b/lib/realm-execution/src/realm-execution/tasks/impl/per_device_op_state_init_task.cc @@ -31,7 +31,7 @@ void per_device_op_state_init_task_body(void const *args, RealmContext ctx{proc}; device_handle_t device_handle = device_handle_t_from_device_specific_managed_ff_handle( - task_args.device_handle, ctx.get_current_device_idx()); + task_args.device_handle, ctx.get_current_global_device_id()); // Patch the invocation to include the provided instances auto map_instance_to_accessor = [&](DynamicValueAttrs const &value) { @@ -55,16 +55,16 @@ void per_device_op_state_init_task_body(void const *args, task_args.profiling_settings, device_handle, task_args.optimizer_attrs, - ctx.get_current_device_idx()); + ctx.get_current_global_device_id()); DeviceSpecificPerDeviceOpState result_state = assert_unwrap(result_invocation.node_attrs.per_device_op_state); // Important: to make sure this doesn't get deallocated, we intentionally leak // the allocation here PerDeviceOpState *result_state_ptr = new PerDeviceOpState{get_per_device_op_state_from_device_specific( - result_state, ctx.get_current_device_idx())}; + result_state, ctx.get_current_global_device_id())}; DeviceSpecificPtr result_device_specific{ - ctx.get_current_device_idx(), result_state_ptr}; + ctx.get_current_global_device_id(), result_state_ptr}; spawn_per_device_op_state_init_return_task(ctx, task_args.origin_proc, result_device_specific, diff --git a/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc b/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc new file mode 100644 index 0000000000..75430da876 --- /dev/null +++ b/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc @@ -0,0 +1,13 @@ +#include +#include "realm-execution/fmt/realm_processor_kind.h" + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("fmt::to_string(Realm::Processor::Kind)") { + std::string result = fmt::to_string(::FlexFlow::Realm::Processor::Kind::TOC_PROC); + std::string correct = ""; + + CHECK(result == correct); + } +} 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 f698d2a07f..57871fb148 100644 --- a/lib/realm-execution/test/src/realm-execution/test_e2e.cc +++ b/lib/realm-execution/test/src/realm-execution/test_e2e.cc @@ -224,8 +224,6 @@ TEST_SUITE(FF_TEST_SUITE) { RealmManager manager = RealmManager{&fake_argc, &fake_argv}; (void)manager.start_controller([](RealmContext &ctx) { - ASSERT(ctx.processors.has_value()); - E2ETrainingConfig cfg = create_e2e_test_case(); Allocator allocator = ctx.get_current_device_allocator(); diff --git a/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc b/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc index 2ad5b54a17..7dc6a2cc5e 100644 --- a/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc +++ b/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc @@ -2,9 +2,9 @@ #include "substitutions/output_graph/output_graph_expr.h" #include "substitutions/sub_parallel_computation_graph.h" #include "utils/bidict/algorithms/bidict_from_pairs.h" -#include "utils/bidict/algorithms/merge_disjoint_bidicts.h" #include "utils/containers/values.h" #include "utils/containers/zip_values_strict.h" +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" namespace FlexFlow { @@ -26,7 +26,7 @@ bidict mapping_for_layer = bidict_from_pairs(values( zip_values_strict(layer_outputs, output_graph_expr_outputs))); - result = merge_disjoint_bidicts(result, mapping_for_layer); + result = binary_merge_disjoint_bidicts(result, mapping_for_layer); } return result; diff --git a/lib/substitutions/src/substitutions/pcg_pattern_match.cc b/lib/substitutions/src/substitutions/pcg_pattern_match.cc index 498fd6c1bf..46d284c4a0 100644 --- a/lib/substitutions/src/substitutions/pcg_pattern_match.cc +++ b/lib/substitutions/src/substitutions/pcg_pattern_match.cc @@ -5,12 +5,12 @@ #include "utils/bidict/algorithms/bidict_from_keys_and_values.h" #include "utils/bidict/algorithms/bidict_from_map.h" #include "utils/bidict/algorithms/exhaustive_relational_join.h" -#include "utils/bidict/algorithms/merge_disjoint_bidicts.h" #include "utils/bidict/algorithms/transform_values.h" #include "utils/containers/is_subseteq_of.h" #include "utils/containers/map_values.h" #include "utils/containers/values.h" #include "utils/containers/zip.h" +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" namespace FlexFlow { @@ -34,7 +34,7 @@ bidict exhaustive_relational_join(pattern_node_outputs.reversed(), matched_layer_output_tensors); - result = merge_disjoint_bidicts(result, mapping); + result = binary_merge_disjoint_bidicts(result, mapping); } return result; diff --git a/lib/task-spec/include/task-spec/device_specific.h b/lib/task-spec/include/task-spec/device_specific.h index 49a2555411..bc09cfc7a7 100644 --- a/lib/task-spec/include/task-spec/device_specific.h +++ b/lib/task-spec/include/task-spec/device_specific.h @@ -1,7 +1,7 @@ #ifndef _FLEXFLOW_LOCAL_EXECUTION_DEVICE_SPECIFIC_H #define _FLEXFLOW_LOCAL_EXECUTION_DEVICE_SPECIFIC_H -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "utils/hash/tuple.h" #include @@ -12,7 +12,7 @@ struct DeviceSpecific { DeviceSpecific() = delete; template - static DeviceSpecific create(device_id_t const &device_idx, + static DeviceSpecific create(global_device_id_t const &device_idx, Args &&...args) { return DeviceSpecific(std::make_shared(std::forward(args)...), device_idx); @@ -26,18 +26,18 @@ struct DeviceSpecific { return this->tie() != other.tie(); } - T const *get(device_id_t const &curr_device_idx) const { + T const *get(global_device_id_t const &curr_device_idx) const { ASSERT(curr_device_idx == this->device_idx); return (T const *)this->ptr.get(); } private: - DeviceSpecific(std::shared_ptr ptr, device_id_t const &device_idx) + DeviceSpecific(std::shared_ptr ptr, global_device_id_t const &device_idx) : ptr(ptr), device_idx(device_idx) {} private: std::shared_ptr ptr; - device_id_t device_idx; + global_device_id_t device_idx; private: std::tuple tie() const { diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml index 7655bd24e1..a25c9fd08b 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml @@ -13,7 +13,7 @@ includes = [ "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h", "task-spec/dynamic_graph/training_operation_attrs.dtg.h", "task-spec/device_specific_per_device_op_state.dtg.h", - "task-spec/device_id_t.dtg.h", + "task-spec/global_device_id_t.dtg.h", "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h", ] @@ -26,8 +26,8 @@ name = "task_type" type = "std::optional<::FlexFlow::DynamicTaskType>" [[fields]] -name = "device_coord" -type = "std::optional<::FlexFlow::device_id_t>" +name = "device_id" +type = "std::optional<::FlexFlow::global_device_id_t>" docstring = ''' \brief The device on which this task should execute. diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h index 6a52773a71..c5a4976db7 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h @@ -1,16 +1,16 @@ #ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" namespace FlexFlow { -bidict +bidict dynamic_node_mapping_bindings_for_slot_name(DynamicNodeMapping const &, TensorSlotName const &); -std::unordered_set +std::unordered_set target_devices_of_dynamic_node_mapping(DynamicNodeMapping const &); } // namespace FlexFlow diff --git a/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h b/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h index b40fabe7bf..a9cb86c185 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/machine_slicing.h @@ -7,11 +7,11 @@ namespace FlexFlow { std::unordered_set perform_machine_slicing_for_invocation(DynamicNodeInvocation const &, - device_id_t const &); + global_device_id_t const &); DynamicOpenDataflowGraph perform_machine_slicing(DynamicOpenDataflowGraph const &, - device_id_t const &); + global_device_id_t const &); } // namespace FlexFlow diff --git a/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml index ba9785bc92..2be195f7c1 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/parallel_tensor_mapping.dtg.toml @@ -11,7 +11,7 @@ features = [ includes = [ "utils/bidict/bidict.h", "op-attrs/parallel_tensor_space_coordinate.dtg.h", - "task-spec/device_id_t.dtg.h", + "task-spec/global_device_id_t.dtg.h", ] src_includes = [ @@ -19,4 +19,4 @@ src_includes = [ [[fields]] name = "raw" -type = "::FlexFlow::bidict<::FlexFlow::ParallelTensorSpaceCoordinate, ::FlexFlow::device_id_t>" +type = "::FlexFlow::bidict<::FlexFlow::ParallelTensorSpaceCoordinate, ::FlexFlow::global_device_id_t>" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml index 43fb8490ec..6638350bfd 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml +++ b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_node_attrs.dtg.toml @@ -11,7 +11,7 @@ features = [ includes = [ "", "task-spec/dynamic_graph/dynamic_task_type.dtg.h", - "task-spec/device_id_t.dtg.h", + "task-spec/global_device_id_t.dtg.h", "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h", "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h", "task-spec/dynamic_graph/training_operation_attrs.dtg.h", @@ -27,8 +27,8 @@ name = "task_type" type = "std::optional<::FlexFlow::DynamicTaskType>" [[fields]] -name = "device_coord" -type = "std::optional<::FlexFlow::device_id_t>" +name = "device_id" +type = "std::optional<::FlexFlow::global_device_id_t>" [[fields]] name = "mapping" diff --git a/lib/task-spec/include/task-spec/device_id_t.dtg.toml b/lib/task-spec/include/task-spec/global_device_id_t.dtg.toml similarity index 91% rename from lib/task-spec/include/task-spec/device_id_t.dtg.toml rename to lib/task-spec/include/task-spec/global_device_id_t.dtg.toml index f6d29ee1a7..5644c33c17 100644 --- a/lib/task-spec/include/task-spec/device_id_t.dtg.toml +++ b/lib/task-spec/include/task-spec/global_device_id_t.dtg.toml @@ -1,5 +1,5 @@ namespace = "FlexFlow" -name = "device_id_t" +name = "global_device_id_t" type = "struct" features = [ "eq", diff --git a/lib/task-spec/include/task-spec/global_device_id_t.h b/lib/task-spec/include/task-spec/global_device_id_t.h new file mode 100644 index 0000000000..2f48046ff2 --- /dev/null +++ b/lib/task-spec/include/task-spec/global_device_id_t.h @@ -0,0 +1,15 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_GLOBAL_DEVICE_ID_T_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_GLOBAL_DEVICE_ID_T_H + +#include "task-spec/global_device_id_t.dtg.h" +#include "task-spec/local_device_id_t.dtg.h" +#include "pcg/node_idx_t.dtg.h" + +namespace FlexFlow { + +global_device_id_t global_device_id_from_local(local_device_id_t const &, node_idx_t); +local_device_id_t local_device_id_from_global(global_device_id_t const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/local_device_id_t.dtg.toml b/lib/task-spec/include/task-spec/local_device_id_t.dtg.toml new file mode 100644 index 0000000000..cb6ae47621 --- /dev/null +++ b/lib/task-spec/include/task-spec/local_device_id_t.dtg.toml @@ -0,0 +1,23 @@ +namespace = "FlexFlow" +name = "local_device_id_t" +type = "struct" +features = [ + "eq", + "ord", + "hash", + "json", + "fmt", +] + +includes = [ + "pcg/device_in_node_idx_t.dtg.h", + "pcg/device_type.dtg.h", +] + +[[fields]] +name = "idx" +type = "::FlexFlow::device_in_node_idx_t" + +[[fields]] +name = "device_type" +type = "::FlexFlow::DeviceType" diff --git a/lib/task-spec/include/task-spec/per_device_op_state.h b/lib/task-spec/include/task-spec/per_device_op_state.h index 8783f902e4..3d9f90c0a6 100644 --- a/lib/task-spec/include/task-spec/per_device_op_state.h +++ b/lib/task-spec/include/task-spec/per_device_op_state.h @@ -9,7 +9,7 @@ namespace FlexFlow { PerDeviceOpState get_per_device_op_state_from_device_specific( - DeviceSpecificPerDeviceOpState const &, device_id_t device_idx); + DeviceSpecificPerDeviceOpState const &, global_device_id_t device_idx); } diff --git a/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h b/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h index 165631889f..776157e644 100644 --- a/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h +++ b/lib/task-spec/include/task-spec/task_argument_accessor/itask_argument_accessor.h @@ -9,7 +9,7 @@ #include "op-attrs/tensor_slot_name.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/concrete_arg_spec.h" -#include "task-spec/device_id_t.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/ops/arg_slot_id_t.dtg.h" #include "task-spec/per_device_op_state.dtg.h" #include "task-spec/privilege_tensor_accessor.h" @@ -37,7 +37,7 @@ struct ITaskArgumentAccessor { virtual OptimizerAttrs get_optimizer_attrs() const = 0; virtual Allocator get_allocator() const = 0; - virtual device_id_t get_device_idx() const = 0; + virtual global_device_id_t get_device_idx() const = 0; }; CHECK_RC_COPY_VIRTUAL_COMPLIANT(ITaskArgumentAccessor); diff --git a/lib/task-spec/include/task-spec/task_argument_accessor/task_argument_accessor.h b/lib/task-spec/include/task-spec/task_argument_accessor/task_argument_accessor.h index 0c63643e6e..9ad56bf69d 100644 --- a/lib/task-spec/include/task-spec/task_argument_accessor/task_argument_accessor.h +++ b/lib/task-spec/include/task-spec/task_argument_accessor/task_argument_accessor.h @@ -58,7 +58,7 @@ struct TaskArgumentAccessor { return this->ptr->get_allocator(); } - device_id_t get_device_idx() const { + global_device_id_t get_device_idx() const { return this->ptr->get_device_idx(); } 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 f611735e26..3487b2979e 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 @@ -61,16 +61,16 @@ bool graph_is_fully_copy_inserted(DynamicOpenDataflowGraph const &g) { static std::pair filter_mapping_to_avoid_degenerate_copies(DynamicValueAttrs const &input, DynamicValueAttrs const &output) { - std::unordered_set> + std::unordered_set> input_mapping = unordered_set_of(assert_unwrap(input.mapping).raw); - std::unordered_set> + std::unordered_set> output_mapping = unordered_set_of(assert_unwrap(output.mapping).raw); // Exclude the point shared between the input and output mappings, because // those will not result in actual copies once shard expansion is performed std::unordered_set< - std::pair> + std::pair> remove = set_intersection(input_mapping, output_mapping); DynamicValueAttrs filtered_input = input; diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc index 9e56dddcf9..4e21da9c2a 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc @@ -4,24 +4,24 @@ namespace FlexFlow { -bidict +bidict dynamic_node_mapping_bindings_for_slot_name( DynamicNodeMapping const &mapping, TensorSlotName const &slot_name) { bidict coord_bindings = get_tensor_bindings_for_slot_name(mapping.op_task_group, slot_name); return transform_values( - coord_bindings, [&](MachineSpaceCoordinate const &coord) -> device_id_t { - return device_id_t{coord, mapping.device_type}; + coord_bindings, [&](MachineSpaceCoordinate const &coord) -> global_device_id_t { + return global_device_id_t{coord, mapping.device_type}; }); } -std::unordered_set +std::unordered_set target_devices_of_dynamic_node_mapping(DynamicNodeMapping const &mapping) { return transform(mapping.op_task_group.get_shard_bindings().left_values(), - [&](MachineSpaceCoordinate const &c) -> device_id_t { - return device_id_t{ + [&](MachineSpaceCoordinate const &c) -> global_device_id_t { + return global_device_id_t{ /*coord=*/c, /*device_type=*/mapping.device_type, }; diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc index 07c49664a5..edb70c2f3a 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc @@ -59,6 +59,15 @@ void check_dynamic_open_dataflow_graph_is_valid( ASSERT(values_produced_multiple_times.size() == 0, keys(values_produced_multiple_times)); + // since DynamicOpenDataflowGraph contains a set of invocations rather than a + // LabelledOpenKwargDataflowGraph, some properties guaranteed by + // LabelledOpenKwargDataflowGraph (e.g., the graph is acyclic, all tensors + // originate from another operator's output unless they're a designated graph + // input, etc.) are not automatically guaranteed. Since + // LabelledOpenKwargDataflowGraph guarantees these properties, the easiest + // way to check them is to try to convert the DynamicOpenDataflowGraph into a + // LabelledOpenKwargDataflowGraph, and if a value is returned without an + // assertion we know the properties hold. labelled_open_kwarg_dataflow_graph_from_dynamic_open_dataflow_graph(g); } diff --git a/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc b/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc index e46f8b1510..4871a3ee4e 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/machine_slicing.cc @@ -6,11 +6,11 @@ namespace FlexFlow { std::unordered_set perform_machine_slicing_for_invocation( DynamicNodeInvocation const &invocation, - device_id_t const &device_coord) { + global_device_id_t const &device_id) { - ASSERT(invocation.node_attrs.device_coord.has_value()); + ASSERT(invocation.node_attrs.device_id.has_value()); - if (invocation.node_attrs.device_coord.value() == device_coord) { + if (invocation.node_attrs.device_id.value() == device_id) { return {invocation}; } else { return {}; @@ -19,12 +19,12 @@ std::unordered_set DynamicOpenDataflowGraph perform_machine_slicing(DynamicOpenDataflowGraph const &g, - device_id_t const &device_coord) { + global_device_id_t const &device_id) { DynamicOpenDataflowGraph result = flatmap_dynamic_invocation_set( g, [&](DynamicNodeInvocation const &invocation) -> std::unordered_set { - return perform_machine_slicing_for_invocation(invocation, device_coord); + return perform_machine_slicing_for_invocation(invocation, device_id); }); return result; diff --git a/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_node_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_node_attrs.cc index d613194d14..7ad4686c73 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_node_attrs.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_node_attrs.cc @@ -7,7 +7,7 @@ SerializableDynamicNodeAttrs dynamic_node_attrs_to_serializable(DynamicNodeAttrs const &attrs) { return SerializableDynamicNodeAttrs{ /*task_type=*/attrs.task_type, - /*device_coord=*/attrs.device_coord, + /*device_id=*/attrs.device_id, /*mapping=*/attrs.mapping, /*op_attrs=*/attrs.op_attrs, /*layer_guid=*/attrs.layer_guid, @@ -18,7 +18,7 @@ DynamicNodeAttrs dynamic_node_attrs_from_serializable( SerializableDynamicNodeAttrs const &attrs) { return DynamicNodeAttrs{ /*task_type=*/attrs.task_type, - /*device_coord=*/attrs.device_coord, + /*device_id=*/attrs.device_id, /*mapping=*/attrs.mapping, /*op_attrs=*/attrs.op_attrs, /*layer_guid=*/attrs.layer_guid, 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 117819a639..15f78944c8 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 @@ -13,7 +13,7 @@ namespace FlexFlow { bool node_is_shard_expanded(DynamicNodeAttrs const &n) { - return n.device_coord.has_value(); + return n.device_id.has_value(); } bool value_is_shard_expanded(DynamicValueAttrs const &n) { @@ -42,9 +42,9 @@ bool graph_is_fully_shard_expanded(DynamicOpenDataflowGraph const &g) { slot_is_shard_expanded); } -static bidict +static bidict restrict_tensor_mapping_keys_to_coord( - bidict const &mapping, + bidict const &mapping, ParallelTensorSpaceCoordinate const ¶llel_tensor_coord) { return filter_keys(mapping, [&](ParallelTensorSpaceCoordinate const &p) { return p == parallel_tensor_coord; @@ -53,7 +53,7 @@ static bidict static DynamicNodeInvocation shard_invocation_for_binding( DynamicNodeInvocation const &i, - device_id_t const &device_coord, + global_device_id_t const &device_id, OperatorAtomicTaskShardBinding const &binding) { auto shard_expand_value_attrs = [&](DynamicTensorSlot const &s, @@ -76,7 +76,7 @@ static DynamicNodeInvocation shard_invocation_for_binding( DynamicNodeAttrs expanded_node_attrs = [&]() { DynamicNodeAttrs result = i.node_attrs; - result.device_coord = device_coord; + result.device_id = device_id; return result; }(); @@ -91,7 +91,7 @@ static std::unordered_set perform_shard_expansion_for_copy(DynamicNodeInvocation const &i) { auto [input_slot, input] = get_only(i.inputs); auto [output_slot, output] = get_only(i.outputs); - bidict input_mapping = + bidict input_mapping = assert_unwrap(input.mapping).raw; require_same(input_mapping.left_values(), assert_unwrap(output.mapping).raw.left_values()); @@ -105,7 +105,7 @@ static std::unordered_set // because we expect this to align with the most efficient way to issue // copies in Realm, although the current Realm backend uses a // centralized controller and thus issues copies all from a single node. - device_id_t machine_coord = input_mapping.at_l(p); + global_device_id_t machine_coord = input_mapping.at_l(p); return shard_invocation_for_binding(i, machine_coord, @@ -125,11 +125,11 @@ std::unordered_set DynamicNodeMapping mapping = assert_unwrap(i.node_attrs.mapping); - std::unordered_set shard_machine_coords = + std::unordered_set shard_machine_coords = target_devices_of_dynamic_node_mapping(mapping); return transform( - shard_machine_coords, [&](device_id_t const &c) -> DynamicNodeInvocation { + shard_machine_coords, [&](global_device_id_t const &c) -> DynamicNodeInvocation { OperatorAtomicTaskShardBinding slot_bindings = mapping.op_task_group.get_shard_bindings().at_l(c.coord); diff --git a/lib/task-spec/src/task-spec/global_device_id_t.cc b/lib/task-spec/src/task-spec/global_device_id_t.cc new file mode 100644 index 0000000000..e32a5f5f8b --- /dev/null +++ b/lib/task-spec/src/task-spec/global_device_id_t.cc @@ -0,0 +1,25 @@ +#include "task-spec/global_device_id_t.h" + +namespace FlexFlow { + +global_device_id_t global_device_id_from_local( + local_device_id_t const &local_device_id, + node_idx_t node_idx) +{ + return global_device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/node_idx.raw, + /*device_idx=*/local_device_id.idx.raw, + }, + /*device_type=*/local_device_id.device_type, + }; +} + +local_device_id_t local_device_id_from_global(global_device_id_t const &global_device_id) { + return local_device_id_t{ + /*idx=*/device_in_node_idx_t{global_device_id.coord.device_idx}, + /*device_type=*/global_device_id.device_type, + }; +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/per_device_op_state.cc b/lib/task-spec/src/task-spec/per_device_op_state.cc index 438cd8886c..db4c3f95d8 100644 --- a/lib/task-spec/src/task-spec/per_device_op_state.cc +++ b/lib/task-spec/src/task-spec/per_device_op_state.cc @@ -5,7 +5,7 @@ namespace FlexFlow { PerDeviceOpState get_per_device_op_state_from_device_specific( DeviceSpecificPerDeviceOpState const &device_specific, - device_id_t device_idx) { + global_device_id_t device_idx) { return device_specific.visit( [&](auto const &x) { return PerDeviceOpState{*(x.get(device_idx))}; }); } diff --git a/lib/task-spec/test/src/task-spec/device_specific.cc b/lib/task-spec/test/src/task-spec/device_specific.cc index 6a42a9b570..983ba4943d 100644 --- a/lib/task-spec/test/src/task-spec/device_specific.cc +++ b/lib/task-spec/test/src/task-spec/device_specific.cc @@ -7,12 +7,12 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("DeviceSpecific") { DeviceSpecific device_specific1 = DeviceSpecific::create( - device_id_t{MachineSpaceCoordinate{0_n, 0_n}, DeviceType::GPU}, + global_device_id_t{MachineSpaceCoordinate{0_n, 0_n}, DeviceType::GPU}, "hello world"); DeviceSpecific device_specific2 = DeviceSpecific::create( - device_id_t{MachineSpaceCoordinate{0_n, 1_n}, DeviceType::GPU}, + global_device_id_t{MachineSpaceCoordinate{0_n, 1_n}, DeviceType::GPU}, "hello world"); std::string result1 = fmt::to_string(device_specific1); 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 58a8a36fcd..97304ebbf0 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 @@ -65,8 +65,8 @@ TEST_SUITE(FF_TEST_SUITE) { }; }; - auto mk_device_id = [&](nonnegative_int device_idx) -> device_id_t { - return device_id_t{ + auto mk_device_id = [&](nonnegative_int device_idx) -> global_device_id_t { + return global_device_id_t{ mk_machine_coord(device_idx), DeviceType::GPU, }; @@ -209,7 +209,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(1_n)}, }, @@ -220,7 +220,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(1_n)}, }, @@ -231,7 +231,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(2_n)}, }, @@ -242,7 +242,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(2_n)}, }, @@ -399,7 +399,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(1_n)}, }, @@ -410,7 +410,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(1_n)}, }, @@ -421,7 +421,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*src_slot=*/TensorSlotName::OUTPUT, /*mapping=*/ ParallelTensorMapping{ - bidict{ + bidict{ {mk_ptensor_coord(0_n), mk_device_id(0_n)}, {mk_ptensor_coord(1_n), mk_device_id(1_n)}, }, diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc index 6bc31999a6..9abb802af2 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc @@ -7,8 +7,8 @@ using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("perform_machine_slicing_for_invocation") { auto mk_device_id = [](nonnegative_int node_idx, - nonnegative_int device_idx) -> device_id_t { - return device_id_t{ + nonnegative_int device_idx) -> global_device_id_t { + return global_device_id_t{ MachineSpaceCoordinate{ /*node_idx=*/node_idx, /*device_idx=*/device_idx, @@ -33,9 +33,9 @@ TEST_SUITE(FF_TEST_SUITE) { }; }; - device_id_t mc1 = mk_device_id(0_n, 0_n); - device_id_t mc2 = mk_device_id(2_n, 0_n); - device_id_t mc3 = mk_device_id(4_n, 0_n); + global_device_id_t mc1 = mk_device_id(0_n, 0_n); + global_device_id_t mc2 = mk_device_id(2_n, 0_n); + global_device_id_t mc3 = mk_device_id(4_n, 0_n); ParallelTensorSpaceCoordinate mc1_input_coord = mk_pt_coord(0_n, 0_n, 0_n, 0_n); diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc index 65870192d7..cc51c838b9 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc @@ -44,14 +44,14 @@ TEST_SUITE(FF_TEST_SUITE) { }; DeviceType device_type = DeviceType::GPU; - auto mk_device_id = [&](MachineSpaceCoordinate const &c) -> device_id_t { - return device_id_t{c, device_type}; + auto mk_device_id = [&](MachineSpaceCoordinate const &c) -> global_device_id_t { + return global_device_id_t{c, device_type}; }; auto mk_value = [&](size_t src_node_id, TensorSlotName src_slot_name, - bidict tensor_binding, + bidict tensor_binding, std::optional const &shard_coord) -> DynamicValueAttrs { if (shard_coord.has_value()) { @@ -155,7 +155,7 @@ TEST_SUITE(FF_TEST_SUITE) { TensorSlotName use_slot_name, std::optional const &shard_coord) -> DynamicValueAttrs { - bidict tensor_binding = + bidict tensor_binding = dynamic_node_mapping_bindings_for_slot_name(node_mapping, use_slot_name); return mk_value( @@ -212,7 +212,7 @@ TEST_SUITE(FF_TEST_SUITE) { perform_shard_expansion_for_invocation(input); auto mk_invocation_shard = - [&](device_id_t const &device_coord, + [&](global_device_id_t const &device_coord, ParallelTensorSpaceCoordinate const &input_shard_coord, ParallelTensorSpaceCoordinate const &weight_shard_coord, ParallelTensorSpaceCoordinate const &output_1_shard_coord, @@ -283,19 +283,19 @@ TEST_SUITE(FF_TEST_SUITE) { } SUBCASE("for copy operator") { - device_id_t dev1 = mk_device_id(mk_machine_coord(0_n, 0_n)); - device_id_t dev2 = mk_device_id(mk_machine_coord(1_n, 0_n)); - device_id_t dev3 = mk_device_id(mk_machine_coord(2_n, 0_n)); - device_id_t dev4 = mk_device_id(mk_machine_coord(3_n, 0_n)); + global_device_id_t dev1 = mk_device_id(mk_machine_coord(0_n, 0_n)); + global_device_id_t dev2 = mk_device_id(mk_machine_coord(1_n, 0_n)); + global_device_id_t dev3 = mk_device_id(mk_machine_coord(2_n, 0_n)); + global_device_id_t dev4 = mk_device_id(mk_machine_coord(3_n, 0_n)); ParallelTensorSpaceCoordinate pt1 = mk_pt_coord(0_n, 0_n, 0_n, 0_n); ParallelTensorSpaceCoordinate pt2 = mk_pt_coord(0_n, 1_n, 0_n, 0_n); - bidict src_binding{ + bidict src_binding{ {pt1, dev1}, {pt2, dev2}, }; - bidict dst_binding{ + bidict dst_binding{ {pt1, dev3}, {pt2, dev4}, }; @@ -331,7 +331,7 @@ TEST_SUITE(FF_TEST_SUITE) { perform_shard_expansion_for_invocation(input); auto mk_invocation_shard = - [&](device_id_t const &device_coord, + [&](global_device_id_t const &device_id, ParallelTensorSpaceCoordinate const &tensor_shard_coord) -> DynamicNodeInvocation { DynamicNodeInvocation result = input; @@ -343,7 +343,7 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; // See perform_shard_expansion_for_copy in shard_expansion.cc for explanation of the choice of device placement. - result.node_attrs.device_coord = device_coord; + result.node_attrs.device_id = device_id; result.outputs = { { mk_slot(TensorSlotName::OUTPUT), diff --git a/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h b/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h new file mode 100644 index 0000000000..607cb08229 --- /dev/null +++ b/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h @@ -0,0 +1,38 @@ +#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_BINARY_MERGE_DISJOINT_BIDICTS_H +#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_BINARY_MERGE_DISJOINT_BIDICTS_H + +#include "utils/bidict/bidict.h" +#include +#include "utils/containers/are_disjoint.h" +#include "utils/bidict/algorithms/right_entries.h" +#include "utils/bidict/algorithms/left_entries.h" + +namespace FlexFlow { + +template +bidict binary_merge_disjoint_bidicts(bidict const &lhs, + bidict const &rhs) { + ASSERT( + are_disjoint(left_entries(lhs), left_entries(rhs)), + "Left entries of {} and {} are non-disjoint", lhs, rhs + ); + + ASSERT( + are_disjoint(right_entries(lhs), right_entries(rhs)), + "Right entries of {} and {} are non-disjoint", lhs, rhs + ); + + bidict result; + for (auto const &kv : lhs) { + result.equate_strict(kv.first, kv.second); + } + for (auto const &kv : rhs) { + result.equate_strict(kv.first, kv.second); + } + + return result; +} + +} // namespace FlexFlow + +#endif diff --git a/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h b/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h index 97e7334c26..1ce7cb3a13 100644 --- a/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h +++ b/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h @@ -1,35 +1,20 @@ #ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_DISJOINT_BIDICTS_H #define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_DISJOINT_BIDICTS_H -#include "utils/bidict/algorithms/left_entries.h" -#include "utils/bidict/algorithms/right_entries.h" #include "utils/bidict/bidict.h" -#include "utils/containers/are_disjoint.h" -#include "utils/exception.h" +#include "utils/containers/foldl.h" +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" namespace FlexFlow { template -bidict merge_disjoint_bidicts(bidict const &lhs, - bidict const &rhs) { - if (!are_disjoint(left_entries(lhs), left_entries(rhs))) { - throw mk_runtime_error( - fmt::format("Left entries of {} and {} are non-disjoint", lhs, rhs)); - } - if (!are_disjoint(right_entries(lhs), right_entries(rhs))) { - throw mk_runtime_error( - fmt::format("Right entries of {} and {} are non-disjoint", lhs, rhs)); - } - - bidict result; - for (auto const &kv : lhs) { - result.equate(kv.first, kv.second); - } - for (auto const &kv : rhs) { - result.equate(kv.first, kv.second); - } - - return result; +bidict merge_disjoint_bidicts(std::set> const &bidicts) { + return foldl( + bidicts, + bidict{}, + [](bidict const &accum, bidict const &x) -> bidict { + return binary_merge_disjoint_bidicts(accum, x); + }); } } // namespace FlexFlow diff --git a/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc b/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc new file mode 100644 index 0000000000..d3deb887b1 --- /dev/null +++ b/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc @@ -0,0 +1,13 @@ +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" +#include "utils/archetypes/ordered_value_type.h" + +namespace FlexFlow { + +using K = ordered_value_type<0>; +using V = ordered_value_type<1>; + +template + bidict binary_merge_disjoint_bidicts(bidict const &, + bidict const &); + +} // namespace FlexFlow diff --git a/lib/utils/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc b/lib/utils/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc index 754b8d2e90..e5ed21bb40 100644 --- a/lib/utils/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc +++ b/lib/utils/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc @@ -1 +1,11 @@ #include "utils/bidict/algorithms/merge_disjoint_bidicts.h" +#include "utils/archetypes/ordered_value_type.h" + +namespace FlexFlow { + +using K = ordered_value_type<0>; +using V = ordered_value_type<1>; + +template bidict merge_disjoint_bidicts(std::set> const &); + +} // namespace FlexFlow diff --git a/lib/utils/test/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc b/lib/utils/test/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc similarity index 72% rename from lib/utils/test/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc rename to lib/utils/test/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc index 0a1babd9f9..8a3371b8d8 100644 --- a/lib/utils/test/src/utils/bidict/algorithms/merge_disjoint_bidicts.cc +++ b/lib/utils/test/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc @@ -1,17 +1,17 @@ -#include "utils/bidict/algorithms/merge_disjoint_bidicts.h" +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" #include using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { - TEST_CASE("merge_disjoint_bidicts") { + TEST_CASE("binary_merge_disjoint_bidicts") { SUBCASE("disjoint keys and values") { bidict bd1 = {{1, "one"}, {2, "two"}}; bidict bd2 = {{3, "three"}, {4, "four"}}; - bidict result = merge_disjoint_bidicts(bd1, bd2); + bidict result = binary_merge_disjoint_bidicts(bd1, bd2); bidict correct = { {1, "one"}, {2, "two"}, {3, "three"}, {4, "four"}}; @@ -22,21 +22,21 @@ TEST_SUITE(FF_TEST_SUITE) { bidict bd1 = {{1, "one"}, {2, "two"}}; bidict bd2 = {{2, "three"}, {3, "four"}}; - CHECK_THROWS(merge_disjoint_bidicts(bd1, bd2)); + CHECK_THROWS(binary_merge_disjoint_bidicts(bd1, bd2)); } SUBCASE("overlapping key, same associated value") { bidict bd1 = {{1, "one"}, {2, "two"}}; bidict bd2 = {{2, "two"}, {3, "three"}}; - CHECK_THROWS(merge_disjoint_bidicts(bd1, bd2)); + CHECK_THROWS(binary_merge_disjoint_bidicts(bd1, bd2)); } SUBCASE("overlapping values") { bidict bd1 = {{1, "one"}, {2, "two"}}; bidict bd2 = {{3, "two"}, {4, "four"}}; - CHECK_THROWS(merge_disjoint_bidicts(bd1, bd2)); + CHECK_THROWS(binary_merge_disjoint_bidicts(bd1, bd2)); } } } From 9ecf61a3581b1e8d7e42c71006bcba2db9b9ba5e Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Fri, 12 Jun 2026 20:55:50 -0700 Subject: [PATCH 4/5] Format and fix checks --- .../unstructured_device_mapping.h | 18 --- .../machine_mapping/allowed_machine_views.cc | 4 +- .../machine_mapping_mutation_set.cc | 4 +- .../compiler/machine_mapping/machine_view.cc | 4 +- .../start_invariant_machine_view.cc | 3 +- .../src/compiler/mcmc/mcmc_over_mapped_pcg.cc | 8 +- .../task_graph_simulator/pcg_task_graph.cc | 5 +- .../task_graph_simulator/task_simulator.cc | 6 +- .../unity_algorithm/unity_algorithm.cc | 3 +- .../machine_mapping/allowed_machine_views.cc | 2 +- .../compiler/machine_mapping/machine_view.cc | 47 ++++--- .../start_invariant_machine_view.cc | 33 ++--- .../src/compiler/mcmc/mcmc_over_mapped_pcg.cc | 2 +- .../task_graph_simulator/task_simulator.cc | 21 +-- .../unity_algorithm/graph_optimize_state.cc | 1 - .../unity_algorithm/unity_algorithm.cc | 2 +- .../computation_graph_instance.h | 2 +- .../cost_estimator/local_cost_estimator.h | 2 +- .../local_task_argument_accessor.h | 2 +- .../per_device_op_state_initialization.h | 2 +- .../cost_estimator/local_cost_estimator.cc | 4 +- .../computation_graph_instance.cc | 31 +++-- .../cost_estimator/local_cost_estimator.cc | 34 +++-- .../local_task_argument_accessor.cc | 10 +- ...ce_specific_managed_per_device_ff_handle.h | 3 +- .../realm-execution/device_specific_ptr.h | 3 +- .../realm-execution/fmt/realm_processor.h | 11 +- .../fmt/realm_processor_kind.h | 10 +- .../include/realm-execution/pcg_instance.h | 2 +- .../include/realm-execution/processor_query.h | 3 +- .../include/realm-execution/realm_context.h | 11 +- .../src/realm-execution/address_space.cc | 9 +- ...uted_per_device_op_state_initialization.cc | 4 +- .../realm-execution/fmt/realm_processor.cc | 3 +- .../fmt/realm_processor_kind.cc | 3 +- .../src/realm-execution/processor_query.cc | 3 +- .../src/realm-execution/realm_context.cc | 127 +++++++++--------- .../fmt/realm_processor_kind.cc | 5 +- .../output_expr_to_result_sub_pcg_mapping.cc | 2 +- .../src/substitutions/pcg_pattern_match.cc | 2 +- .../dynamic_graph/dynamic_node_mapping.h | 2 +- .../include/task-spec/global_device_id_t.h | 5 +- .../task-spec/dynamic_graph/copy_insertion.cc | 9 +- .../dynamic_graph/dynamic_node_mapping.cc | 3 +- .../dynamic_graph/shard_expansion.cc | 16 ++- .../src/task-spec/global_device_id_t.cc | 24 ++-- .../test/src/task-spec/device_specific.cc | 6 +- .../dynamic_graph/shard_expansion.cc | 12 +- .../binary_merge_disjoint_bidicts.h | 24 ++-- .../algorithms/merge_disjoint_bidicts.h | 12 +- lib/utils/include/utils/variant.h | 2 +- .../binary_merge_disjoint_bidicts.cc | 5 +- .../include/test/utils/rapidcheck/doctest.h | 2 +- .../non_normal_sp_decomposition.cc | 2 +- 54 files changed, 274 insertions(+), 301 deletions(-) delete mode 100644 lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.h diff --git a/lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.h b/lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.h deleted file mode 100644 index 8c1333fabc..0000000000 --- a/lib/compiler/include/compiler/machine_mapping/unstructured_device_mapping.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef _FLEXFLOW_COMPILER_MACHINE_MAPPING_UNSTRUCTURED_DEVICE_MAPPING_H -#define _FLEXFLOW_COMPILER_MACHINE_MAPPING_UNSTRUCTURED_DEVICE_MAPPING_H - -#include "compiler/machine_mapping/machine_mapping.dtg.h" -#include "compiler/machine_mapping/unstructured_device_mapping.dtg.h" -#include "pcg/machine_compute_specification.dtg.h" -#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h" - -namespace FlexFlow { - -UnstructuredDeviceMapping get_unstructured_device_mapping( - MachineMapping const &machine_mapping, - MachineComputeSpecification const &machine_spec, - ParallelComputationGraph const &pcg); - -} // namespace FlexFlow - -#endif diff --git a/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc b/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc index 9194c2e982..4cd4abb056 100644 --- a/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc +++ b/lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc @@ -91,13 +91,11 @@ static std::unordered_set auto get_candidate_starts = [](MachineComputeResourceSlice const &slice) -> std::unordered_set { - std::unordered_set result; for (nonnegative_int node_idx : nonnegative_range(slice.num_nodes)) { for (nonnegative_int device_idx : nonnegative_range(slice.num_gpus_per_node)) { - result.insert( - MachineSpaceCoordinate{node_idx, device_idx}); + result.insert(MachineSpaceCoordinate{node_idx, device_idx}); } } return result; diff --git a/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc b/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc index d6cdca97d1..41fbce8a78 100644 --- a/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc +++ b/lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc @@ -17,8 +17,8 @@ std::optional for (parallel_layer_guid_t layer : layers) { OperatorTaskSpace task = get_operator_task_space(pcg, layer); std::unordered_set allowed_machine_views = - get_allowed_machine_views( - compute_slice_from_specification(resources), task); + get_allowed_machine_views(compute_slice_from_specification(resources), + task); if (allowed_machine_views.empty()) { return std::nullopt; } diff --git a/lib/compiler/src/compiler/machine_mapping/machine_view.cc b/lib/compiler/src/compiler/machine_mapping/machine_view.cc index 5c38a66901..baf40faf12 100644 --- a/lib/compiler/src/compiler/machine_mapping/machine_view.cc +++ b/lib/compiler/src/compiler/machine_mapping/machine_view.cc @@ -125,8 +125,8 @@ MachineSpaceCoordinate compute_index(machine_view.start.node_idx, inter_dimension_indices); nonnegative_int device_idx = compute_index(machine_view.start.device_idx, intra_dimension_indices); - MachineSpaceCoordinate ms_coord = MachineSpaceCoordinate{ - node_idx, device_idx}; + MachineSpaceCoordinate ms_coord = + MachineSpaceCoordinate{node_idx, device_idx}; return ms_coord; } diff --git a/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc b/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc index 4a2d66acc1..cd9309d886 100644 --- a/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc +++ b/lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc @@ -54,8 +54,7 @@ MachineSpaceOffset get_machine_space_offset( StartInvariantMachineView const &start_inv_machine_view, TaskSpaceCoordinate const &coord) { - MachineSpaceCoordinate dummy_start = - MachineSpaceCoordinate{0_n, 0_n}; + MachineSpaceCoordinate dummy_start = MachineSpaceCoordinate{0_n, 0_n}; MachineView mv = machine_view_from_start_invariant(start_inv_machine_view, dummy_start); diff --git a/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc b/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc index 0d2c1e4455..0dd1607d9b 100644 --- a/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc +++ b/lib/compiler/src/compiler/mcmc/mcmc_over_mapped_pcg.cc @@ -21,8 +21,8 @@ SearchResult MCMCOverMappedPCGConfig const &search_config) { MachineComputeSpecification compute_spec = machine_spec.compute_specification; std::vector substitutions = get_substitution_set(compute_spec); - MachineMapping random_mapping = assert_unwrap( - get_random_mapping(pcg, compute_spec)); + MachineMapping random_mapping = + assert_unwrap(get_random_mapping(pcg, compute_spec)); SearchResult starting_state = SearchResult{pcg, random_mapping}; auto sampler = [&](SearchResult mapped_pcg) -> std::optional { @@ -42,8 +42,8 @@ SearchResult mapped_pcg, random_substitution, match); }); } else { - MachineMapping new_machine_mapping = assert_unwrap(get_random_mutation( - mapped_pcg, compute_spec)); + MachineMapping new_machine_mapping = + assert_unwrap(get_random_mutation(mapped_pcg, compute_spec)); return SearchResult{mapped_pcg.pcg, new_machine_mapping}; } }; diff --git a/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc b/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc index 058f5b72bb..d3666a2451 100644 --- a/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc +++ b/lib/compiler/src/compiler/task_graph_simulator/pcg_task_graph.cc @@ -24,7 +24,8 @@ PCGTaskGraph DiGraph digraph = DiGraph::create(); bidict node_to_task; bidict node_to_layer; - std::unordered_map> node_to_devices; + std::unordered_map> + node_to_devices; for (parallel_layer_guid_t const &layer : get_parallel_layers(pcg)) { MachineView mv = machine_mapping.machine_views.at(layer); @@ -35,7 +36,7 @@ PCGTaskGraph node_to_layer.equate(node, layer); node_to_devices[node] = get_machine_space_coordinates(get_operator_task_space(pcg, layer), - machine_mapping.machine_views.at(layer)); + machine_mapping.machine_views.at(layer)); } for (ParallelComputationGraphEdge const &edge : get_edges(pcg)) { diff --git a/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc b/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc index 28a6a3efae..843470a107 100644 --- a/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc +++ b/lib/compiler/src/compiler/task_graph_simulator/task_simulator.cc @@ -52,13 +52,15 @@ milliseconds_t task_simulator_estimate_forward_pass_time( } assert(current_task.is_operator()); - auto get_devices = [&](Node const &n) -> std::unordered_set { + auto get_devices = + [&](Node const &n) -> std::unordered_set { return task_graph.node_to_devices.at(n); }; std::unordered_set devices_occupied = set_union(transform(in_progress_tasks, get_devices)); - std::unordered_set required_devices = get_devices(task); + std::unordered_set required_devices = + get_devices(task); return set_intersection(devices_occupied, required_devices).empty(); }; diff --git a/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc b/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc index 240ad0b1c4..dd95f9ea26 100644 --- a/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc +++ b/lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc @@ -67,8 +67,7 @@ SearchResult graph_optimize(ParallelComputationGraph &pcg, OperatorTaskSpace op_task_space = get_operator_task_space_for_runtime_only_op_cost_estimate_key(key); - return get_allowed_machine_views( - resources, op_task_space); + return get_allowed_machine_views(resources, op_task_space); }, }; diff --git a/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc b/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc index 6a867b16f3..c465f88c51 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/allowed_machine_views.cc @@ -1,11 +1,11 @@ #include "compiler/machine_mapping/allowed_machine_views.h" -#include #include "utils/containers/extend.h" #include "utils/containers/range.h" #include "utils/containers/transform.h" #include "utils/containers/unordered_set_of.h" #include "utils/containers/zip.h" #include "utils/fmt/unordered_set.h" +#include #include using namespace FlexFlow; diff --git a/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc b/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc index 29fdfff3f6..ce4c9b4b25 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/machine_view.cc @@ -162,8 +162,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 0_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/2_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/2_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -171,8 +171,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/4_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/4_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -180,8 +180,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/2_n, /*device_idx=*/2_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/2_n, /*device_idx=*/2_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -189,8 +189,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/2_n, /*device_idx=*/4_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/2_n, /*device_idx=*/4_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -237,8 +237,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 0_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/0_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/0_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -246,8 +246,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/4_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/4_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -255,8 +255,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/1_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/1_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -264,8 +264,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/5_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/5_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -302,8 +302,7 @@ TEST_SUITE(FF_TEST_SUITE) { }}, }; MachineView mv = MachineView{ - MachineSpaceCoordinate{ - /*node_idx=*/0_n, /*device_idx=*/1_n}, + MachineSpaceCoordinate{/*node_idx=*/0_n, /*device_idx=*/1_n}, {MachineViewDimension{stride_t{1_p}, MachineSpecificationDimension::INTER_NODE}, MachineViewDimension{stride_t{2_p}, @@ -313,8 +312,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n, 0_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/0_n, /*device_idx=*/3_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/0_n, /*device_idx=*/3_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -322,8 +321,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n, 1_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/5_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/5_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); @@ -331,8 +330,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n, 1_n}); - MachineSpaceCoordinate correct = MachineSpaceCoordinate{ - /*node_idx=*/1_n, /*device_idx=*/7_n}; + MachineSpaceCoordinate correct = + MachineSpaceCoordinate{/*node_idx=*/1_n, /*device_idx=*/7_n}; MachineSpaceCoordinate result = get_machine_space_coordinate(task, mv, coord); CHECK(correct == result); diff --git a/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc b/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc index e3b08e3805..be04954f18 100644 --- a/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc +++ b/lib/compiler/test/src/compiler/machine_mapping/start_invariant_machine_view.cc @@ -36,8 +36,7 @@ TEST_SUITE(FF_TEST_SUITE) { } TEST_CASE("StartInvariantMachineView - conversions") { - MachineSpaceCoordinate start = - MachineSpaceCoordinate{1_n, 2_n}; + MachineSpaceCoordinate start = MachineSpaceCoordinate{1_n, 2_n}; std::vector dimensions = { MachineViewDimension{stride_t{2_p}, MachineSpecificationDimension::INTER_NODE}, @@ -45,8 +44,7 @@ TEST_SUITE(FF_TEST_SUITE) { MachineSpecificationDimension::INTRA_NODE}}; MachineView mv = MachineView{start, dimensions}; - StartInvariantMachineView simv = - StartInvariantMachineView{dimensions}; + StartInvariantMachineView simv = StartInvariantMachineView{dimensions}; SUBCASE("start_invariant_from_machine_view") { StartInvariantMachineView result = start_invariant_from_machine_view(mv); @@ -93,9 +91,9 @@ TEST_SUITE(FF_TEST_SUITE) { 3_ge2, }}, }; - StartInvariantMachineView simv = StartInvariantMachineView{ - {MachineViewDimension{stride_t{2_p}, - MachineSpecificationDimension::INTRA_NODE}}}; + StartInvariantMachineView simv = + StartInvariantMachineView{{MachineViewDimension{ + stride_t{2_p}, MachineSpecificationDimension::INTRA_NODE}}}; MachineComputeSpecification ms = MachineComputeSpecification{ /*num_nodes=*/1_p, /*num_cpus_per_node=*/6_p, @@ -105,8 +103,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("get_machine_space_offset") { SUBCASE("Task with TaskSpaceCoordinate = (0,)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{0, 0}; + MachineSpaceOffset correct = MachineSpaceOffset{0, 0}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -114,8 +111,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{0, 2}; + MachineSpaceOffset correct = MachineSpaceOffset{0, 2}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -123,8 +119,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (2,)") { TaskSpaceCoordinate coord = make_task_space_coordinate({2_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{0, 4}; + MachineSpaceOffset correct = MachineSpaceOffset{0, 4}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -178,8 +173,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("get_machine_space_offset") { SUBCASE("Task with TaskSpaceCoordinate = (0,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 0_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{0, 0}; + MachineSpaceOffset correct = MachineSpaceOffset{0, 0}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -187,8 +181,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (0,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({0_n, 1_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{0, 2}; + MachineSpaceOffset correct = MachineSpaceOffset{0, 2}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -196,8 +189,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,0)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 0_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{1, 0}; + MachineSpaceOffset correct = MachineSpaceOffset{1, 0}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); @@ -205,8 +197,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("Task with TaskSpaceCoordinate = (1,1)") { TaskSpaceCoordinate coord = make_task_space_coordinate({1_n, 1_n}); - MachineSpaceOffset correct = - MachineSpaceOffset{1, 2}; + MachineSpaceOffset correct = MachineSpaceOffset{1, 2}; MachineSpaceOffset result = get_machine_space_offset(task, simv, coord); CHECK(correct == result); diff --git a/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc b/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc index 0fe4277ab6..f006c21ba3 100644 --- a/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc +++ b/lib/compiler/test/src/compiler/mcmc/mcmc_over_mapped_pcg.cc @@ -1,6 +1,5 @@ #include "compiler/mcmc/mcmc_over_mapped_pcg.h" #include "compiler/task_graph_simulator/task_simulator.h" -#include #include "internal/runtime_only_cost_estimator_for_test.h" #include "op-attrs/parallel_tensor_dims.h" #include "op-attrs/parallel_tensor_shape.dtg.h" @@ -10,6 +9,7 @@ #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h" #include "pcg/pcg_from_computation_graph.h" #include "utils/integer_conversions.h" +#include using namespace FlexFlow; diff --git a/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc b/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc index a7f8886846..b46619ab20 100644 --- a/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc +++ b/lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc @@ -65,10 +65,8 @@ TEST_SUITE(FF_TEST_SUITE) { std::vector dims = {}; ParallelComputationGraph pcg = b.pcg; - MachineView mv1 = - MachineView{MachineSpaceCoordinate{0_n, 0_n}, dims}; - MachineView mv2 = - MachineView{MachineSpaceCoordinate{0_n, 1_n}, dims}; + MachineView mv1 = MachineView{MachineSpaceCoordinate{0_n, 0_n}, dims}; + MachineView mv2 = MachineView{MachineSpaceCoordinate{0_n, 1_n}, dims}; MachineMapping device_mapping = MachineMapping{{ {layer0, mv1}, @@ -146,14 +144,10 @@ TEST_SUITE(FF_TEST_SUITE) { std::vector dims = {}; SUBCASE("all different devices") { - MachineView mv0 = MachineView{ - MachineSpaceCoordinate{0_n, 0_n}, dims}; - MachineView mv1 = MachineView{ - MachineSpaceCoordinate{0_n, 1_n}, dims}; - MachineView mv2 = MachineView{ - MachineSpaceCoordinate{1_n, 0_n}, dims}; - MachineView mv3 = MachineView{ - MachineSpaceCoordinate{1_n, 1_n}, dims}; + MachineView mv0 = MachineView{MachineSpaceCoordinate{0_n, 0_n}, dims}; + MachineView mv1 = MachineView{MachineSpaceCoordinate{0_n, 1_n}, dims}; + MachineView mv2 = MachineView{MachineSpaceCoordinate{1_n, 0_n}, dims}; + MachineView mv3 = MachineView{MachineSpaceCoordinate{1_n, 1_n}, dims}; MachineMapping device_mapping = MachineMapping{{ {layer0, mv0}, @@ -206,8 +200,7 @@ TEST_SUITE(FF_TEST_SUITE) { } SUBCASE("all the same device") { - MachineView mv = MachineView{ - MachineSpaceCoordinate{0_n, 0_n}, dims}; + MachineView mv = MachineView{MachineSpaceCoordinate{0_n, 0_n}, dims}; MachineMapping device_mapping = MachineMapping{{ {layer0, mv}, {layer1, mv}, diff --git a/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc b/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc index cc7ca9425f..a9ee3f34e8 100644 --- a/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc +++ b/lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc @@ -3,7 +3,6 @@ #include "compiler/machine_mapping/machine_mapping.h" #include "compiler/machine_mapping/machine_view.dtg.h" #include "compiler/machine_mapping/machine_view.h" -#include #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.h" #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h" #include "test/utils/doctest/check_without_stringify.h" diff --git a/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc b/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc index 0d4123d381..65a1860034 100644 --- a/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc +++ b/lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc @@ -1,6 +1,5 @@ #include "compiler/unity_algorithm/unity_algorithm.h" #include "compiler/cost_estimator/runtime_only_cost_estimator_from_cost_estimator.h" -#include #include "internal/cost_estimator_for_test.h" #include "op-attrs/parallel_tensor_dims.h" #include "op-attrs/parallel_tensor_shape.dtg.h" @@ -10,6 +9,7 @@ #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h" #include "pcg/pcg_from_computation_graph.h" #include "utils/integer_conversions.h" +#include using namespace FlexFlow; diff --git a/lib/local-execution/include/local-execution/computation_graph_instance.h b/lib/local-execution/include/local-execution/computation_graph_instance.h index e6c7ded9c1..23467f8da5 100644 --- a/lib/local-execution/include/local-execution/computation_graph_instance.h +++ b/lib/local-execution/include/local-execution/computation_graph_instance.h @@ -7,12 +7,12 @@ #include "kernels/profiling_settings.dtg.h" #include "local-execution/loss_config.dtg.h" #include "pcg/computation_graph.dtg.h" -#include "task-spec/global_device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "utils/units/milliseconds_t.h" #include #include diff --git a/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h b/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h index 15aa267eec..8c4af368b0 100644 --- a/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h +++ b/lib/local-execution/include/local-execution/cost_estimator/local_cost_estimator.h @@ -5,8 +5,8 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "task-spec/global_device_id_t.dtg.h" #include "pcg/machine_interconnect_specification.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" namespace FlexFlow { diff --git a/lib/local-execution/include/local-execution/local_task_argument_accessor.h b/lib/local-execution/include/local-execution/local_task_argument_accessor.h index 38b1023cd2..e2d093967c 100644 --- a/lib/local-execution/include/local-execution/local_task_argument_accessor.h +++ b/lib/local-execution/include/local-execution/local_task_argument_accessor.h @@ -2,8 +2,8 @@ #define _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_LOCAL_TASK_ARGUMENT_ACCESSOR_H #include "kernels/accessor.h" -#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/task_argument_accessor/itask_argument_accessor.h" #include "task-spec/task_argument_accessor/task_tensor_parameter.dtg.h" #include diff --git a/lib/local-execution/include/local-execution/per_device_op_state_initialization.h b/lib/local-execution/include/local-execution/per_device_op_state_initialization.h index 90a1ffc688..4c546992e4 100644 --- a/lib/local-execution/include/local-execution/per_device_op_state_initialization.h +++ b/lib/local-execution/include/local-execution/per_device_op_state_initialization.h @@ -4,9 +4,9 @@ #include "kernels/allocation.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/profiling_settings.dtg.h" -#include "task-spec/global_device_id_t.dtg.h" #include "pcg/optimizer_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" namespace FlexFlow { diff --git a/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc b/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc index f151ab1d17..6c74264dff 100644 --- a/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc +++ b/lib/local-execution/src/local-execution/cost_estimator/local_cost_estimator.cc @@ -103,8 +103,8 @@ OpCostMetrics LocalCostEstimator::estimate_cost( // allocate memory std::shared_ptr tracked_allocator_ptr = - std::make_shared(create_local_allocator_for_device_type( - this->device_idx.device_type)); + std::make_shared( + create_local_allocator_for_device_type(this->device_idx.device_type)); layer_guid_t layer_guid = layer_guid_t{Node{0}}; diff --git a/lib/local-execution/test/src/local-execution/computation_graph_instance.cc b/lib/local-execution/test/src/local-execution/computation_graph_instance.cc index 69091b5d2e..ae8365b127 100644 --- a/lib/local-execution/test/src/local-execution/computation_graph_instance.cc +++ b/lib/local-execution/test/src/local-execution/computation_graph_instance.cc @@ -140,14 +140,13 @@ TEST_SUITE(FF_TEST_SUITE) { /*weight_decay=*/0.001}}; device_handle_t ff_handle = cpu_make_device_handle_t(); global_device_id_t global_device_id = global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - }, - /*device_type=*/DeviceType::CPU, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::CPU, }; - std::unordered_map input_tensors; ComputationGraphInstance computation_graph_instance = @@ -313,11 +312,11 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }, }; global_device_id_t device_idx = global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - }, - /*device_type=*/DeviceType::GPU, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::GPU, }; device_handle_t ff_handle = gpu_make_device_handle_t(managed_handle.raw_handle()); @@ -434,11 +433,11 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }; global_device_id_t device_idx = global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - }, - /*device_type=*/DeviceType::GPU, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::GPU, }; device_handle_t ff_handle = diff --git a/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc b/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc index 8d1506b5fc..3d1b691a08 100644 --- a/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc +++ b/lib/local-execution/test/src/local-execution/cost_estimator/local_cost_estimator.cc @@ -19,11 +19,11 @@ TEST_SUITE(FF_TEST_SUITE) { Allocator allocator = create_local_cpu_memory_allocator(); device_handle_t ff_handle = cpu_make_device_handle_t(); global_device_id_t device_idx = global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - }, - /*device_type=*/DeviceType::CPU, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::CPU, }; OptimizerAttrs optimizer_attrs = OptimizerAttrs{ @@ -69,10 +69,9 @@ TEST_SUITE(FF_TEST_SUITE) { /*output_shapes=*/{{TensorSlotName::OUTPUT, output_shape}}, /*optimizer_attrs=*/optimizer_attrs, /*machine_view=*/ - make_1d_machine_view( - MachineSpaceCoordinate{0_n, 0_n}, - MachineSpecificationDimension::INTRA_NODE, - stride_t{1_p}), + make_1d_machine_view(MachineSpaceCoordinate{0_n, 0_n}, + MachineSpecificationDimension::INTRA_NODE, + stride_t{1_p}), }; OpCostMetrics result = cost_estimator.estimate_cost(op_cost_estimate_key); @@ -94,11 +93,11 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { Allocator allocator = create_local_cuda_memory_allocator(); global_device_id_t device_idx = global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - }, - /*device_type=*/DeviceType::GPU, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::GPU, }; device_handle_t ff_handle = @@ -168,10 +167,9 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*output_shapes=*/{{TensorSlotName::OUTPUT, output_shape}}, /*optimizer_attrs=*/optimizer_attrs, /*machine_view=*/ - make_1d_machine_view( - MachineSpaceCoordinate{0_n, 0_n}, - MachineSpecificationDimension::INTRA_NODE, - stride_t{1_p}), + make_1d_machine_view(MachineSpaceCoordinate{0_n, 0_n}, + MachineSpecificationDimension::INTRA_NODE, + stride_t{1_p}), }; OpCostMetrics result = cost_estimator.estimate_cost(op_cost_estimate_key); diff --git a/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc b/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc index dba4464611..8800711fef 100644 --- a/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc +++ b/lib/local-execution/test/src/local-execution/local_task_argument_accessor.cc @@ -52,11 +52,11 @@ TEST_SUITE(FF_TEST_SUITE) { }; global_device_id_t device_idx = global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/0_n, - /*device_idx=*/0_n, - }, - /*device_type=*/DeviceType::CPU, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/DeviceType::CPU, }; LocalTaskArgumentAccessor acc = LocalTaskArgumentAccessor{ diff --git a/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h b/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h index a7e265f777..cc08f963e7 100644 --- a/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h +++ b/lib/realm-execution/include/realm-execution/device_specific_managed_per_device_ff_handle.h @@ -11,7 +11,8 @@ namespace FlexFlow { DeviceSpecificPtr make_device_specific_managed_ff_handle( - global_device_id_t const &, std::optional const &); + global_device_id_t const &, + std::optional const &); device_handle_t device_handle_t_from_device_specific_managed_ff_handle( DeviceSpecificPtr const &, global_device_id_t); diff --git a/lib/realm-execution/include/realm-execution/device_specific_ptr.h b/lib/realm-execution/include/realm-execution/device_specific_ptr.h index f1126f64c1..fe8da0f2e6 100644 --- a/lib/realm-execution/include/realm-execution/device_specific_ptr.h +++ b/lib/realm-execution/include/realm-execution/device_specific_ptr.h @@ -32,7 +32,8 @@ template struct DeviceSpecificPtr { public: DeviceSpecificPtr() = delete; - explicit DeviceSpecificPtr(global_device_id_t device_idx, std::optional ptr) + explicit DeviceSpecificPtr(global_device_id_t device_idx, + std::optional ptr) : device_idx(device_idx), ptr(ptr) {} std::optional get(global_device_id_t device_idx) const { diff --git a/lib/realm-execution/include/realm-execution/fmt/realm_processor.h b/lib/realm-execution/include/realm-execution/fmt/realm_processor.h index e2fadb6b18..12c8787b2f 100644 --- a/lib/realm-execution/include/realm-execution/fmt/realm_processor.h +++ b/lib/realm-execution/include/realm-execution/fmt/realm_processor.h @@ -8,10 +8,10 @@ namespace fmt { template -struct formatter< - ::FlexFlow::Realm::Processor, - Char, - std::enable_if_t::value>> +struct formatter<::FlexFlow::Realm::Processor, + Char, + std::enable_if_t::value>> : formatter<::std::string> { template auto format(::FlexFlow::Realm::Processor const &m, FormatContext &ctx) @@ -27,7 +27,8 @@ struct formatter< namespace FlexFlow { -std::ostream &operator<<(std::ostream &s, ::FlexFlow::Realm::Processor const &m); +std::ostream &operator<<(std::ostream &s, + ::FlexFlow::Realm::Processor const &m); } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h b/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h index d75ab067cc..e591eb9adf 100644 --- a/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h +++ b/lib/realm-execution/include/realm-execution/fmt/realm_processor_kind.h @@ -3,16 +3,13 @@ #include "realm-execution/realm.h" #include -#include #include +#include namespace fmt { template -struct formatter< - ::FlexFlow::Realm::Processor::Kind, - Char - > +struct formatter<::FlexFlow::Realm::Processor::Kind, Char> : formatter<::std::string> { template auto format(::FlexFlow::Realm::Processor::Kind const &m, FormatContext &ctx) @@ -76,7 +73,8 @@ struct formatter< namespace FlexFlow { -std::ostream &operator<<(std::ostream &, ::FlexFlow::Realm::Processor::Kind const &); +std::ostream &operator<<(std::ostream &, + ::FlexFlow::Realm::Processor::Kind const &); } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/pcg_instance.h b/lib/realm-execution/include/realm-execution/pcg_instance.h index 6c82e4c61c..d9d49a780d 100644 --- a/lib/realm-execution/include/realm-execution/pcg_instance.h +++ b/lib/realm-execution/include/realm-execution/pcg_instance.h @@ -11,10 +11,10 @@ #include "realm-execution/per_device_op_state_backing.dtg.h" #include "realm-execution/realm_context.h" #include "realm-execution/tensor_instance_backing.dtg.h" -#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h" #include "task-spec/dynamic_graph/dynamic_tensor_accessor.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" #include "utils/units/milliseconds_t.h" #include diff --git a/lib/realm-execution/include/realm-execution/processor_query.h b/lib/realm-execution/include/realm-execution/processor_query.h index 693d488390..acc3d79440 100644 --- a/lib/realm-execution/include/realm-execution/processor_query.h +++ b/lib/realm-execution/include/realm-execution/processor_query.h @@ -5,7 +5,8 @@ namespace FlexFlow { -std::set processor_set_from_query(Realm::Machine::ProcessorQuery const &); +std::set + processor_set_from_query(Realm::Machine::ProcessorQuery const &); } // namespace FlexFlow diff --git a/lib/realm-execution/include/realm-execution/realm_context.h b/lib/realm-execution/include/realm-execution/realm_context.h index e06ac522d7..9ed8352f97 100644 --- a/lib/realm-execution/include/realm-execution/realm_context.h +++ b/lib/realm-execution/include/realm-execution/realm_context.h @@ -10,9 +10,9 @@ #include "realm-execution/realm.h" #include "realm-execution/tasks/task_id_t.dtg.h" #include "task-spec/global_device_id_t.dtg.h" +#include "task-spec/local_device_id_t.dtg.h" #include #include -#include "task-spec/local_device_id_t.dtg.h" namespace FlexFlow { @@ -35,7 +35,8 @@ struct RealmContext { ///\{ Realm::Processor processor_from_global_device_id(global_device_id_t const &); global_device_id_t global_device_id_from_processor(Realm::Processor); - Realm::Processor processor_from_local_device_id(local_device_id_t const &) const; + Realm::Processor + processor_from_local_device_id(local_device_id_t const &) const; local_device_id_t local_device_id_from_processor(Realm::Processor) const; static Realm::Memory get_nearest_memory(Realm::Processor); @@ -113,7 +114,8 @@ struct RealmContext { */ Realm::Runtime get_runtime(); - bidict const &get_global_machine_topology(); + bidict const & + get_global_machine_topology(); private: Realm::Runtime runtime; @@ -121,7 +123,8 @@ struct RealmContext { Allocator allocator; std::vector outstanding_events; bidict local_machine_topology; - std::optional> cached_global_machine_topology = std::nullopt; + std::optional> + cached_global_machine_topology = std::nullopt; }; } // namespace FlexFlow diff --git a/lib/realm-execution/src/realm-execution/address_space.cc b/lib/realm-execution/src/realm-execution/address_space.cc index 2886f16312..b1f229cce8 100644 --- a/lib/realm-execution/src/realm-execution/address_space.cc +++ b/lib/realm-execution/src/realm-execution/address_space.cc @@ -2,11 +2,12 @@ namespace FlexFlow { -node_idx_t node_idx_from_realm_address_space(Realm::AddressSpace address_space) { +node_idx_t + node_idx_from_realm_address_space(Realm::AddressSpace address_space) { return node_idx_t{ - nonnegative_int{ - static_cast(address_space), - }, + nonnegative_int{ + static_cast(address_space), + }, }; } diff --git a/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc b/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc index a4a4240ff3..7c47701cc9 100644 --- a/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc +++ b/lib/realm-execution/src/realm-execution/distributed_per_device_op_state_initialization.cc @@ -39,8 +39,8 @@ PerDeviceOpStateBacking perform_distributed_per_device_op_state_initialization( invocation); DeviceSpecificPtr *device_state_ptr = - new DeviceSpecificPtr{ctx.get_current_global_device_id(), - std::nullopt}; + new DeviceSpecificPtr{ + ctx.get_current_global_device_id(), std::nullopt}; std::optional completion_event = spawn_per_device_op_state_init_task(ctx, diff --git a/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc b/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc index a4afe3cd1b..efbe36d5cb 100644 --- a/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc +++ b/lib/realm-execution/src/realm-execution/fmt/realm_processor.cc @@ -2,7 +2,8 @@ namespace FlexFlow { -std::ostream &operator<<(std::ostream &s, ::FlexFlow::Realm::Processor const &m) { +std::ostream &operator<<(std::ostream &s, + ::FlexFlow::Realm::Processor const &m) { return s << fmt::to_string(m); } diff --git a/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc b/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc index 2552ec7ef5..b631ef5027 100644 --- a/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc +++ b/lib/realm-execution/src/realm-execution/fmt/realm_processor_kind.cc @@ -2,7 +2,8 @@ namespace FlexFlow { -std::ostream &operator<<(std::ostream &s, ::FlexFlow::Realm::Processor::Kind const &k) { +std::ostream &operator<<(std::ostream &s, + ::FlexFlow::Realm::Processor::Kind const &k) { return (s << fmt::to_string(k)); } diff --git a/lib/realm-execution/src/realm-execution/processor_query.cc b/lib/realm-execution/src/realm-execution/processor_query.cc index b5dd3e12da..f122552e52 100644 --- a/lib/realm-execution/src/realm-execution/processor_query.cc +++ b/lib/realm-execution/src/realm-execution/processor_query.cc @@ -2,7 +2,8 @@ namespace FlexFlow { -std::set processor_set_from_query(Realm::Machine::ProcessorQuery const &pq) { +std::set + processor_set_from_query(Realm::Machine::ProcessorQuery const &pq) { std::set result; for (Realm::Processor p : pq) { result.insert(p); diff --git a/lib/realm-execution/src/realm-execution/realm_context.cc b/lib/realm-execution/src/realm-execution/realm_context.cc index 07d60aa73c..869e71d616 100644 --- a/lib/realm-execution/src/realm-execution/realm_context.cc +++ b/lib/realm-execution/src/realm-execution/realm_context.cc @@ -4,67 +4,62 @@ #include "op-attrs/datatype.h" #include "op-attrs/parallel_tensor_shape.h" #include "op-attrs/tensor_dims.dtg.h" +#include "realm-execution/address_space.h" +#include "realm-execution/fmt/realm_processor.h" +#include "realm-execution/fmt/realm_processor_kind.h" #include "realm-execution/processor_kind.h" +#include "realm-execution/processor_query.h" #include "realm-execution/realm_allocator.h" #include "realm-execution/tasks/task_id_t.dtg.h" #include "realm-execution/tasks/task_id_t.h" +#include "task-spec/global_device_id_t.h" +#include "utils/bidict/algorithms/bidict_from_enumerating.h" +#include "utils/bidict/algorithms/merge_disjoint_bidicts.h" +#include "utils/bidict/algorithms/transform_values.h" +#include "utils/containers/are_all_same.h" #include "utils/containers/contains_key.h" +#include "utils/containers/group_by.h" +#include "utils/containers/set_of.h" #include "utils/containers/transform.h" #include "utils/exception.h" #include "utils/nonnegative_int/nonnegative_int.h" #include "utils/one_to_many/one_to_many.h" #include "utils/optional.h" #include "utils/positive_int/positive_int.h" -#include "utils/bidict/algorithms/merge_disjoint_bidicts.h" -#include "realm-execution/address_space.h" -#include "utils/bidict/algorithms/bidict_from_enumerating.h" -#include "utils/bidict/algorithms/transform_values.h" -#include "utils/containers/group_by.h" -#include "task-spec/global_device_id_t.h" -#include "utils/containers/are_all_same.h" -#include "utils/containers/set_of.h" -#include "realm-execution/processor_query.h" -#include "realm-execution/fmt/realm_processor.h" -#include "realm-execution/fmt/realm_processor_kind.h" namespace FlexFlow { - -bidict - build_local_machine_topology(std::set const &local_procs) -{ +bidict build_local_machine_topology( + std::set const &local_procs) { { bool procs_are_local = are_all_same( - transform(local_procs, - [&](Realm::Processor p) -> Realm::AddressSpace { - return p.address_space(); - })); + transform(local_procs, [&](Realm::Processor p) -> Realm::AddressSpace { + return p.address_space(); + })); ASSERT(procs_are_local); } - OneToMany by_proc_kind = - group_by(local_procs, - [](Realm::Processor p) -> Realm::Processor::Kind { - return p.kind(); - }); + OneToMany by_proc_kind = + group_by(local_procs, [](Realm::Processor p) -> Realm::Processor::Kind { + return p.kind(); + }); auto local_machine_topology_for_proc_kind = [&](Realm::Processor::Kind k) - -> bidict - { + -> bidict { if (!contains(by_proc_kind.left_values(), k)) { return {}; } - bidict enumerated = - bidict_from_enumerating(set_of(by_proc_kind.at_l(k).unwrap_as_unordered_set())).reversed(); + bidict enumerated = + bidict_from_enumerating( + set_of(by_proc_kind.at_l(k).unwrap_as_unordered_set())) + .reversed(); - bidict result = - transform_values( - enumerated, - [&](nonnegative_int idx) -> local_device_id_t { + bidict result = transform_values( + enumerated, [&](nonnegative_int idx) -> local_device_id_t { return local_device_id_t{ - /*idx=*/device_in_node_idx_t{idx}, - /*device_type=*/device_type_from_processor_kind(k), + /*idx=*/device_in_node_idx_t{idx}, + /*device_type=*/device_type_from_processor_kind(k), }; }); @@ -72,48 +67,48 @@ bidict }; return binary_merge_disjoint_bidicts( - local_machine_topology_for_proc_kind(Realm::Processor::Kind::LOC_PROC), - local_machine_topology_for_proc_kind(Realm::Processor::Kind::TOC_PROC)); + local_machine_topology_for_proc_kind(Realm::Processor::Kind::LOC_PROC), + local_machine_topology_for_proc_kind(Realm::Processor::Kind::TOC_PROC)); } static bidict - build_global_machine_topology(std::set const &global_procs) -{ - OneToMany by_node_idx = - group_by(global_procs, - [](Realm::Processor p) -> node_idx_t { - return node_idx_from_realm_address_space(p.address_space()); - }); - - auto build_global_machine_topology_for_node = [&](node_idx_t const &node_idx) - -> bidict - { - std::set procs_for_node = set_of(by_node_idx.at_l(node_idx).unwrap_as_unordered_set()); + build_global_machine_topology( + std::set const &global_procs) { + OneToMany by_node_idx = + group_by(global_procs, [](Realm::Processor p) -> node_idx_t { + return node_idx_from_realm_address_space(p.address_space()); + }); + + auto build_global_machine_topology_for_node = [&](node_idx_t const &node_idx) + -> bidict { + std::set procs_for_node = + set_of(by_node_idx.at_l(node_idx).unwrap_as_unordered_set()); - bidict - local_topology_for_node = build_local_machine_topology(procs_for_node); + bidict local_topology_for_node = + build_local_machine_topology(procs_for_node); return transform_values( - local_topology_for_node, - [&](local_device_id_t const &local_device_id) -> global_device_id_t { - return global_device_id_from_local(local_device_id, node_idx); - }); + local_topology_for_node, + [&](local_device_id_t const &local_device_id) -> global_device_id_t { + return global_device_id_from_local(local_device_id, node_idx); + }); }; return merge_disjoint_bidicts( - transform( - set_of(by_node_idx.left_values()), - build_global_machine_topology_for_node)); + transform(set_of(by_node_idx.left_values()), + build_global_machine_topology_for_node)); } -static bidict discover_local_machine_topology(Realm::Processor local_processor) { +static bidict + discover_local_machine_topology(Realm::Processor local_processor) { Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); pq.same_address_space_as(local_processor); return build_local_machine_topology(processor_set_from_query(pq)); } -static bidict discover_global_machine_topology() { +static bidict + discover_global_machine_topology() { Realm::Machine::ProcessorQuery pq(Realm::Machine::get_machine()); return build_global_machine_topology(processor_set_from_query(pq)); @@ -123,8 +118,7 @@ RealmContext::RealmContext(Realm::Processor processor) : processor(processor), allocator(get_realm_allocator( processor, RealmContext::get_nearest_memory(processor))), - local_machine_topology(discover_local_machine_topology(processor)) -{ } + local_machine_topology(discover_local_machine_topology(processor)) {} RealmContext::~RealmContext() { if (!this->outstanding_events.empty()) { @@ -139,8 +133,8 @@ Realm::Processor RealmContext::processor_from_global_device_id( return this->get_global_machine_topology().at_r(global_device_id); } -global_device_id_t RealmContext::global_device_id_from_processor( - Realm::Processor processor) { +global_device_id_t + RealmContext::global_device_id_from_processor(Realm::Processor processor) { return this->get_global_machine_topology().at_l(processor); } @@ -182,8 +176,8 @@ global_device_id_t RealmContext::get_current_global_device_id() const { Realm::Processor proc = this->get_current_processor(); return global_device_id_from_local( - this->local_device_id_from_processor(proc), - node_idx_from_realm_address_space(proc.address_space())); + this->local_device_id_from_processor(proc), + node_idx_from_realm_address_space(proc.address_space())); } Realm::Event @@ -397,7 +391,8 @@ Realm::Event RealmContext::merge_outstanding_events() { return result; } -bidict const &RealmContext::get_global_machine_topology() { +bidict const & + RealmContext::get_global_machine_topology() { if (!this->cached_global_machine_topology.has_value()) { this->cached_global_machine_topology = discover_global_machine_topology(); } diff --git a/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc b/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc index 75430da876..3897be92ff 100644 --- a/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc +++ b/lib/realm-execution/test/src/realm-execution/fmt/realm_processor_kind.cc @@ -1,11 +1,12 @@ -#include #include "realm-execution/fmt/realm_processor_kind.h" +#include using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("fmt::to_string(Realm::Processor::Kind)") { - std::string result = fmt::to_string(::FlexFlow::Realm::Processor::Kind::TOC_PROC); + std::string result = + fmt::to_string(::FlexFlow::Realm::Processor::Kind::TOC_PROC); std::string correct = ""; CHECK(result == correct); diff --git a/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc b/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc index 7dc6a2cc5e..4374a951f8 100644 --- a/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc +++ b/lib/substitutions/src/substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.cc @@ -2,9 +2,9 @@ #include "substitutions/output_graph/output_graph_expr.h" #include "substitutions/sub_parallel_computation_graph.h" #include "utils/bidict/algorithms/bidict_from_pairs.h" +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" #include "utils/containers/values.h" #include "utils/containers/zip_values_strict.h" -#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" namespace FlexFlow { diff --git a/lib/substitutions/src/substitutions/pcg_pattern_match.cc b/lib/substitutions/src/substitutions/pcg_pattern_match.cc index 46d284c4a0..85a0493e33 100644 --- a/lib/substitutions/src/substitutions/pcg_pattern_match.cc +++ b/lib/substitutions/src/substitutions/pcg_pattern_match.cc @@ -4,13 +4,13 @@ #include "substitutions/unlabelled/unlabelled_graph_pattern.h" #include "utils/bidict/algorithms/bidict_from_keys_and_values.h" #include "utils/bidict/algorithms/bidict_from_map.h" +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" #include "utils/bidict/algorithms/exhaustive_relational_join.h" #include "utils/bidict/algorithms/transform_values.h" #include "utils/containers/is_subseteq_of.h" #include "utils/containers/map_values.h" #include "utils/containers/values.h" #include "utils/containers/zip.h" -#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" namespace FlexFlow { diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h index c5a4976db7..3c5b27c344 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h @@ -1,8 +1,8 @@ #ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H -#include "task-spec/global_device_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" +#include "task-spec/global_device_id_t.dtg.h" namespace FlexFlow { diff --git a/lib/task-spec/include/task-spec/global_device_id_t.h b/lib/task-spec/include/task-spec/global_device_id_t.h index 2f48046ff2..428f3f7a3c 100644 --- a/lib/task-spec/include/task-spec/global_device_id_t.h +++ b/lib/task-spec/include/task-spec/global_device_id_t.h @@ -1,13 +1,14 @@ #ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_GLOBAL_DEVICE_ID_T_H #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_GLOBAL_DEVICE_ID_T_H +#include "pcg/node_idx_t.dtg.h" #include "task-spec/global_device_id_t.dtg.h" #include "task-spec/local_device_id_t.dtg.h" -#include "pcg/node_idx_t.dtg.h" namespace FlexFlow { -global_device_id_t global_device_id_from_local(local_device_id_t const &, node_idx_t); +global_device_id_t global_device_id_from_local(local_device_id_t const &, + node_idx_t); local_device_id_t local_device_id_from_global(global_device_id_t const &); } // 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 3487b2979e..2c00d7412a 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 @@ -61,10 +61,12 @@ bool graph_is_fully_copy_inserted(DynamicOpenDataflowGraph const &g) { static std::pair filter_mapping_to_avoid_degenerate_copies(DynamicValueAttrs const &input, DynamicValueAttrs const &output) { - std::unordered_set> + std::unordered_set< + std::pair> input_mapping = unordered_set_of(assert_unwrap(input.mapping).raw); - std::unordered_set> + std::unordered_set< + std::pair> output_mapping = unordered_set_of(assert_unwrap(output.mapping).raw); // Exclude the point shared between the input and output mappings, because @@ -214,8 +216,7 @@ std::unordered_map [&](InternalDynamicSlotSite const &s) -> ParallelTensorMapping { return ParallelTensorMapping{ dynamic_node_mapping_bindings_for_slot_name( - assert_unwrap(i.node_attrs.mapping), - s.slot_name.slot_name), + assert_unwrap(i.node_attrs.mapping), s.slot_name.slot_name), }; }); }; diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc index 4e21da9c2a..b2a6e71af8 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc @@ -11,7 +11,8 @@ bidict get_tensor_bindings_for_slot_name(mapping.op_task_group, slot_name); return transform_values( - coord_bindings, [&](MachineSpaceCoordinate const &coord) -> global_device_id_t { + coord_bindings, + [&](MachineSpaceCoordinate const &coord) -> global_device_id_t { return global_device_id_t{coord, mapping.device_type}; }); } 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 15f78944c8..887c53998f 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 @@ -44,7 +44,8 @@ bool graph_is_fully_shard_expanded(DynamicOpenDataflowGraph const &g) { static bidict restrict_tensor_mapping_keys_to_coord( - bidict const &mapping, + bidict const + &mapping, ParallelTensorSpaceCoordinate const ¶llel_tensor_coord) { return filter_keys(mapping, [&](ParallelTensorSpaceCoordinate const &p) { return p == parallel_tensor_coord; @@ -128,13 +129,14 @@ std::unordered_set std::unordered_set shard_machine_coords = target_devices_of_dynamic_node_mapping(mapping); - return transform( - shard_machine_coords, [&](global_device_id_t const &c) -> DynamicNodeInvocation { - OperatorAtomicTaskShardBinding slot_bindings = - mapping.op_task_group.get_shard_bindings().at_l(c.coord); + return transform(shard_machine_coords, + [&](global_device_id_t const &c) -> DynamicNodeInvocation { + OperatorAtomicTaskShardBinding slot_bindings = + mapping.op_task_group.get_shard_bindings().at_l( + c.coord); - return shard_invocation_for_binding(i, c, slot_bindings); - }); + return shard_invocation_for_binding(i, c, slot_bindings); + }); } DynamicOpenDataflowGraph diff --git a/lib/task-spec/src/task-spec/global_device_id_t.cc b/lib/task-spec/src/task-spec/global_device_id_t.cc index e32a5f5f8b..af0a52293a 100644 --- a/lib/task-spec/src/task-spec/global_device_id_t.cc +++ b/lib/task-spec/src/task-spec/global_device_id_t.cc @@ -2,23 +2,23 @@ namespace FlexFlow { -global_device_id_t global_device_id_from_local( - local_device_id_t const &local_device_id, - node_idx_t node_idx) -{ +global_device_id_t + global_device_id_from_local(local_device_id_t const &local_device_id, + node_idx_t node_idx) { return global_device_id_t{ - /*coord=*/MachineSpaceCoordinate{ - /*node_idx=*/node_idx.raw, - /*device_idx=*/local_device_id.idx.raw, - }, - /*device_type=*/local_device_id.device_type, + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/node_idx.raw, + /*device_idx=*/local_device_id.idx.raw, + }, + /*device_type=*/local_device_id.device_type, }; } -local_device_id_t local_device_id_from_global(global_device_id_t const &global_device_id) { +local_device_id_t + local_device_id_from_global(global_device_id_t const &global_device_id) { return local_device_id_t{ - /*idx=*/device_in_node_idx_t{global_device_id.coord.device_idx}, - /*device_type=*/global_device_id.device_type, + /*idx=*/device_in_node_idx_t{global_device_id.coord.device_idx}, + /*device_type=*/global_device_id.device_type, }; } diff --git a/lib/task-spec/test/src/task-spec/device_specific.cc b/lib/task-spec/test/src/task-spec/device_specific.cc index 983ba4943d..2d4082bb1a 100644 --- a/lib/task-spec/test/src/task-spec/device_specific.cc +++ b/lib/task-spec/test/src/task-spec/device_specific.cc @@ -7,12 +7,14 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("DeviceSpecific") { DeviceSpecific device_specific1 = DeviceSpecific::create( - global_device_id_t{MachineSpaceCoordinate{0_n, 0_n}, DeviceType::GPU}, + global_device_id_t{MachineSpaceCoordinate{0_n, 0_n}, + DeviceType::GPU}, "hello world"); DeviceSpecific device_specific2 = DeviceSpecific::create( - global_device_id_t{MachineSpaceCoordinate{0_n, 1_n}, DeviceType::GPU}, + global_device_id_t{MachineSpaceCoordinate{0_n, 1_n}, + DeviceType::GPU}, "hello world"); std::string result1 = fmt::to_string(device_specific1); diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc index cc51c838b9..0c50416e7c 100644 --- a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc +++ b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc @@ -44,14 +44,16 @@ TEST_SUITE(FF_TEST_SUITE) { }; DeviceType device_type = DeviceType::GPU; - auto mk_device_id = [&](MachineSpaceCoordinate const &c) -> global_device_id_t { + auto mk_device_id = + [&](MachineSpaceCoordinate const &c) -> global_device_id_t { return global_device_id_t{c, device_type}; }; auto mk_value = [&](size_t src_node_id, TensorSlotName src_slot_name, - bidict tensor_binding, + bidict + tensor_binding, std::optional const &shard_coord) -> DynamicValueAttrs { if (shard_coord.has_value()) { @@ -155,9 +157,9 @@ TEST_SUITE(FF_TEST_SUITE) { TensorSlotName use_slot_name, std::optional const &shard_coord) -> DynamicValueAttrs { - bidict tensor_binding = - dynamic_node_mapping_bindings_for_slot_name(node_mapping, - use_slot_name); + bidict + tensor_binding = dynamic_node_mapping_bindings_for_slot_name( + node_mapping, use_slot_name); return mk_value( src_node_id, src_slot_name, tensor_binding, shard_coord); }; diff --git a/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h b/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h index 607cb08229..bcc59c1e6b 100644 --- a/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h +++ b/lib/utils/include/utils/bidict/algorithms/binary_merge_disjoint_bidicts.h @@ -1,26 +1,26 @@ #ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_BINARY_MERGE_DISJOINT_BIDICTS_H #define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_BINARY_MERGE_DISJOINT_BIDICTS_H +#include "utils/bidict/algorithms/left_entries.h" +#include "utils/bidict/algorithms/right_entries.h" #include "utils/bidict/bidict.h" -#include #include "utils/containers/are_disjoint.h" -#include "utils/bidict/algorithms/right_entries.h" -#include "utils/bidict/algorithms/left_entries.h" +#include namespace FlexFlow { template bidict binary_merge_disjoint_bidicts(bidict const &lhs, bidict const &rhs) { - ASSERT( - are_disjoint(left_entries(lhs), left_entries(rhs)), - "Left entries of {} and {} are non-disjoint", lhs, rhs - ); - - ASSERT( - are_disjoint(right_entries(lhs), right_entries(rhs)), - "Right entries of {} and {} are non-disjoint", lhs, rhs - ); + ASSERT(are_disjoint(left_entries(lhs), left_entries(rhs)), + "Left entries of {} and {} are non-disjoint", + lhs, + rhs); + + ASSERT(are_disjoint(right_entries(lhs), right_entries(rhs)), + "Right entries of {} and {} are non-disjoint", + lhs, + rhs); bidict result; for (auto const &kv : lhs) { diff --git a/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h b/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h index 1ce7cb3a13..add49b3089 100644 --- a/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h +++ b/lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h @@ -1,20 +1,20 @@ #ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_DISJOINT_BIDICTS_H #define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_DISJOINT_BIDICTS_H +#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" #include "utils/bidict/bidict.h" #include "utils/containers/foldl.h" -#include "utils/bidict/algorithms/binary_merge_disjoint_bidicts.h" namespace FlexFlow { template bidict merge_disjoint_bidicts(std::set> const &bidicts) { return foldl( - bidicts, - bidict{}, - [](bidict const &accum, bidict const &x) -> bidict { - return binary_merge_disjoint_bidicts(accum, x); - }); + bidicts, + bidict{}, + [](bidict const &accum, bidict const &x) -> bidict { + return binary_merge_disjoint_bidicts(accum, x); + }); } } // namespace FlexFlow diff --git a/lib/utils/include/utils/variant.h b/lib/utils/include/utils/variant.h index ba689d57be..3c63a28c11 100644 --- a/lib/utils/include/utils/variant.h +++ b/lib/utils/include/utils/variant.h @@ -1,9 +1,9 @@ #ifndef _FLEXFLOW_UTILS_VARIANT_H #define _FLEXFLOW_UTILS_VARIANT_H -#include #include "utils/type_traits.h" #include +#include #include #include diff --git a/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc b/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc index d3deb887b1..d0d2aa3b8f 100644 --- a/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc +++ b/lib/utils/src/utils/bidict/algorithms/binary_merge_disjoint_bidicts.cc @@ -6,8 +6,7 @@ namespace FlexFlow { using K = ordered_value_type<0>; using V = ordered_value_type<1>; -template - bidict binary_merge_disjoint_bidicts(bidict const &, - bidict const &); +template bidict binary_merge_disjoint_bidicts(bidict const &, + bidict const &); } // namespace FlexFlow diff --git a/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h b/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h index 15ea9fc663..bca939efc2 100644 --- a/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h +++ b/lib/utils/test/common/include/test/utils/rapidcheck/doctest.h @@ -1,8 +1,8 @@ #ifndef _FLEXFLOW_UTILS_TEST_COMMON_INCLUDE_TEST_UTILS_RAPIDCHECK_DOCTEST_H #define _FLEXFLOW_UTILS_TEST_COMMON_INCLUDE_TEST_UTILS_RAPIDCHECK_DOCTEST_H -#include #include +#include namespace FlexFlow { diff --git a/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc b/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc index 4d91d13c05..eb0268e3bd 100644 --- a/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc +++ b/lib/utils/test/src/utils/graph/series_parallel/non_normal_sp_decomposition.cc @@ -1,6 +1,6 @@ #include "utils/graph/series_parallel/non_normal_sp_decomposition.h" -#include #include "utils/graph/series_parallel/series_parallel_decomposition.dtg.h" +#include using namespace ::FlexFlow; From fb7280a276eba6e87b62f87a4843081dbe2abe4a Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Fri, 12 Jun 2026 21:13:11 -0700 Subject: [PATCH 5/5] Minor fix in run-model --- bin/run-model/src/run-model/main.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bin/run-model/src/run-model/main.cc b/bin/run-model/src/run-model/main.cc index 46e1186b6b..a6f1d026ef 100644 --- a/bin/run-model/src/run-model/main.cc +++ b/bin/run-model/src/run-model/main.cc @@ -100,7 +100,8 @@ int main(int argc, char **argv) { /*loss=*/std::nullopt, /*input_tensors=*/input_tensors, /*profiling_settings=*/ProfilingSettings{0, 0}, - /*device_handle=*/device_handle); + /*device_handle=*/device_handle, + /*device_type=*/DeviceType::GPU); // begin training loop int num_epochs = 5;