implement get_operator_attrs

wmdi · Sep 8, 2023 · e3b633f · e3b633f
1 parent 21b8549
commit e3b633f
Show file tree

Hide file tree

Showing 10 changed files with 342 additions and 51 deletions.
diff --git a/lib/substitutions/CMakeLists.txt b/lib/substitutions/CMakeLists.txt
@@ -10,6 +10,7 @@ ff_add_library(
   DEPS
     utils
     op-attrs
+    pcg
 )
 
 add_subdirectory(ffi)

diff --git a/lib/substitutions/include/substitutions/operator_pattern.h b/lib/substitutions/include/substitutions/operator_pattern.h
@@ -23,7 +23,7 @@ enum class OperatorAttributeKey {
   STRIDE_W,
   PADDING_H,
   PADDING_W,
-  AGGR_MODE,
+  AGGR,
   NUM_ENTRIES,
   OUT_CHANNELS,
   ACTIVATION,
@@ -42,14 +42,49 @@ enum class OperatorAttributeKey {
   PARALLEL_DIM,
   PARALLEL_DEGREE,
   PAD,
+  EMBED_DIM,
+  KDIM,
+  VDIM,
+  DROPOUT,
+  BIAS,
+  ADD_BIAS_KV,
+  ADD_ZERO_ATTN,
+  A_SEQ_LENGTH_DIM,
+  B_SEQ_LENGTH_DIM,
+  RELU,
+  TARGET_DIMS,
+  RATE,
+  SEED,
+  SHOULD_BROADCAST_LHS,
+  SHOULD_BROADCAST_RHS,
+  DIM,
+  ELEMENTWISE_AFFINE,
+  REGULARIZER,
+  SHAPE,
+  SPLITS,
+  K,
+  SORTED,
 };
 
-using OperatorAttributeValue =
-    variant<int, float, bool, std::vector<int>, OperatorType, Activation>;
+using OperatorAttributeValue = variant<int,
+                                       float,
+                                       bool,
+                                       stack_vector<int, MAX_TENSOR_DIM>,
+                                       OperatorType,
+                                       Activation,
+                                       ff_dim_t,
+                                       unsigned long long,
+                                       AggregateOp,
+                                       stack_vector<ff_dim_t, MAX_TENSOR_DIM>,
+                                       RegularizerAttrs,
+                                       PoolOp,
+                                       TensorShape>;
 
-using OperatorAttributeConstraint = AttributeConstraint<OperatorAttributeKey, OperatorAttributeValue>;
+using OperatorAttributeConstraint =
+    AttributeConstraint<OperatorAttributeKey, OperatorAttributeValue>;
 
-using OperatorPattern = AttributePattern<OperatorAttributeKey, OperatorAttributeValue>;
+using OperatorPattern =
+    AttributePattern<OperatorAttributeKey, OperatorAttributeValue>;
 
 optional<OperatorAttributeValue>
     evaluate_attribute_expr(Operator const &attrs,

diff --git a/lib/substitutions/include/substitutions/output_graph.h b/lib/substitutions/include/substitutions/output_graph.h
@@ -15,8 +15,7 @@ struct AttrConstant {
   OperatorAttributeValue value;
 };
 
-using OperatorAttributeExpr =
-    variant<OperatorAttrAccess, AttrConstant>;
+using OperatorAttributeExpr = variant<OperatorAttrAccess, AttrConstant>;
 
 // NOTE(@wmdi): Not sure if it aligns with other design. Or alternatively we can
 // define the assignment for each operator type.

diff --git a/lib/substitutions/include/substitutions/parallel_tensor_pattern.h b/lib/substitutions/include/substitutions/parallel_tensor_pattern.h
@@ -13,7 +13,8 @@ using TensorAttributeValue = variant<int, std::vector<int>>;
 using TensorAttributeConstraint =
     AttributeConstraint<TensorAttributeKey, TensorAttributeValue>;
 
-using ParallelTensorPattern = AttributePattern<TensorAttributeKey, TensorAttributeValue>;
+using ParallelTensorPattern =
+    AttributePattern<TensorAttributeKey, TensorAttributeValue>;
 
 optional<TensorAttributeValue>
     evaluate_attribute_expr(ParallelTensor const &tensor_shape,

diff --git a/lib/substitutions/src/graph_pattern.cc b/lib/substitutions/src/graph_pattern.cc
@@ -9,23 +9,45 @@
 
 namespace FlexFlow {
 
-template <typename T, typename V>
-optional<V> evaluate_list_index_access(ListIndexAccess<T> const &index_access,
-                                       optional<V> const &v) {
+optional<OperatorAttributeValue> evaluate_list_index_access(int index,
+                                       optional<OperatorAttributeValue> const &v) {
+  if (!v.has_value() ||
+      !holds_alternative<stack_vector<int, MAX_TENSOR_DIM>>(v.value()) ||
+      !holds_alternative<stack_vector<ff_dim_t, MAX_TENSOR_DIM>>(v.value())) {
+    return nullopt;
+  }
+
+  if (index >= MAX_TENSOR_DIM) {
+    return nullopt;
+  }
+
+  if (holds_alternative<stack_vector<int, MAX_TENSOR_DIM>>(v.value())) {
+    return get<stack_vector<int, MAX_TENSOR_DIM>>(v.value()).at(index);
+  } else {
+    return get<stack_vector<ff_dim_t, MAX_TENSOR_DIM>>(v.value()).at(index);
+  }
+}
+
+optional<TensorAttributeValue> evaluate_list_index_access(int const &index,
+                                       optional<TensorAttributeValue> const &v) {
   if (!v.has_value() || !holds_alternative<std::vector<int>>(v.value())) {
     return nullopt;
   }
 
   auto vec = get<std::vector<int>>(v.value());
-  if (index_access.index >= vec.size()) {
+
+  if (index >= vec.size()) {
     return nullopt;
   }
 
-  return vec.at(index_access.index);
+  return vec.at(index);
 }
 
-template <typename V>
-optional<V> evaluate_list_size(optional<V> const &v) {
+optional<OperatorAttributeValue> evaluate_list_size(optional<OperatorAttributeValue> const &v) {
+  return MAX_TENSOR_DIM;
+}
+
+optional<TensorAttributeValue> evaluate_list_size(optional<TensorAttributeValue> const &v) {
   if (!v.has_value() || !holds_alternative<std::vector<int>>(v.value())) {
     return nullopt;
   }
@@ -44,7 +66,7 @@ struct EvaluateOperatorAttributeExpr {
       operator()(ListIndexAccess<OperatorAttributeKey> const &index_access) {
     optional<OperatorAttributeValue> v =
         get_attribute(this->attrs, index_access.attribute_key);
-    return evaluate_list_index_access(index_access, v);
+    return evaluate_list_index_access(index_access.index, v);
   }
 
   optional<OperatorAttributeValue>
@@ -94,8 +116,9 @@ struct EvaluateTensorAttributeExpr {
 
   optional<TensorAttributeValue>
       operator()(ListIndexAccess<TensorAttributeKey> const &index_access) {
-    auto v = this->evaluate(index_access.attribute_key);
-    return evaluate_list_index_access(index_access, v);
+    optional<TensorAttributeValue> v =
+        this->evaluate(index_access.attribute_key);
+    return evaluate_list_index_access(index_access.index, v);
   }
 
   optional<TensorAttributeValue>
@@ -184,6 +207,13 @@ optional<bool> satisfies(ParallelTensor const &params,
       [&](TensorAttributeConstraint const &c) { return satisfies(params, c); });
 }
 
+struct AlwaysTrueCriterion {
+  template <typename T>
+  bool operator()(T const &t) const {
+    return true;
+  }
+};
+
 bool assignment_satisfies(ParallelComputationGraph const &pcg,
                           GraphPattern const &pattern,
                           MultiDiGraphPatternMatch const &patternMatch) {
@@ -192,19 +222,19 @@ bool assignment_satisfies(ParallelComputationGraph const &pcg,
     auto patternNode = kv.first;
     auto pcgNode = kv.second;
     optional<bool> constraintResult =
-        satisfies(pcg->at(pcgNode), pattern->at(patternNode));
+        satisfies(pcg.value().at(pcgNode), pattern.value().at(patternNode));
     result &= constraintResult.value_or(false);
   }
 
   for (auto const &kv : patternMatch.edge_assignment) {
     auto patternEdge = kv.first;
     auto pcgEdge = kv.second;
     optional<bool> constraintResult =
-        satisfies(pcg->at(pcgEdge), pattern->at(patternEdge));
+        satisfies(pcg.value().at(pcgEdge), pattern.value().at(patternEdge));
     result &= constraintResult.value_or(false);
   }
 
-  result &= pattern_matches(pattern, pcg, patternMatch);
+  result &= pattern_matches(pattern, pcg, patternMatch, AlwaysTrueCriterion{});
 
   return result;
 }