[CINN]improve shape constraints in lowering group (#70436)

Hongqing-work · web-flow · commit 9d5cd6aae569 · 2024-12-27T16:39:36.000+08:00
* [CINN]improve shape constraints in lowering group

* fix comment
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/collect_sym_expr.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/collect_sym_expr.cc
@@ -235,6 +235,7 @@ CreateGroupShapeOrDataExprs(
 
   local_shape_analysis.RegisterSymbolConstraintFromShapeAnalysis(
       global_shape_analysis);
+  local_shape_analysis.SubstituteInConstraint(dim_expr_map);
   for (const auto& item : dim_expr_map) {
     local_shape_analysis.AddEqualCstr(item.first, item.second);
   }
diff --git a/paddle/pir/include/dialect/shape/utils/constraints_manager.h b/paddle/pir/include/dialect/shape/utils/constraints_manager.h
@@ -99,7 +99,6 @@ class IR_API ConstraintsManager {
 
   const InputRangeConstraints& input_ranges() const { return input_ranges_; }
 
- private:
   void SubstituteInConstraint(const DimExpr& lhs, const DimExpr& rhs);
 
  private:
diff --git a/paddle/pir/include/dialect/shape/utils/shape_analysis.h b/paddle/pir/include/dialect/shape/utils/shape_analysis.h
@@ -125,6 +125,10 @@ class IR_API InferSymbolicShapeContext {
   bool IsBroadcastable(const symbol::DimExpr& lhs,
                        const symbol::DimExpr& rhs) const;
 
+  void SubstituteInConstraint(
+      const std::unordered_map<symbol::DimExpr, symbol::DimExpr>&
+          substitution_pattern);
+
   bool HasPredefinedRange(const symbol::DimExpr& dim_expr) const;
 
   void PrintShapeOrDatas() const;
@@ -215,6 +219,10 @@ class IR_API ShapeConstraintIRAnalysis final
   bool IsBroadcastable(const symbol::DimExpr& lhs,
                        const symbol::DimExpr& rhs) const;
 
+  void SubstituteInConstraint(
+      const std::unordered_map<symbol::DimExpr, symbol::DimExpr>&
+          substitution_pattern);
+
   // Used to debug
   void PrintShapeOrDatas() const;
 
diff --git a/paddle/pir/src/dialect/shape/utils/dim_expr_util.cc b/paddle/pir/src/dialect/shape/utils/dim_expr_util.cc
@@ -1251,6 +1251,15 @@ IR_API PriorityComparisonStatus CompareDimExprPriority(const DimExpr& lhs,
                                        : PriorityComparisonStatus::LOWER;
   }
 
+  const auto& HasSubset = [&](const DimExpr& lhs, const DimExpr& rhs) {
+    std::unordered_map<DimExpr, DimExpr> subsitute_patterns;
+    subsitute_patterns[rhs] = DimExpr{"NS"};
+    if (SubstituteDimExpr(lhs, subsitute_patterns) != lhs) return true;
+    return false;
+  };
+  if (HasSubset(lhs, rhs)) return PriorityComparisonStatus::LOWER;
+  if (HasSubset(rhs, lhs)) return PriorityComparisonStatus::HIGHER;
+
   auto CompareForEqualPriority = common::Overloaded{
       [](const std::string& lhs, const std::string& rhs) {
         if (lhs.size() != rhs.size()) {
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -278,6 +278,21 @@ bool InferSymbolicShapeContext::IsBroadcastable(
   return constraints_manager_.IsBroadcastable(lhs, rhs);
 }
 
+void InferSymbolicShapeContext::SubstituteInConstraint(
+    const std::unordered_map<symbol::DimExpr, symbol::DimExpr>&
+        substitution_pattern) {
+  for (const auto& item : substitution_pattern) {
+    constraints_manager_.SubstituteInConstraint(item.first, item.second);
+  }
+  std::unordered_map<symbol::DimExpr, symbol::DimExpr> new_substitution_pattern;
+  for (const auto& old_pattern : substitution_pattern_) {
+    new_substitution_pattern.emplace(
+        symbol::SubstituteDimExpr(old_pattern.first, substitution_pattern),
+        symbol::SubstituteDimExpr(old_pattern.second, substitution_pattern));
+  }
+  substitution_pattern_ = new_substitution_pattern;
+}
+
 bool InferSymbolicShapeContext::HasPredefinedRange(
     const symbol::DimExpr& dim_expr) const {
   return constraints_manager_.IsBoundedInput(dim_expr);
@@ -361,26 +376,8 @@ InferSymbolicShapeContext::SimplifyBroadcastForShapeOrData(
       });
 }
 
-namespace {
-
-bool CanSubstituteInShapeAnalysis(const symbol::DimExpr& lhs,
-                                  const symbol::DimExpr& rhs) {
-  auto CanSubstitutePredictor = ::common::Overloaded{
-      [](std::int64_t lhs, const auto& rhs) { return true; },
-      [](const std::string& lhs, const std::string& rhs) { return true; },
-      [](const std::string& lhs,
-         const symbol::Broadcast<symbol::DimExpr>& rhs) { return true; },
-      [](const auto& lhs, const auto& rhs) { return false; }};
-  return std::visit(CanSubstitutePredictor, lhs.variant(), rhs.variant()) ||
-         std::visit(CanSubstitutePredictor, rhs.variant(), lhs.variant());
-}
-
-}  // namespace
-
 void InferSymbolicShapeContext::SubstituteDimExpr(
     const symbol::DimExpr& origin, const symbol::DimExpr& substituted) {
-  if (!CanSubstituteInShapeAnalysis(origin, substituted)) return;
-
   substitution_pattern_[origin] = substituted;
   for (auto& val : substitution_pattern_) {
     if (val.second == origin) {
@@ -682,6 +679,12 @@ bool ShapeConstraintIRAnalysis::IsBroadcastable(
   return context_.IsBroadcastable(lhs, rhs);
 }
 
+void ShapeConstraintIRAnalysis::SubstituteInConstraint(
+    const std::unordered_map<symbol::DimExpr, symbol::DimExpr>&
+        substitution_pattern) {
+  context_.SubstituteInConstraint(substitution_pattern);
+}
+
 void ShapeConstraintIRAnalysis::PrintShapeOrDatas() const {
   context_.PrintShapeOrDatas();
 }

Original file line number	Diff line number	Diff line change
`@@ -235,6 +235,7 @@ CreateGroupShapeOrDataExprs(`
`235`	`235`
`236`	`236`	`local_shape_analysis.RegisterSymbolConstraintFromShapeAnalysis(`
`237`	`237`	`global_shape_analysis);`
	`238`	`+ local_shape_analysis.SubstituteInConstraint(dim_expr_map);`
`238`	`239`	`for (const auto& item : dim_expr_map) {`
`239`	`240`	`local_shape_analysis.AddEqualCstr(item.first, item.second);`
`240`	`241`	`}`