[CINN] Fix dimexpr order check in SortOperands (#72132)

huangjiyi · web-flow · commit 461ad02bfae9 · 2025-04-09T13:03:03.000+08:00
diff --git a/paddle/pir/src/dialect/shape/utils/dim_expr_util.cc b/paddle/pir/src/dialect/shape/utils/dim_expr_util.cc
@@ -246,7 +246,7 @@ struct IsListLhsBeforeListRhsStruct {
       if (lhs_operands->at(i) == rhs_operands->at(i)) continue;
       return IsLhsBeforeRhs(lhs_operands->at(i), rhs_operands->at(i));
     }
-    return true;
+    return false;
   }
 };
 
@@ -1255,6 +1255,9 @@ void DoPass(bool* rewritten, DimExpr* expr) {
 }
 
 DimExpr Simplify(const DimExpr& expr) {
+  if (expr.isa<std::int64_t>() || expr.isa<std::string>()) {
+    return expr;
+  }
   DimExpr ret = expr;
   for (bool keep_rewrite = true; keep_rewrite;) {
     keep_rewrite = false;
diff --git a/test/cpp/pir/shape_dialect/simplify_dim_expr_test.cc b/test/cpp/pir/shape_dialect/simplify_dim_expr_test.cc
@@ -290,4 +290,19 @@ TEST(Simplify, Case2) {
   ASSERT_TRUE((SimplifyDimExpr(dim_expr)) == expected);
 }
 
+TEST(Simplify, Case3) {
+  DimExpr S3{"S3"};
+  DimExpr S4{"S4"};
+  DimExpr S5{"S5"};
+  DimExpr S7{"S7"};
+  DimExpr S8{"S8"};
+  DimExpr dim_expr = Mul<DimExpr>{List<DimExpr>{
+      Div<DimExpr>{Mul<DimExpr>{List<DimExpr>{S3, S4, S5}},
+                   Mul<DimExpr>{List<DimExpr>{S7, S8}}},
+      Div<DimExpr>{Mul<DimExpr>{List<DimExpr>{S3, S4, S5}},
+                   Div<DimExpr>{Mul<DimExpr>{List<DimExpr>{S3, S4, S5}},
+                                Mul<DimExpr>{List<DimExpr>{S7, S8}}}}}};
+  ASSERT_TRUE((SimplifyDimExpr(dim_expr) == dim_expr));  // Need to simplify
+}
+
 }  // namespace symbol::test

Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,7 @@ struct IsListLhsBeforeListRhsStruct {`
`246`	`246`	`if (lhs_operands->at(i) == rhs_operands->at(i)) continue;`
`247`	`247`	`return IsLhsBeforeRhs(lhs_operands->at(i), rhs_operands->at(i));`
`248`	`248`	`}`
`249`		`- return true;`
	`249`	`+ return false;`
`250`	`250`	`}`
`251`	`251`	`};`
`252`	`252`
`@@ -1255,6 +1255,9 @@ void DoPass(bool* rewritten, DimExpr* expr) {`
`1255`	`1255`	`}`
`1256`	`1256`
`1257`	`1257`	`DimExpr Simplify(const DimExpr& expr) {`
	`1258`	`+ if (expr.isa<std::int64_t>() \|\| expr.isa<std::string>()) {`
	`1259`	`+ return expr;`
	`1260`	`+ }`
`1258`	`1261`	`DimExpr ret = expr;`
`1259`	`1262`	`for (bool keep_rewrite = true; keep_rewrite;) {`
`1260`	`1263`	`keep_rewrite = false;`