[DimExpr] Move SimplifyDimExpr, from cinn to shape_dialect (#61133)

jiahy0825 · web-flow · commit 9adcafe47912 · 2024-01-27T16:00:28.000+08:00
* [DimExpr] Move SimplifyDimExpr, from cinn to shape_dialect

* Fix cmake

* Fix compile error

* Fix gcc12 compile error

* Fix gcc12 compile error

* Fix windows compile error

* Add IR_API for windows platform
diff --git a/paddle/cinn/adt/dim_expr_simplifier.h b/paddle/cinn/adt/dim_expr_simplifier.h
diff --git a/paddle/cinn/adt/simplify_value.cc b/paddle/cinn/adt/simplify_value.cc
@@ -15,13 +15,13 @@
 #include <optional>
 #include <typeinfo>
 
-#include "paddle/cinn/adt/dim_expr_simplifier.h"
 #include "paddle/cinn/adt/equation.h"
 #include "paddle/cinn/adt/equation_value_match_trait.h"
 #include "paddle/cinn/adt/get_sub_reshape_dim_ranges.h"
 #include "paddle/cinn/adt/index_expr_infer_context.h"
 #include "paddle/cinn/adt/match.h"
 #include "paddle/cinn/adt/simplify_value.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr_simplify.h"
 
 namespace cinn::adt {
 
@@ -64,7 +64,7 @@ struct SimplifyRedundantBroadcastedIterator {
       return SimplifyValue(outter_iterator, ctx);
     } else {
       const auto& bd = MakeBroadcastedDim(outter_dim, inner_dim);
-      const auto& simplified_bd = DimExpr{SimplifyDimExpr(bd)};
+      const auto& simplified_bd = DimExpr{symbol::SimplifyDimExpr(bd)};
       return BroadcastedIterator<Value, DimExpr>{inner_iterator, simplified_bd};
     }
     LOG(FATAL) << "Dead code";
diff --git a/paddle/cinn/common/CMakeLists.txt b/paddle/cinn/common/CMakeLists.txt
@@ -22,7 +22,6 @@ gather_srcs(
   python_interpreter_guard.cc
   nvgpu_dev_info.cc
   integer_set.cc
-  dim_expr_simplify.cc
   dim_expr_converter.cc
   broadcast_tree.cc
   dim_expr_util.cc)
@@ -51,8 +50,6 @@ if(WITH_CUDA)
 endif()
 if(NOT CINN_ONLY)
   cinn_cc_test(dim_expr_util_test SRCS dim_expr_util_test.cc DEPS cinncore)
-  cinn_cc_test(dim_expr_simplify_test SRCS dim_expr_simplify_test.cc DEPS
-               cinncore)
   cinn_cc_test(dim_expr_converter_test SRCS dim_expr_converter_test.cc DEPS
                cinncore)
   cinn_cc_test(broadcast_tree_test SRCS broadcast_tree_test.cc DEPS cinncore)
diff --git a/paddle/cinn/common/broadcast_tree.cc b/paddle/cinn/common/broadcast_tree.cc
@@ -17,8 +17,8 @@
 #include <optional>
 #include <unordered_map>
 
-#include "paddle/cinn/common/dim_expr_simplify.h"
 #include "paddle/cinn/common/dim_expr_util.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr_simplify.h"
 
 namespace cinn::common {
 
@@ -208,23 +208,26 @@ symbol::DimExpr GetCstrLhsEqRhsDimExpr(
     const symbol::DimExpr& dim_expr) {
   const auto& pattern2replacement =
       ConstructCstrLhsEqRhsReplacement(broadcastable_condition);
-  return SimplifyDimExpr(SubstituteDimExpr(dim_expr, pattern2replacement));
+  return symbol::SimplifyDimExpr(
+      SubstituteDimExpr(dim_expr, pattern2replacement));
 }
 
 symbol::DimExpr GetCstrLhsEqOneDimExpr(
     const symbol::Broadcastable<symbol::DimExpr>& broadcastable_condition,
     const symbol::DimExpr& dim_expr) {
   const auto& pattern2replacement =
       ConstructCstrLhsEqOneReplacement(broadcastable_condition);
-  return SimplifyDimExpr(SubstituteDimExpr(dim_expr, pattern2replacement));
+  return symbol::SimplifyDimExpr(
+      SubstituteDimExpr(dim_expr, pattern2replacement));
 }
 
 symbol::DimExpr GetCstrRhsEqOneDimExpr(
     const symbol::Broadcastable<symbol::DimExpr>& broadcastable_condition,
     const symbol::DimExpr& dim_expr) {
   const auto& pattern2replacement =
       ConstructCstrRhsEqOneReplacement(broadcastable_condition);
-  return SimplifyDimExpr(SubstituteDimExpr(dim_expr, pattern2replacement));
+  return symbol::SimplifyDimExpr(
+      SubstituteDimExpr(dim_expr, pattern2replacement));
 }
 
 typedef symbol::DimExpr (*ConvertDimExprT)(
diff --git a/paddle/cinn/hlir/dialect/operator/ir/manual_op.cc b/paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
@@ -16,7 +16,6 @@
 
 #include <vector>
 #include "glog/logging.h"
-#include "paddle/cinn/common/dim_expr_simplify.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
 #include "paddle/common/ddim.h"
 #include "paddle/common/enforce.h"
@@ -27,6 +26,7 @@
 #include "paddle/pir/core/builtin_type.h"
 #include "paddle/pir/core/op_base.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr_simplify.h"
 
 namespace cinn {
 namespace dialect {
@@ -427,7 +427,7 @@ bool GenerateShapeOp::InferSymbolicShape(
     for (const auto& attr_dim_expr : attr_dim_exprs) {
       const auto& substituted =
           SubstituteDimExpr(attr_dim_expr, DimExprs4SymbolName);
-      const auto& simplified = common::SimplifyDimExpr(substituted);
+      const auto& simplified = symbol::SimplifyDimExpr(substituted);
       dim_exprs.push_back(simplified);
     }
     return dim_exprs;
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.cc
@@ -16,14 +16,14 @@
 
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.h"
 
-#include "paddle/cinn/common/dim_expr_simplify.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
 #include "paddle/cinn/hlir/dialect/runtime/ir/runtime_dialect.h"
 #include "paddle/cinn/runtime/flags.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.h"
 #include "paddle/pir/core/builtin_type.h"
 #include "paddle/pir/dialect/shape/utils/dim_expr.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr_simplify.h"
 #include "paddle/utils/flags.h"
 
 PD_DECLARE_string(cinn_convert_dynamic_dim_to_static_dim);
@@ -168,7 +168,7 @@ class DynamicToStaticConverter {
         shape_analysis_->GetShapeOrDataForValue(value).shape();
     std::vector<std::int64_t> static_shapes{};
     VisitEachDimExpr(dynamic_shapes, [&](const symbol::DimExpr& dim_expr) {
-      const auto& static_shape = cinn::common::SimplifyDimExpr(
+      const auto& static_shape = symbol::SimplifyDimExpr(
           cinn::dialect::SubstituteDimExpr(dim_expr, DimExpr4SymbolName_));
       CHECK(static_shape.Has<std::int64_t>());
       static_shapes.push_back(static_shape.Get<std::int64_t>());
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc
@@ -14,7 +14,6 @@
 
 #include "paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.h"
 
-#include "paddle/cinn/common/dim_expr_simplify.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
@@ -26,6 +25,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/dialect/shape/utils/dim_expr.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr_simplify.h"
 #include "paddle/pir/pass/pass.h"
 #include "paddle/pir/pattern_rewrite/pattern_applicator.h"
 #include "paddle/pir/pattern_rewrite/pattern_match.h"
@@ -339,7 +339,7 @@ class SplitGenerateShapeIntoShapeOps
       CachedDimExprToValueConverter* converter) const {
     std::vector<pir::Value> ret;
     for (const auto& dim_expr : dim_exprs) {
-      const auto& simplified = cinn::common::SimplifyDimExpr(dim_expr);
+      const auto& simplified = symbol::SimplifyDimExpr(dim_expr);
       pir::Value value = converter->ConvertToValue(simplified);
       ret.push_back(value);
     }
diff --git a/paddle/cinn/hlir/pe/map_expr_to_ir.cc b/paddle/cinn/hlir/pe/map_expr_to_ir.cc
@@ -18,7 +18,6 @@
 #include <vector>
 
 #include "paddle/cinn/adt/dim_expr.h"
-#include "paddle/cinn/adt/dim_expr_simplifier.h"
 #include "paddle/cinn/adt/equation_value_match_trait.h"
 #include "paddle/cinn/adt/inline_translator.h"
 #include "paddle/cinn/adt/map_expr.h"
diff --git a/paddle/pir/dialect/shape/utils/dim_expr_simplify.cc b/paddle/pir/dialect/shape/utils/dim_expr_simplify.cc
@@ -12,12 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/cinn/common/dim_expr_simplify.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr_simplify.h"
 #include <numeric>
 
-namespace cinn::common {
-
-using namespace symbol;  // NOLINT
+namespace symbol {
 
 namespace {
 
@@ -202,8 +200,8 @@ struct IsLhsBeforeRhsStruct<std::string, std::string> {
 template <>
 struct IsLhsBeforeRhsStruct<Negative<DimExpr>, Negative<DimExpr>> {
   static bool Call(const Negative<DimExpr>& lhs, const Negative<DimExpr>& rhs) {
-    const auto& [lhs_operand] = *lhs;
-    const auto& [rhs_operand] = *rhs;
+    const auto& lhs_operand = lhs->data;
+    const auto& rhs_operand = rhs->data;
     return IsLhsBeforeRhs(lhs_operand, rhs_operand);
   }
 };
@@ -212,8 +210,8 @@ template <>
 struct IsLhsBeforeRhsStruct<Reciprocal<DimExpr>, Reciprocal<DimExpr>> {
   static bool Call(const Reciprocal<DimExpr>& lhs,
                    const Reciprocal<DimExpr>& rhs) {
-    const auto& [lhs_operand] = *lhs;
-    const auto& [rhs_operand] = *rhs;
+    const auto& lhs_operand = lhs->data;
+    const auto& rhs_operand = rhs->data;
     return IsLhsBeforeRhs(lhs_operand, rhs_operand);
   }
 };
@@ -266,7 +264,7 @@ struct SortOperands {
 
   bool IsSorted(const List<DimExpr>& operands) {
     CHECK(!operands->empty());
-    for (int i = 0; i < operands->size() - 1; ++i) {
+    for (std::size_t i = 0; i < operands->size() - 1; ++i) {
       if (IsLhsBeforeRhs(operands->at(i + 1), operands->at(i))) {
         return false;
       }
@@ -277,7 +275,7 @@ struct SortOperands {
 
 std::int64_t GetInteger(const DimExpr& expr) {
   if (expr.Has<Negative<DimExpr>>()) {
-    const auto& [integer] = *expr.Get<Negative<DimExpr>>();
+    const auto& integer = expr.Get<Negative<DimExpr>>()->data;
     CHECK(integer.Has<std::int64_t>());
     return -integer.Get<std::int64_t>();
   }
@@ -496,7 +494,7 @@ struct FoldOperandTrait<Add> {
       return true;
     }
     if (dim_expr.Has<Negative<DimExpr>>()) {
-      const auto& [operand] = *dim_expr.Get<Negative<DimExpr>>();
+      const auto& operand = dim_expr.Get<Negative<DimExpr>>()->data;
       return operand.Has<std::int64_t>();
     }
     return false;
@@ -520,7 +518,7 @@ struct FoldOperandTrait<Add> {
 
   static bool IsInversedPair(const DimExpr& lhs, const DimExpr& rhs) {
     if (lhs.Has<Negative<DimExpr>>()) {
-      const auto& [lhs_operand] = *lhs.Get<Negative<DimExpr>>();
+      const auto& lhs_operand = lhs.Get<Negative<DimExpr>>()->data;
       return lhs_operand == rhs;
     }
     if (rhs.Has<Negative<DimExpr>>()) {
@@ -539,22 +537,28 @@ ConstRational SimplifiedConstRational(int64_t num, int64_t dem) {
 }
 
 template <typename T>
-ConstRational GetConstRationalImpl(const T& expr) {
+std::optional<ConstRational> GetConstRationalImpl(const T& expr) {
   LOG(FATAL) << "not supported.";
+  return std::nullopt;
 }
 
-ConstRational GetConstRationalImpl(std::int64_t value) {
+std::optional<ConstRational> GetConstRationalImpl(std::int64_t value) {
   return ConstRational{value, 1};
 }
 
-ConstRational GetConstRationalImpl(const Reciprocal<DimExpr>& value) {
+std::optional<ConstRational> GetConstRationalImpl(
+    const Reciprocal<DimExpr>& value) {
   const auto& [denominator] = *value;
   return ConstRational{1, denominator.Get<std::int64_t>()};
 }
 
 ConstRational GetConstRational(const DimExpr& expr) {
   return std::visit(
-      [&](const auto& impl) { return GetConstRationalImpl(impl); },
+      [&](const auto& impl) {
+        std::optional<ConstRational> opt_ret = GetConstRationalImpl(impl);
+        CHECK(opt_ret.has_value());
+        return opt_ret.value();
+      },
       expr.variant());
 }
 
@@ -580,7 +584,7 @@ struct FoldOperandTrait<Mul> {
       return true;
     }
     if (dim_expr.Has<Reciprocal<DimExpr>>()) {
-      const auto& [operand] = *dim_expr.Get<Reciprocal<DimExpr>>();
+      const auto& operand = dim_expr.Get<Reciprocal<DimExpr>>()->data;
       return operand.Has<std::int64_t>();
     }
     return false;
@@ -610,11 +614,11 @@ struct FoldOperandTrait<Mul> {
   }
   static bool IsInversedPair(const DimExpr& lhs, const DimExpr& rhs) {
     if (lhs.Has<Reciprocal<DimExpr>>()) {
-      const auto& [lhs_operand] = *lhs.Get<Reciprocal<DimExpr>>();
+      const auto& lhs_operand = lhs.Get<Reciprocal<DimExpr>>()->data;
       return lhs_operand == rhs;
     }
     if (rhs.Has<Reciprocal<DimExpr>>()) {
-      const auto& [rhs_operand] = *rhs.Get<Reciprocal<DimExpr>>();
+      const auto& rhs_operand = rhs.Get<Reciprocal<DimExpr>>()->data;
       return rhs_operand == lhs;
     }
     return false;
@@ -669,8 +673,8 @@ struct FoldInversedPairToUnit {
   using dim_expr_type = Op<DimExpr>;
 
   struct SearchResult {
-    int value_pos;
-    int inverse_value_pos;
+    std::size_t value_pos;
+    std::size_t inverse_value_pos;
   };
 
   DimExpr Rewrite(const DimExpr& expr) {
@@ -704,8 +708,8 @@ struct FoldInversedPairToUnit {
 
   std::optional<SearchResult> SearchInversedPair(
       const List<DimExpr>& operands) {
-    for (int i = 0; i < operands->size(); ++i) {
-      for (int j = 0; j < operands->size(); ++j) {
+    for (std::size_t i = 0; i < operands->size(); ++i) {
+      for (std::size_t j = 0; j < operands->size(); ++j) {
         if (i == j) {
           continue;
         }
@@ -758,7 +762,8 @@ struct FoldRedundantSymbolicBroadcast {
 
   std::optional<MaxInt64> SearchMaxInt64(const List<DimExpr>& operands) {
     std::optional<MaxInt64> ret;
-    for (int i = 0; i < operands->size(); ++i) {
+    int operands_size = static_cast<int>(operands->size());
+    for (int i = 0; i < operands_size; ++i) {
       const auto& expr = operands->at(i);
       if (!expr.Has<std::int64_t>()) {
         continue;
@@ -788,8 +793,8 @@ struct FoldRedundantBroadcast {
   using dim_expr_type = Broadcast<DimExpr>;
 
   struct SearchResult {
-    int value_pos;
-    int same_value_pos;
+    std::size_t value_pos;
+    std::size_t same_value_pos;
   };
 
   DimExpr Rewrite(const DimExpr& expr) {
@@ -816,8 +821,8 @@ struct FoldRedundantBroadcast {
 
   std::optional<SearchResult> SearchInversedPair(
       const List<DimExpr>& operands) {
-    for (int i = 0; i < operands->size(); ++i) {
-      for (int j = 0; j < operands->size(); ++j) {
+    for (std::size_t i = 0; i < operands->size(); ++i) {
+      for (std::size_t j = 0; j < operands->size(); ++j) {
         if (i == j) {
           continue;
         }
@@ -871,4 +876,4 @@ DimExpr Simplify(const DimExpr& expr) {
 
 DimExpr SimplifyDimExpr(const DimExpr& expr) { return Simplify(expr); }
 
-}  // namespace cinn::common
+}  // namespace symbol
diff --git a/paddle/pir/dialect/shape/utils/dim_expr_simplify.h b/paddle/pir/dialect/shape/utils/dim_expr_simplify.h
@@ -16,8 +16,8 @@
 
 #include "paddle/pir/dialect/shape/utils/dim_expr.h"
 
-namespace cinn::common {
+namespace symbol {
 
-symbol::DimExpr SimplifyDimExpr(const symbol::DimExpr&);
+IR_API DimExpr SimplifyDimExpr(const DimExpr& dim_expr);
 
 }
diff --git a/test/cpp/pir/shape_dialect/CMakeLists.txt b/test/cpp/pir/shape_dialect/CMakeLists.txt
@@ -1,4 +1,5 @@
 paddle_test(symbol_dim_expr_test SRCS symbol_dim_expr_test.cc)
+paddle_test(simplify_dim_expr_test SRCS simplify_dim_expr_test.cc)
 
 if(WITH_CINN)
   paddle_test(shape_analysis_test SRCS shape_analysis_test.cc)
diff --git a/test/cpp/pir/shape_dialect/simplify_dim_expr_test.cc b/test/cpp/pir/shape_dialect/simplify_dim_expr_test.cc

Original file line number	Diff line number	Diff line change
`@@ -16,8 +16,8 @@`
`16`	`16`
`17`	`17`	`#include "paddle/pir/dialect/shape/utils/dim_expr.h"`
`18`	`18`
`19`		`-namespace cinn::common {`
	`19`	`+namespace symbol {`
`20`	`20`
`21`		`-symbol::DimExpr SimplifyDimExpr(const symbol::DimExpr&);`
	`21`	`+IR_API DimExpr SimplifyDimExpr(const DimExpr& dim_expr);`
`22`	`22`
`23`	`23`	`}`