【CINN】Simplify dynamic shape (#70933)

liuruyan · web-flow · commit a6c9bd4e86c3 · 2025-01-23T15:12:30.000+08:00
* simplify dynamic shape

* fix ci
diff --git a/paddle/cinn/common/ir_util.cc b/paddle/cinn/common/ir_util.cc
@@ -25,6 +25,7 @@
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
 #include "paddle/cinn/ir/utils/ir_compare.h"
+#include "paddle/cinn/ir/utils/ir_copy.h"
 #include "paddle/common/enforce.h"
 namespace cinn {
 namespace common {
@@ -719,12 +720,16 @@ ir::IndexExpr ConstructIndexExprByNodeType(const ir::IrNodeTy &ty,
 ir::IndexExpr ChangeSeqOfDivMod(const ir::IndexExpr &expr) {
   switch (expr.node_type()) {
     case ir::IrNodeTy::IntImm:
-    case ir::IrNodeTy::_Var_: {
+    case ir::IrNodeTy::_Var_:
+    case ir::IrNodeTy::Cast:
+    case ir::IrNodeTy::Load: {
       return expr;
     }
     case ir::IrNodeTy::Add:
     case ir::IrNodeTy::Sub:
     case ir::IrNodeTy::Mul:
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max:
     case ir::IrNodeTy::Div: {
       auto lhs = ChangeSeqOfDivMod(expr.operand(0));
       auto rhs = ChangeSeqOfDivMod(expr.operand(1));
@@ -739,6 +744,9 @@ ir::IndexExpr ChangeSeqOfDivMod(const ir::IndexExpr &expr) {
       } else {
         auto lhs = ChangeSeqOfDivMod(expr.operand(0));
         auto rhs = ChangeSeqOfDivMod(expr.operand(1));
+        if (lhs.node_type() == ir::IrNodeTy::Div) {
+          return (lhs.operand(0) % (lhs.operand(1) * rhs)) / lhs.operand(1);
+        }
         return ConstructIndexExprByNodeType(expr.node_type(), lhs, rhs, false);
       }
     }
@@ -747,5 +755,60 @@ ir::IndexExpr ChangeSeqOfDivMod(const ir::IndexExpr &expr) {
           "Unsupported type of expr in ChangeSeqOfDivMod which is: %s", expr));
   }
 }
+std::optional<ir::IndexExpr> DivByPartMul(const ir::IndexExpr &lhs,
+                                          const ir::IndexExpr &rhs,
+                                          ir::IrNodeTy ty) {
+  std::vector<ir::IndexExpr> elems = common::GetFlattenExprs<ir::Mul>(rhs);
+
+  ir::IndexExpr result = ir::ir_utils::IRCopy(lhs);
+
+  for (const auto &elem : elems) {
+    if (common::IsDivisiblieBySymbol(result, elem, ty)) {
+      result = common::SimplifySymbolicDivide(result, elem, ty);
+    } else {
+      return std::nullopt;
+    }
+  }
+  return result;
+}
+
+std::optional<ir::IndexExpr> SimplifyComplexMod(const ir::IndexExpr &lhs,
+                                                const ir::IndexExpr &rhs) {
+  if (lhs == rhs) return ir::IndexExpr(lhs.type(), 0);
+  switch (lhs.node_type()) {
+    case ir::IrNodeTy::Add: {
+      auto simplify_lhs = SimplifyComplexMod(lhs.operand(0), rhs);
+      auto simplify_rhs = SimplifyComplexMod(lhs.operand(1), rhs);
+      if (simplify_lhs.has_value() && simplify_rhs.has_value())
+        return (simplify_lhs.value() + simplify_rhs.value());
+      return std::nullopt;
+    }
+    case ir::IrNodeTy::Mul: {
+      // (S0 % 4 * S1 % 8) % 4 != S0 % 4 * S1 % 4;
+      if (DivByPartMul(lhs, rhs, ir::IrNodeTy::Mod))
+        return ir::IndexExpr(lhs.type(), 0);
+      return std::nullopt;
+    }
+    case ir::IrNodeTy::Div:
+    case ir::IrNodeTy::IntImm:
+    case ir::IrNodeTy::_Var_:
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max:
+    case ir::IrNodeTy::Load:
+    case ir::IrNodeTy::Cast: {
+      return std::nullopt;
+    }
+    case ir::IrNodeTy::Mod: {
+      if (DivByPartMul(lhs.operand(1), rhs, ir::IrNodeTy::Mod)) {
+        return lhs.operand(0) % rhs;
+      }
+      return std::nullopt;
+    }
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type of expr in SimplifyComplexMod which is: %s", lhs));
+  }
+  return std::nullopt;
+}
 }  // namespace common
 }  // namespace cinn
diff --git a/paddle/cinn/common/ir_util.h b/paddle/cinn/common/ir_util.h
@@ -381,5 +381,38 @@ enum IndexType {
  * information in some scenarios.
  */
 IndexType VerifyIndex(const ir::Expr &expr);
+
+/*!
+ * \brief The multiplication in rhs is broken down and each sub-part is
+ * independently determined to be divisible.
+ * \param lhs The dividend.
+ * \param rhs The divisor.
+ * \param ty  ty is `Mod` or `Div`.
+ * \return A optional index expression indicating whether the `lhs`
+ * is divisible, nullopt indicating not divisible.
+ *
+ * For example:
+ * 1. i * S0 * S1 * S2 / (S0 * S1) ==> i / S2
+ * 2. i * S0 * S1 / S0 ==> i * S1
+ * 3. i * S0 / (S0 + 1) ==> nullopt
+ */
+std::optional<ir::IndexExpr> DivByPartMul(const ir::IndexExpr &lhs,
+                                          const ir::IndexExpr &rhs,
+                                          ir::IrNodeTy ty);
+
+/*!
+ * \brief Simplify complex modulo expressions.
+ * \param lhs The dividend.
+ * \param rhs The divisor.
+ * \return A optional index expression indicating whether simplified
+ *
+ * For example:
+ * 1. (i / S0 * S0 + i % (S0 * S1)) % S0 ==> i % S0
+ * 2. (i / S0 * S0 * S1 + i % (S0 * S1 * S2)) % (S0 * S1) ==> i % (S0 * S1)
+ * 3. i % (S0 * S1) % S0 ==> i % S0
+ * 4. i * S0 * S1 % (S0 * S1) ==> 0
+ */
+std::optional<ir::IndexExpr> SimplifyComplexMod(const ir::IndexExpr &lhs,
+                                                const ir::IndexExpr &rhs);
 }  // namespace common
 }  // namespace cinn
diff --git a/paddle/cinn/common/simplify_special_pattern.cc b/paddle/cinn/common/simplify_special_pattern.cc
@@ -92,8 +92,10 @@ static std::optional<ir::IndexExpr> MergeMulModInner(
                                 mod_r_expr);
       if (overall_mult == inner_div_ptr->b().as_index() &&
           overall_mult == mod_r_expr &&
-          ProveDivisible(inner_div_ptr->a().as_index() - mod_l_expr,
-                         mod_r_expr)) {
+          (ProveDivisible(inner_div_ptr->a().as_index() - mod_l_expr,
+                          mod_r_expr) ||
+           inner_div_ptr->a().as_index() % overall_mult ==
+               mod_l_expr % mod_r_expr)) {
         // Found!
         return no_opt_sum.get()
                    ? no_opt_sum * mult_outer + inner_div_ptr->a().as_index()
@@ -211,10 +213,13 @@ std::optional<ir::IndexExpr> DivMulAddModCornerCase(const ir::IndexExpr& lhs,
   auto innerDiv = inner.As<ir::Div>();
   if (!innerDiv) return std::nullopt;
   if (innerDiv->b().as_index() == rhsMod->b().as_index() &&
-      innerDiv->b().as_index() == mult_outer &&
-      ProveDivisible(rhsMod->a().as_index() - innerDiv->a().as_index(),
-                     mult_outer)) {
-    return innerDiv->a().as_index();
+      innerDiv->b().as_index() == mult_outer) {
+    // The second condition is to adapt to the dynamic shape:
+    // f % (S0 * S1) / S0 * S0 + f % S0 ==> f % (S0 * S1)
+    if (ProveDivisible(rhsMod->a().as_index() - innerDiv->a().as_index(),
+                       mult_outer) ||
+        innerDiv->a().as_index() % mult_outer == rhs)
+      return innerDiv->a().as_index();
   }
   return std::nullopt;
 }
@@ -263,6 +268,8 @@ std::optional<ir::IndexExpr> SubModCornerCase(const ir::IndexExpr& lhs,
                                               bool isDiv) {
   auto flatten = GetFlattenExprs<ir::Add>(lhs);
 
+  if (flatten.size() < 2) return std::nullopt;
+
   for (int64_t i = 0, e = flatten.size(); i < e; ++i) {
     // Check if negation
     ir::IndexExpr beforeNegation = flatten[i];
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
@@ -536,7 +536,9 @@ IndexExpr Simplify(const IndexExpr &expr, IndexExpr::OptLevel level) {
 }
 
 IndexExpr IndexExpr::Normalize(OptLevel level) const {
-  return Simplify(*this, level);
+  auto res = Simplify(*this, level);
+  res = ChangeSeqOfDivMod(res);
+  return Simplify(res, level);
 }
 
 int32_t IndexExpr::as_int32() const {
diff --git a/paddle/cinn/ir/op/ir_operators.cc b/paddle/cinn/ir/op/ir_operators.cc
@@ -533,24 +533,17 @@ static IndexExpr SimplifyDiv(const IndexExpr &lhs, const IndexExpr &rhs) {
         }
       }
     }
-
-    // S0 / 2 / 5 ===> S0 / 10.
-    if (lhsDiv) {
-      if (auto lrhs = lhsDiv->b().as_index().As<IntImm>()) {
-        return lhsDiv->a().as_index() /
-               IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);
-      }
-    }
   } else {
     // dynamic branch!
-    if (common::IsDivisiblieBySymbol(lhs, rhs, ir::IrNodeTy::Div)) {
-      return common::SimplifySymbolicDivide(lhs, rhs, ir::IrNodeTy::Div);
+    if (auto res = DivByPartMul(lhs, rhs, ir::IrNodeTy::Div)) {
+      return res.value();
     }
+  }
 
-    // TODO(liujinnan): Deal dynamic shape, e.g. S0 / S1 / S2 ===> S0 / (S1 *
-    // S2). if (auto lhsDiv = lhs.As<Div>()) {
-    //   return lhsDiv->a().as_index() / (lhsDiv->b().as_index() * rhs);
-    // }
+  // static and dynamic common branch!
+  // S0 / S1 / 2 ===> S0 / (S1 * 2).
+  if (auto lhsDiv = lhs.As<Div>()) {
+    return lhsDiv->a().as_index() / (lhsDiv->b().as_index() * rhs);
   }
 
   return Div::Make(lhs, rhs);
@@ -593,8 +586,14 @@ static IndexExpr SimplifyMod(const IndexExpr &lhs, const IndexExpr &rhs) {
     }
   } else {
     // dynamic branch!
-    if (common::IsDivisiblieBySymbol(lhs, rhs, ir::IrNodeTy::Mod))
+    if (auto res = DivByPartMul(lhs, rhs, ir::IrNodeTy::Mod)) {
       return IndexExpr(0);
+    }
+  }
+
+  // static and dynamic common branch!
+  if (auto res = SimplifyComplexMod(lhs, rhs)) {
+    return res.value();
   }
 
   return Mod::Make(lhs, rhs);
diff --git a/test/cpp/pir/cinn/adt/index_expr_test.cc b/test/cpp/pir/cinn/adt/index_expr_test.cc
@@ -34,9 +34,15 @@ class TestIndexExpr : public ::testing::Test {
              .set_index(true);
     S7 = ir::Var(ir::Expr(static_cast<int64_t>(1)), ir::Expr(INT32_MAX), "S7")
              .set_index(true);
+    S8 = ir::Var(ir::Expr(static_cast<int64_t>(1)), ir::Expr(INT32_MAX), "S8")
+             .set_index(true);
+    S9 = ir::Var(ir::Expr(static_cast<int64_t>(1)), ir::Expr(INT32_MAX), "S9")
+             .set_index(true);
+
+    f = ir::Var(ir::Expr(static_cast<int64_t>(1)), ir::Expr(INT32_MAX), "f");
   };
 
-  ir::Var S4, S5, S6, S7;
+  ir::Var S4, S5, S6, S7, S8, S9, f;
 };
 TEST_F(TestIndexExpr, IndexExpr_0) {
   ir::IndexExpr a(14);
@@ -166,7 +172,8 @@ TEST_F(TestIndexExpr, IndexExpr_3) {
   EXPECT_EQ(q14.as_index().Normalize(), ir::IndexExpr(S4 + S5));
   EXPECT_EQ(q15.as_index().Normalize(),
             ir::IndexExpr((S4 * 256 + S5 + S6 * 1024)) % 25088);
-  EXPECT_EQ(q16.as_index().Normalize(), ir::IndexExpr(S4 * 256 + S5));
+  EXPECT_EQ(q16.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+            ir::IndexExpr(S4 * 256 + S5));
 }
 
 TEST_F(TestIndexExpr, Change_Seq_Of_Div_Mod) {
@@ -204,5 +211,75 @@ TEST_F(TestIndexExpr, Test_ConstructIndexExprByNodeType) {
   EXPECT_EQ(result_min, ir::Min::Make(S4, S5));
   EXPECT_EQ(result_max, ir::Max::Make(S4, S5));
 }
+
+TEST_F(TestIndexExpr, Test_dynamic) {
+  ir::Expr q =
+      ((((((((((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) % S5) * S6) +
+                ((((S7 * 1024) + S8) + (S9 * 4096)) % S6)) +
+               (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) % 640) %
+                  S4) *
+                 S6) *
+                S5)) +
+              (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) / 640) *
+                 S5) *
+                S6) *
+               S4)) /
+             ((S5 * S6) * S4)) *
+            S4) +
+           (((((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) % S5) * S6) +
+                ((((S7 * 1024) + S8) + (S9 * 4096)) % S6)) +
+               (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) % 640) %
+                  S4) *
+                 S6) *
+                S5)) +
+              (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) / 640) *
+                 S5) *
+                S6) *
+               S4)) /
+             (S5 * S6)) %
+            S4)) *
+          S5) +
+         (((((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) % S5) * S6) +
+              ((((S7 * 1024) + S8) + (S9 * 4096)) % S6)) +
+             (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) % 640) % S4) *
+               S6) *
+              S5)) +
+            (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) / 640) * S5) *
+              S6) *
+             S4)) /
+           S6) %
+          S5)) *
+        S6) +
+       ((((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) % S5) * S6) +
+           ((((S7 * 1024) + S8) + (S9 * 4096)) % S6)) +
+          (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) % 640) % S4) *
+            S6) *
+           S5)) +
+         (((((((((S7 * 1024) + S8) + (S9 * 4096)) / S6) / S5) / 640) * S5) *
+           S6) *
+          S4)) %
+        S6));
+
+  ir::Expr q1 =
+      ((((((f % ((S5 * S6) * 640)) % ((S5 * S6) * S4)) / (S5 * S6)) * S6) *
+        S5) +
+       (f % (S5 * S6)));
+  ir::Expr q2 = ((f % ((S5 * S6) * 640)) % ((S5 * S6) * S4)) % (S5 * S6);
+  ir::Expr q3 = (S5 * S6) * S4 / (S5 * S6);
+  ir::Expr q4 = (S5 * S6) * S4 % (S5 * S6);
+
+  EXPECT_EQ(
+      q.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+      ((((((((S7 * 1024) + S8) + (S9 * 4096)) / ((S5 * S6) * 640)) * S5) * S6) *
+        S4) +
+       (((((S7 * 1024) + S8) + (S9 * 4096)) % ((S5 * S6) * 640)) %
+        ((S5 * S6) * S4))));
+  EXPECT_EQ(q1.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+            ((f % ((S5 * S6) * 640)) % ((S5 * S6) * S4)));
+  EXPECT_EQ(q2.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+            f % (S5 * S6));
+  EXPECT_EQ(q3.as_index().Normalize(ir::IndexExpr::OptLevel::Level2), Expr(S4));
+  EXPECT_EQ(q4.as_index().Normalize(ir::IndexExpr::OptLevel::Level2), Expr(0));
+}
 }  // namespace common
 }  // namespace cinn

Original file line number	Diff line number	Diff line change
`@@ -536,7 +536,9 @@ IndexExpr Simplify(const IndexExpr &expr, IndexExpr::OptLevel level) {`
`536`	`536`	`}`
`537`	`537`
`538`	`538`	`IndexExpr IndexExpr::Normalize(OptLevel level) const {`
`539`		`- return Simplify(*this, level);`
	`539`	`+ auto res = Simplify(*this, level);`
	`540`	`+ res = ChangeSeqOfDivMod(res);`
	`541`	`+ return Simplify(res, level);`
`540`	`542`	`}`
`541`	`543`
`542`	`544`	`int32_t IndexExpr::as_int32() const {`