Fri Dec 27 03:20:32 PM PST 2024

Author: lockshaw

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.cclib/kernels/src/legion_dim.cc renamed to 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/ff_dim.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/ff_dim_t.h

@@ -2,12 +2,12 @@
 #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/parallel_tensor_dim_degrees.dtg.h"
+#include "op-attrs/tensor_num_dims.dtg.h"
 
 namespace FlexFlow {
 
 OperatorSpaceParallelTensorSpaceMapping
-  get_identity_mapping(ParallelTensorDimDegrees const &);
+  get_identity_mapping(TensorNumDims const &);
 
 } // namespace FlexFlow
 

…ttrs/include/op-attrs/ff_dim.struct.toml …rs/include/op-attrs/ff_dim_t.struct.tomllib/op-attrs/include/op-attrs/ff_dim.struct.toml renamed to lib/op-attrs/include/op-attrs/ff_dim_t.struct.toml lib/op-attrs/include/op-attrs/ff_dim.struct.toml renamed to lib/op-attrs/include/op-attrs/ff_dim_t.struct.toml

@@ -6,10 +6,12 @@
 #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 {
 
@@ -27,6 +29,10 @@ 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>
@@ -51,7 +57,7 @@ tl::expected<OperatorSpaceParallelTensorSpaceMapping, std::string>
                            ParallelTensorDimDegrees const &input);
 tl::expected<OperatorSpaceParallelTensorSpaceMapping, std::string>
     get_output_space_mapping(LinearAttrs const &attrs, 
-                             ParallelTensorDimDegrees const &input);
+                             TensorNumDims const &input_num_dims);
 
 } // namespace FlexFlow
 

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_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/ops/linear.h

@@ -0,0 +1,8 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PARALLEL_TENSOR_SPACE_MAPPING_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PARALLEL_TENSOR_SPACE_MAPPING_H
+
+namespace FlexFlow {
+
+} // namespace FlexFlow
+
+#endif

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

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