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

.proj.toml

@@ -5,28 +5,28 @@ header_extension = ".h"
 
 build_targets = [
   "utils",
-  "op-attrs",
-  "kernels",
-  "pcg",
-  "substitutions",
-  "compiler",
-  "substitution-generator",
-  "local-execution",
-  "models",
-  "export-model-arch",
-  "substitution-to-dot",
+  # "op-attrs",
+  # "kernels",
+  # "pcg",
+  # "substitutions",
+  # "compiler",
+  # "substitution-generator",
+  # "local-execution",
+  # "models",
+  # "export-model-arch",
+  # "substitution-to-dot",
 ]
 
 test_targets = [
   # "kernels-tests",
   "utils-tests",
-  "op-attrs-tests",
-  "pcg-tests",
-  "substitutions-tests",
-  "compiler-tests",
-  "substitution-generator-tests",
-  "local-execution-tests",
-  "models-tests",
+  # "op-attrs-tests",
+  # "pcg-tests",
+  # "substitutions-tests",
+  # "compiler-tests",
+  # "substitution-generator-tests",
+  # "local-execution-tests",
+  # "models-tests",
 ]
 
 [cmake_flags_extra]

lib/compiler/include/compiler/allowed_machine_views.h

@@ -3,7 +3,7 @@
 
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/machine_view.dtg.h"
-#include "pcg/operator_task_space.dtg.h"
+#include "op-attrs/operator_task_space.dtg.h"
 
 namespace FlexFlow {
 

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/src/compiler/allowed_machine_views.cc

@@ -2,7 +2,7 @@
 #include "pcg/machine_specification.h"
 #include "pcg/machine_view.h"
 #include "pcg/multi_dimensional_stride.dtg.h"
-#include "pcg/operator_task_space.h"
+#include "op-attrs/operator_task_space.h"
 #include "utils/containers/all_of.h"
 #include "utils/containers/cartesian_product.h"
 #include "utils/containers/extend.h"

lib/compiler/src/compiler/cost_estimator/network_cost_model.cc

@@ -0,0 +1,13 @@
+#include "compiler/cost_estimator/network_cost_model.h"
+
+namespace FlexFlow {
+
+float estimate_communication_cost(MachineSpecification const &machine_spec, 
+                                  TensorSetMovement const &tensor_set_movement) {
+  NOT_IMPLEMENTED(); // TODO @lockshaw
+  // for (SingleTensorMovement const &single_tensor_movement : tensor_set_movement.single_tensor_movements) {
+  //   for 
+  // }
+}
+
+} // namespace FlexFlow

lib/kernels/include/kernels/legion_dim.h

@@ -6,6 +6,8 @@
 
 namespace FlexFlow {
 
+std::set<legion_dim_t> legion_dim_range(int end);
+
 legion_dim_t add_to_legion_dim(legion_dim_t legion_dim, int value);
 
 legion_dim_t legion_dim_from_ff_dim(ff_dim_t, int num_dimensions);

lib/kernels/src/legion_dim.cc → lib/kernels/src/kernels/legion_dim.cc

@@ -1,7 +1,14 @@
 #include "kernels/legion_dim.h"
+#include "utils/containers/range.h"
+#include "utils/containers/transform.h"
+#include "utils/containers/set_of.h"
 
 namespace FlexFlow {
 
+std::set<legion_dim_t> legion_dim_range(int end) {
+  return set_of(transform(range(end), [](int i) { return ff_dim_t{i}; }));
+}
+
 legion_dim_t add_to_legion_dim(legion_dim_t legion_dim, int value) {
   return legion_dim_t(legion_dim.value + value);
 }

lib/op-attrs/include/op-attrs/dim_ordered/dim_ordered.h

@@ -123,7 +123,7 @@ struct DimOrdered {
   }
 
   reverse_iterator rend() {
-    return this->contents.crend();
+    return this->contents.rend();
   }
 
   const_reverse_iterator rend() const {

lib/op-attrs/include/op-attrs/dim_ordered/enumerate.h

@@ -3,7 +3,7 @@
 
 #include "op-attrs/dim_ordered/dim_ordered.h"
 #include "utils/bidict/bidict.h"
-#include "utils/containers/count.h"
+#include "utils/containers/range.h"
 
 namespace FlexFlow {
 
@@ -18,7 +18,7 @@ namespace FlexFlow {
 template <typename T>
 std::map<ff_dim_t, T> enumerate(FFOrdered<T> const &ff_ordered) {
   std::map<ff_dim_t, T> result;
-  for (int raw_ff_dim : count(ff_ordered.size())) {
+  for (int raw_ff_dim : range(ff_ordered.size())) {
     ff_dim_t ff_dim = ff_dim_t{raw_ff_dim};
     result.insert({ff_dim, ff_ordered.at(ff_dim)});
   }

lib/op-attrs/include/op-attrs/dim_ordered/get_idxs.h

@@ -2,14 +2,14 @@
 #define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_DIM_ORDERED_GET_IDXS_H
 
 #include "op-attrs/dim_ordered/dim_ordered.h"
-#include "utils/containers/count.h"
+#include "utils/containers/range.h"
 #include "utils/containers/transform.h"
 
 namespace FlexFlow {
 
 template <typename T>
-std::vector<ff_dim_t> get_idxs(FFOrdered<T> const &d) {
-  return transform(count(d.size()), [](int i) { return ff_dim_t{i}; });
+std::set<ff_dim_t> get_idxs(FFOrdered<T> const &d) {
+  return transform(set_of(range(d.size())), [](int i) { return ff_dim_t{i}; });
 }
 
 } // namespace FlexFlow

lib/op-attrs/include/op-attrs/ff_dim_t.h

@@ -0,0 +1,20 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_FF_DIM_T_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_FF_DIM_T_H
+
+#include "op-attrs/ff_dim.dtg.h"
+#include "rapidcheck.h"
+
+namespace FlexFlow {
+
+std::set<ff_dim_t> ff_dim_range(int end);
+
+} // namespace FlexFlow
+
+namespace rc {
+template <>
+struct Arbitrary<FlexFlow::ff_dim_t> {
+  static Gen<FlexFlow::ff_dim_t> arbitrary();
+};
+} // namespace rc
+
+#endif 

lib/op-attrs/include/op-attrs/operator_space_parallel_tensor_space_mapping.h

@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_OPERATOR_SPACE_PARALLEL_TENSOR_SPACE_MAPPING_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_OPERATOR_SPACE_PARALLEL_TENSOR_SPACE_MAPPING_H
+
+#include "op-attrs/operator_space_parallel_tensor_space_mapping.dtg.h"
+#include "op-attrs/tensor_num_dims.dtg.h"
+
+namespace FlexFlow {
+
+OperatorSpaceParallelTensorSpaceMapping
+  get_identity_mapping(TensorNumDims const &);
+
+} // namespace FlexFlow
+
+#endif

lib/op-attrs/include/op-attrs/operator_space_parallel_tensor_space_mapping.struct.toml

@@ -0,0 +1,17 @@
+namespace = "FlexFlow"
+name = "OperatorSpaceParallelTensorSpaceMapping"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "utils/orthotope/dim_projection.dtg.h",
+  "op-attrs/operator_task_space_dim_idx_t.dtg.h",
+  "op-attrs/parallel_tensor_dim_idx_t.dtg.h",
+]
+
+[[fields]]
+name = "raw_projection"
+type = "::FlexFlow::DimProjection<::FlexFlow::operator_task_space_dim_idx_t, ::FlexFlow::parallel_tensor_dim_idx_t>"

lib/pcg/include/pcg/operator_task_space.h → lib/op-attrs/include/op-attrs/operator_task_space.h

@@ -1,8 +1,8 @@
-#ifndef _FLEXFLOW_PCG_INCLUDE_OPERATOR_TASK_SPACE_H
-#define _FLEXFLOW_PCG_INCLUDE_OPERATOR_TASK_SPACE_H
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_OPERATOR_TASK_SPACE_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_OPERATOR_TASK_SPACE_H
 
-#include "pcg/operator_task_space.dtg.h"
-#include "pcg/task_space_coordinate.dtg.h"
+#include "op-attrs/operator_task_space.dtg.h"
+#include "op-attrs/task_space_coordinate.dtg.h"
 #include <cstddef>
 #include <unordered_set>
 

lib/pcg/include/pcg/operator_task_space.struct.toml → lib/op-attrs/include/op-attrs/operator_task_space.struct.toml

@@ -10,14 +10,9 @@ features = [
 ]
 
 includes = [ 
-  "<vector>",
-]
-
-src_includes = [
-  "utils/fmt/vector.h",
-  "utils/hash/vector.h"
+  "utils/orthotope/orthotope.dtg.h",
 ]
 
 [[fields]]
-name = "degrees"
-type = "std::vector<int>"
+name = "raw_orthotope"
+type = "::FlexFlow::Orthotope"

lib/op-attrs/include/op-attrs/operator_task_space_dim_idx_t.struct.toml

@@ -0,0 +1,13 @@
+namespace = "FlexFlow"
+name = "operator_task_space_dim_idx_t"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+]
+
+[[fields]]
+name = "raw_idx"
+type = "int"

lib/op-attrs/include/op-attrs/ops/linear.h

@@ -4,10 +4,14 @@
 #include "op-attrs/incoming_tensor_role.dtg.h"
 #include "op-attrs/ops/core.h"
 #include "op-attrs/ops/linear_attrs.dtg.h"
+#include "op-attrs/parallel_tensor_dim_degrees.dtg.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
+#include "op-attrs/tensor_num_dims.dtg.h"
 #include "op-attrs/tensor_shape.dtg.h"
 #include "utils/record_formatter.h"
 #include <tl/expected.hpp>
+#include "op-attrs/operator_space_parallel_tensor_space_mapping.dtg.h"
+#include "op-attrs/parallel_tensor_space_mapping.dtg.h"
 
 namespace FlexFlow {
 
@@ -25,6 +29,17 @@ tl::expected<TensorShape, std::string> get_bias_shape(LinearAttrs const &attrs,
 tl::expected<TensorShape, std::string>
     get_output_shape(LinearAttrs const &attrs, TensorShape const &input);
 
+tl::expected<ParallelTensorSpaceMapping, std::string>
+    get_projection_to_output_parallel_dim_mapping(LinearAttrs const &attrs, 
+                                                  ParallelTensorDimDegrees const &input);
+
+tl::expected<ParallelTensorDimDegrees, std::string>
+    get_projection_parallel_dim_degrees(LinearAttrs const &attrs, ParallelTensorDimDegrees const &input);
+tl::expected<ParallelTensorDimDegrees, std::string>
+    get_bias_parallel_dim_degrees(LinearAttrs const &attrs, ParallelTensorDimDegrees const &input);
+tl::expected<ParallelTensorDimDegrees, std::string>
+    get_output_parallel_dim_degrees(LinearAttrs const &attrs, ParallelTensorDimDegrees const &input);
+
 tl::expected<ParallelTensorShape, std::string>
     get_projection_shape(LinearAttrs const &attrs,
                          ParallelTensorShape const &input);
@@ -34,6 +49,16 @@ tl::expected<ParallelTensorShape, std::string>
     get_output_shape(LinearAttrs const &attrs,
                      ParallelTensorShape const &input);
 
+tl::expected<OperatorSpaceParallelTensorSpaceMapping, std::string>
+    get_projection_space_mapping(LinearAttrs const &attrs,
+                                 ParallelTensorDimDegrees const &input);
+tl::expected<OperatorSpaceParallelTensorSpaceMapping, std::string>
+    get_bias_space_mapping(LinearAttrs const &attrs,
+                           ParallelTensorDimDegrees const &input);
+tl::expected<OperatorSpaceParallelTensorSpaceMapping, std::string>
+    get_output_space_mapping(LinearAttrs const &attrs, 
+                             TensorNumDims const &input_num_dims);
+
 } // namespace FlexFlow
 
 #endif

lib/op-attrs/include/op-attrs/parallel_tensor_dim_degrees.h

@@ -0,0 +1,20 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PARALLEL_TENSOR_DIM_DEGREES_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PARALLEL_TENSOR_DIM_DEGREES_H
+
+#include "op-attrs/parallel_tensor_dim_degrees.dtg.h"
+#include "op-attrs/parallel_tensor_dim_idx_t.dtg.h"
+#include "op-attrs/parallel_tensor_space_coordinate.dtg.h"
+
+namespace FlexFlow {
+
+std::set<parallel_tensor_dim_idx_t> get_nontrivial_parallel_tensor_dim_indices(ParallelTensorDimDegrees const &);
+
+std::unordered_map<parallel_tensor_dim_idx_t, int>
+  get_parallel_tensor_degree_map(ParallelTensorDimDegrees const &);
+
+std::unordered_set<ParallelTensorSpaceCoordinate>
+  get_parallel_tensor_space_coordinates(ParallelTensorDimDegrees const &);
+
+} // namespace FlexFlow
+
+#endif

lib/op-attrs/include/op-attrs/parallel_tensor_dim_idx_t.h

@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PARALLEL_TENSOR_DIM_IDX_T_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PARALLEL_TENSOR_DIM_IDX_T_H
+
+#include "op-attrs/parallel_tensor_dim_idx_t.dtg.h"
+
+namespace FlexFlow {
+
+parallel_tensor_dim_idx_t sum_dim_idx();
+parallel_tensor_dim_idx_t discard_copy_dim_idx();
+parallel_tensor_dim_idx_t shard_dim_idx(ff_dim_t);
+
+} // namespace FlexFlow
+
+#endif

lib/op-attrs/include/op-attrs/parallel_tensor_dim_idx_t.variant.toml

@@ -15,6 +15,9 @@ includes = [
 
 [[values]]
 type = "::FlexFlow::ff_dim_t"
+key = "shard_dim"
+
 
 [[values]]
 type = "::FlexFlow::ReplicaType"
+key = "replica_dim"