Generalize index space constructs and add data dependence info to op-attrs

.proj.toml

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

CMakeLists.txt

@@ -102,7 +102,6 @@ include(spdlog)
 include(doctestlib) # named doctestlib to avoid a name collision with doctest.cmake in rapidcheck
 include(gbenchmark)
 include(libassert)
-include(visit_struct)
 include(CTest)
 include(fmt)
 include(legion)

flake.nix

@@ -93,9 +93,6 @@
                                 -DFF_USE_EXTERNAL_EXPECTED=ON \
                                 -DFF_USE_EXTERNAL_GBENCHMARK=ON \
                                 -DFF_USE_EXTERNAL_LIBASSERT=ON \
-                                -DFF_USE_EXTERNAL_RANGEV3=ON \
-                                -DFF_USE_EXTERNAL_BOOST_PREPROCESSOR=ON \
-                                -DFF_USE_EXTERNAL_TYPE_INDEX=ON"
           '';
 
           buildInputs = builtins.concatLists [

lib/CMakeLists.txt

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

lib/compiler/include/compiler/allowed_machine_views.h

@@ -1,18 +1,18 @@
 #ifndef _FLEXFLOW_COMPILER_ALLOWED_MACHINE_VIEWS_H
 #define _FLEXFLOW_COMPILER_ALLOWED_MACHINE_VIEWS_H
 
-#include "pcg/machine_specification.dtg.h"
+#include "op-attrs/operator_task_space.dtg.h"
+#include "pcg/machine_compute_specification.dtg.h"
 #include "pcg/machine_view.dtg.h"
-#include "pcg/operator_task_space.dtg.h"
 
 namespace FlexFlow {
 
 bool is_valid_machine_view(MachineView const &mv,
                            OperatorTaskSpace const &task,
-                           MachineSpecification const &ms);
+                           MachineComputeSpecification const &ms);
 
 std::unordered_set<MachineView>
-    get_allowed_machine_views(MachineSpecification const &machine_spec,
+    get_allowed_machine_views(MachineComputeSpecification const &machine_spec,
                               OperatorTaskSpace const &task,
                               DeviceType device_type);
 

lib/compiler/include/compiler/cost_estimator/communication_edge.h

@@ -0,0 +1,50 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COMMUNICATION_EDGE_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COMMUNICATION_EDGE_H
+
+#include "pcg/machine_space_coordinate.dtg.h"
+
+namespace FlexFlow {
+
+struct CommunicationEdge {
+  CommunicationEdge() = delete;
+
+  CommunicationEdge(MachineSpaceCoordinate const &src,
+                    MachineSpaceCoordinate const &dst);
+
+  bool operator==(CommunicationEdge const &) const;
+  bool operator!=(CommunicationEdge const &) const;
+
+  bool operator<(CommunicationEdge const &) const;
+  bool operator>(CommunicationEdge const &) const;
+  bool operator<=(CommunicationEdge const &) const;
+  bool operator>=(CommunicationEdge const &) const;
+
+  MachineSpaceCoordinate const &get_src() const;
+  MachineSpaceCoordinate const &get_dst() const;
+private:
+  MachineSpaceCoordinate src;
+  MachineSpaceCoordinate dst;
+private:
+  std::tuple<
+    decltype(src) const &,
+    decltype(dst) const &
+  > tie() const;
+
+  friend struct ::std::hash<CommunicationEdge>;
+};
+
+std::string format_as(CommunicationEdge const &);
+std::ostream &operator<<(std::ostream &, CommunicationEdge const &);
+
+} // namespace FlexFlow
+
+namespace std {
+
+template <>
+struct hash<::FlexFlow::CommunicationEdge> {
+  size_t operator()(::FlexFlow::CommunicationEdge const &) const;
+};
+
+}
+
+#endif

lib/compiler/include/compiler/cost_estimator/network_cost_model.h

@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_NETWORK_COST_MODEL_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_NETWORK_COST_MODEL_H
+
+#include "compiler/cost_estimator/tensor_set_movement.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+
+namespace FlexFlow {
+
+float estimate_communication_cost(MachineSpecification const &,
+                                  TensorSetMovement const &);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/cost_estimator/op_cost_metrics.h

@@ -6,6 +6,9 @@
 
 namespace FlexFlow {
 
+bool is_pareto_optimal_in(OpCostMetrics const &,
+                          std::unordered_set<OpCostMetrics> const &);
+
 OpCostMetrics make_op_cost_metrics_from_runtime_only(
     RuntimeOnlyOpCostMetrics const &runtime_only,
     num_bytes_t const &memory_usage);

lib/compiler/include/compiler/cost_estimator/parallel_tensor_space_to_machine_space_mapping.cc

@@ -0,0 +1,45 @@
+#include "compiler/cost_estimator/parallel_tensor_space_to_machine_space_mapping.h"
+#include "op-attrs/operator_space_to_parallel_tensor_space_mapping.h"
+#include "op-attrs/parallel_tensor_dim_degrees.h"
+#include "utils/bidict/algorithms/exhaustive_relational_join.h"
+#include <libassert/assert.hpp>
+#include "op-attrs/task_space_coordinate.h"
+#include "op-attrs/parallel_tensor_space_coordinate.h"
+#include "utils/bidict/algorithms/transform_keys.h"
+#include "utils/bidict/algorithms/transform_values.h"
+
+namespace FlexFlow {
+
+ParallelTensorSpaceToMachineSpaceMapping
+  ptensor_machine_map_from_composition(
+    OperatorSpaceToMachineSpaceMapping const &op_task_to_machine_space_mapping, 
+    OperatorSpaceToParallelTensorSpaceMapping const &op_task_to_parallel_tensor_space_mapping) {
+  ASSERT(
+    op_task_to_machine_space_mapping.operator_task_space
+    ==
+    get_operator_task_space_for_mapping(op_task_to_parallel_tensor_space_mapping)
+  );
+
+  bidict<ParallelTensorSpaceCoordinate, TaskSpaceCoordinate>
+    pt_to_op_coord_map =
+      transform_keys( 
+        transform_values(
+          op_task_to_parallel_tensor_space_mapping.raw_mapping.coord_mapping.reversed(),
+          task_space_coordinate_from_dim_coord),
+        parallel_tensor_space_coord_from_dim_coord);
+
+  bidict<TaskSpaceCoordinate, MachineSpaceCoordinate>
+    op_to_ms_coord_map =
+      op_task_to_machine_space_mapping.raw_mapping;
+
+  return ParallelTensorSpaceToMachineSpaceMapping{
+    /*raw_mapping=*/exhaustive_relational_join(
+      pt_to_op_coord_map,
+      op_to_ms_coord_map),
+    /*parallel_tensor_space=*/
+      parallel_tensor_dim_degrees_from_dim_domain(
+        op_task_to_parallel_tensor_space_mapping.raw_mapping.r_domain),
+  };
+};
+
+} // namespace FlexFlow

lib/compiler/include/compiler/cost_estimator/parallel_tensor_space_to_machine_space_mapping.h

@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_PARALLEL_TENSOR_SPACE_TO_MACHINE_SPACE_MAPPING_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_PARALLEL_TENSOR_SPACE_TO_MACHINE_SPACE_MAPPING_H
+
+#include "compiler/cost_estimator/parallel_tensor_space_to_machine_space_mapping.dtg.h"
+#include "pcg/operator_space_to_machine_space_mapping.dtg.h"
+#include "op-attrs/operator_space_to_parallel_tensor_space_mapping.dtg.h"
+
+namespace FlexFlow {
+
+ParallelTensorSpaceToMachineSpaceMapping
+  ptensor_machine_map_from_composition(
+    OperatorSpaceToMachineSpaceMapping const &op_task_to_machine_space_mapping, 
+    OperatorSpaceToParallelTensorSpaceMapping const &op_task_to_parallel);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/cost_estimator/parallel_tensor_space_to_machine_space_mapping.struct.toml

@@ -0,0 +1,22 @@
+namespace = "FlexFlow"
+name = "ParallelTensorSpaceToMachineSpaceMapping"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "op-attrs/parallel_tensor_space_coordinate.dtg.h",
+  "pcg/machine_space_coordinate.dtg.h",
+  "op-attrs/parallel_tensor_dim_degrees.dtg.h",
+  "utils/bidict/bidict.h",
+]
+
+[[fields]]
+name = "raw_mapping"
+type = "::FlexFlow::bidict<::FlexFlow::ParallelTensorSpaceCoordinate, ::FlexFlow::MachineSpaceCoordinate>"
+
+[[fields]]
+name = "parallel_tensor_space"
+type = "::FlexFlow::ParallelTensorDimDegrees"

lib/compiler/include/compiler/cost_estimator/single_communication.struct.toml

@@ -0,0 +1,21 @@
+namespace = "FlexFlow"
+name = "SingleCommunication"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "fmt",
+  "json",
+]
+
+includes = [
+  "pcg/machine_space_coordinate.dtg.h",
+]
+
+[[fields]]
+name = "src_machine_coord"
+type = "::FlexFlow::MachineSpaceCoordinate"
+
+[[fields]]
+name = "dst_machine_coord"
+type = "::FlexFlow::MachineSpaceCoordinate"

lib/compiler/include/compiler/cost_estimator/single_tensor_movement.struct.toml

@@ -7,24 +7,16 @@ features = [
 ]
 
 includes = [
-  "op-attrs/parallel_tensor_shape.dtg.h",
-  "pcg/machine_view.dtg.h",
-  "<unordered_set>",
+  "compiler/cost_estimator/communication_edge.h",
+  "utils/units/num_bytes_t.h",
+  "<unordered_map>",
 ]
 
 src_includes = [
-  "utils/hash/unordered_set.h",
-  "utils/fmt/unordered_set.h",
+  "utils/fmt/unordered_map.h",
+  "utils/hash/unordered_map.h",
 ]
 
 [[fields]]
-name = "parallel_tensor_shape"
-type = "::FlexFlow::ParallelTensorShape"
-
-[[fields]]
-name = "src_machine_views"
-type = "std::unordered_set<::FlexFlow::MachineView>"
-
-[[fields]]
-name = "dst_machine_views"
-type = "std::unordered_set<::FlexFlow::MachineView>"
+name = "edge_to_size"
+type = "std::unordered_map<::FlexFlow::CommunicationEdge, ::FlexFlow::num_bytes_t>"

lib/compiler/include/compiler/cost_estimator/tensor_set_movement.h

@@ -8,7 +8,11 @@
 
 namespace FlexFlow {
 
-TensorSetMovement get_tensor_set_movement_from_pcg_edge(
+TensorSetMovement
+  empty_tensor_set_movement();
+
+TensorSetMovement 
+  get_tensor_set_movement_from_pcg_edge(
     ParallelComputationGraphEdge const &edge,
     ParallelComputationGraph const &pcg,
     MachineView const &src_mv,

lib/compiler/include/compiler/cost_estimator/tensor_set_movement.struct.toml

@@ -7,15 +7,16 @@ features = [
 ]
 
 includes = [
-  "compiler/cost_estimator/single_tensor_movement.dtg.h",
-  "<unordered_set>",
+  "compiler/cost_estimator/communication_edge.h",
+  "utils/units/num_bytes_t.h",
+  "<unordered_map>",
 ]
 
 src_includes = [
-  "utils/fmt/unordered_multiset.h",
-  "utils/hash/unordered_multiset.h",
+  "utils/fmt/unordered_map.h",
+  "utils/hash/unordered_map.h",
 ]
 
 [[fields]]
-name = "single_tensor_movements"
-type = "std::unordered_multiset<::FlexFlow::SingleTensorMovement>"
+name = "edge_to_size"
+type = "std::unordered_map<::FlexFlow::CommunicationEdge, ::FlexFlow::num_bytes_t>"

lib/compiler/include/compiler/graph_optimize_result.struct.toml

@@ -1,16 +1,17 @@
 namespace = "FlexFlow"
 name = "GraphOptimizeResult"
-features = [ ]
+features = [
+  "fmt",
+]
 
 includes = [ 
-  "compiler/machine_mapping/machine_mapping.dtg.h",
-  "pcg/parallel_computation_graph/parallel_computation_graph.h"
+  "compiler/mapped_parallel_computation_graph.dtg.h",
 ]
 
-[[fields]]
-name = "pcg"
-type = "::FlexFlow::ParallelComputationGraph"
+src_includes = [
+  "compiler/mapped_parallel_computation_graph.h",
+]
 
 [[fields]]
-name = "machine_mapping"
-type = "::FlexFlow::MachineMapping"
+name = "mapped_pcg"
+type = "::FlexFlow::MappedParallelComputationGraph"