【CINN】Merge dynamic expr by multiplication distributive law (#71018)

liuruyan · web-flow · commit 6821f2e225be · 2025-02-07T17:27:00.000+08:00
* MergeExpr by common factor

* simplify common factor

* add unit test
diff --git a/paddle/cinn/ir/op/ir_operators.cc b/paddle/cinn/ir/op/ir_operators.cc
@@ -435,6 +435,7 @@ static IndexExpr SimplifyAdd(const IndexExpr &lhs, const IndexExpr &rhs) {
     // dynamic branch!
     if (common::IsSumPartialBySymbol(lhs, rhs))
       return common::SimplifySymbolicAdd(lhs, rhs);
+
     if (auto rhs_mul = rhs.As<ir::Mul>()) {
       if (rhs_mul->b().as_index().is_constant()) {
         if (common::IsSumPartialBySymbol(lhs, rhs_mul->a().as_index())) {
@@ -443,6 +444,66 @@ static IndexExpr SimplifyAdd(const IndexExpr &lhs, const IndexExpr &rhs) {
         }
       }
     }
+
+    // (S0 * S1 * S2) + (S2 * S1 * S3) ==> (S0 + S3) * (S1 * S2)
+    auto MergeByCommonFactor =
+        [](const ir::IndexExpr &lhs,
+           const ir::IndexExpr &rhs) -> std::optional<ir::IndexExpr> {
+      auto flatten_mul_lhs = common::GetFlattenExprs<ir::Mul>(lhs);
+      auto flatten_mul_rhs = common::GetFlattenExprs<ir::Mul>(rhs);
+
+      if (flatten_mul_lhs.size() > 1 && flatten_mul_rhs.size() > 1) {
+        ir::IndexExpr common_factor(lhs.type(), 1);
+        std::unordered_map<ir::IndexExpr, int> lhs_count;
+        std::unordered_map<ir::IndexExpr, int> rhs_count;
+
+        for (auto &l : flatten_mul_lhs) lhs_count[l]++;
+        for (auto &r : flatten_mul_rhs) rhs_count[r]++;
+        // Find common factor
+        for (auto &l : flatten_mul_lhs) {
+          if (rhs_count[l] > 0) {
+            common_factor = l * common_factor;
+            rhs_count[l]--;
+            lhs_count[l]--;
+          }
+        }
+        // Find Lhs remainder
+        ir::IndexExpr lhs_remainder(lhs.type(), 1);
+        for (auto &l : flatten_mul_lhs) {
+          while (lhs_count[l] > 0) {
+            lhs_remainder = l * lhs_remainder;
+            lhs_count[l]--;
+          }
+        }
+        // Find Rhs remainder
+        ir::IndexExpr rhs_remainder(rhs.type(), 1);
+        for (auto &r : flatten_mul_rhs) {
+          while (rhs_count[r] > 0) {
+            rhs_remainder = r * rhs_remainder;
+            rhs_count[r]--;
+          }
+        }
+
+        if (common_factor != ir::IndexExpr(1))
+          return (lhs_remainder + rhs_remainder) * common_factor;
+      }
+      return std::nullopt;
+    };
+
+    auto flatten_add_lhs = common::GetFlattenExprs<ir::Add>(lhs);
+
+    bool found = false;
+    ir::IndexExpr res(lhs.type(), 0);
+    for (size_t i = 0; i < flatten_add_lhs.size(); ++i) {
+      auto merge_res = MergeByCommonFactor(flatten_add_lhs[i], rhs);
+      if (!found && merge_res.has_value()) {
+        res = res + merge_res.value();
+        found = true;
+      } else {
+        res = res + flatten_add_lhs[i];
+      }
+    }
+    if (found) return res;
   }
 
   return Add::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
@@ -302,5 +302,198 @@ TEST_F(TestIndexExpr, Test_dynamic) {
             (((((f / ((S5 * S6) * 640)) * S4) * S5) * S6) +
              ((f % ((S5 * S6) * 640)) % ((S5 * S6) * S4))));
 }
+
+TEST_F(TestIndexExpr, CommonFactor) {
+  ir::Var S0 = ir::Var("S0");
+  ir::Var S1 = ir::Var("S1");
+  ir::Var S2 = ir::Var("S2");
+  ir::Var S3 = ir::Var("S3");
+  ir::Var S4 = ir::Var("S4");
+  ir::Var S5 = ir::Var("S5");
+  ir::Var S6 = ir::Var("S6");
+  ir::Var S7 = ir::Var("S7");
+  ir::Var S8 = ir::Var("S8");
+  ir::Var S9 = ir::Var("S9");
+  ir::Var S13 = ir::Var("S13");
+  ir::Var S17 = ir::Var("S17");
+  ir::Var S21 = ir::Var("S21");
+  ir::Var tx = ir::Var("tx");
+  ir::Var bx = ir::Var("bx");
+
+  ir::Expr q = ((((((((S1 + S13) + S17) + S21) + S5) + S9)) * S2) * S3);
+  ir::Expr q1 = (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) + ((S3 * S21) * S2)) +
+                   ((S2 * S3) * S17)) +
+                  ((S2 * S3) * S13)) +
+                 ((S2 * S3) * S1));
+  ir::Expr q2 =
+      (((((((((f * 1024) + tx) + (bx * 4096)) %
+            ((((((((((((((((((((((((((S3 * S5) * S2) * S0) +
+                                   (((S3 * S9) * S2) * S0)) +
+                                  (((S3 * S21) * S2) * S0)) +
+                                 (((S2 * S3) * S17) * S0)) +
+                                (((S2 * S3) * S13) * S0)) +
+                               (((S2 * S3) * S1) * S0)) /
+                              4096) *
+                             4096) +
+                            ((S3 * S5) * S2)) +
+                           ((S3 * S9) * S2)) +
+                          ((S3 * S21) * S2)) +
+                         ((S2 * S3) * S17)) +
+                        ((S2 * S3) * S13)) +
+                       ((S2 * S3) * S1)) +
+                      4095) /
+                     (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) +
+                         ((S3 * S21) * S2)) +
+                        ((S2 * S3) * S17)) +
+                       ((S2 * S3) * S13)) +
+                      ((S2 * S3) * S1))) *
+                    S3) *
+                   S5) *
+                  S2) +
+                 (((((((((((((((((((((S3 * S5) * S2) * S0) +
+                                   (((S3 * S9) * S2) * S0)) +
+                                  (((S3 * S21) * S2) * S0)) +
+                                 (((S2 * S3) * S17) * S0)) +
+                                (((S2 * S3) * S13) * S0)) +
+                               (((S2 * S3) * S1) * S0)) /
+                              4096) *
+                             4096) +
+                            ((S3 * S5) * S2)) +
+                           ((S3 * S9) * S2)) +
+                          ((S3 * S21) * S2)) +
+                         ((S2 * S3) * S17)) +
+                        ((S2 * S3) * S13)) +
+                       ((S2 * S3) * S1)) +
+                      4095) /
+                     (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) +
+                         ((S3 * S21) * S2)) +
+                        ((S2 * S3) * S17)) +
+                       ((S2 * S3) * S13)) +
+                      ((S2 * S3) * S1))) *
+                    S3) *
+                   S9) *
+                  S2)) +
+                (((((((((((((((((((((S3 * S5) * S2) * S0) +
+                                  (((S3 * S9) * S2) * S0)) +
+                                 (((S3 * S21) * S2) * S0)) +
+                                (((S2 * S3) * S17) * S0)) +
+                               (((S2 * S3) * S13) * S0)) +
+                              (((S2 * S3) * S1) * S0)) /
+                             4096) *
+                            4096) +
+                           ((S3 * S5) * S2)) +
+                          ((S3 * S9) * S2)) +
+                         ((S3 * S21) * S2)) +
+                        ((S2 * S3) * S17)) +
+                       ((S2 * S3) * S13)) +
+                      ((S2 * S3) * S1)) +
+                     4095) /
+                    (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) +
+                        ((S3 * S21) * S2)) +
+                       ((S2 * S3) * S17)) +
+                      ((S2 * S3) * S13)) +
+                     ((S2 * S3) * S1))) *
+                   S3) *
+                  S21) *
+                 S2)) +
+               (((((((((((((((((((((S3 * S5) * S2) * S0) +
+                                 (((S3 * S9) * S2) * S0)) +
+                                (((S3 * S21) * S2) * S0)) +
+                               (((S2 * S3) * S17) * S0)) +
+                              (((S2 * S3) * S13) * S0)) +
+                             (((S2 * S3) * S1) * S0)) /
+                            4096) *
+                           4096) +
+                          ((S3 * S5) * S2)) +
+                         ((S3 * S9) * S2)) +
+                        ((S3 * S21) * S2)) +
+                       ((S2 * S3) * S17)) +
+                      ((S2 * S3) * S13)) +
+                     ((S2 * S3) * S1)) +
+                    4095) /
+                   (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) +
+                       ((S3 * S21) * S2)) +
+                      ((S2 * S3) * S17)) +
+                     ((S2 * S3) * S13)) +
+                    ((S2 * S3) * S1))) *
+                  S2) *
+                 S3) *
+                S17)) +
+              (((((((((((((((((((((S3 * S5) * S2) * S0) +
+                                (((S3 * S9) * S2) * S0)) +
+                               (((S3 * S21) * S2) * S0)) +
+                              (((S2 * S3) * S17) * S0)) +
+                             (((S2 * S3) * S13) * S0)) +
+                            (((S2 * S3) * S1) * S0)) /
+                           4096) *
+                          4096) +
+                         ((S3 * S5) * S2)) +
+                        ((S3 * S9) * S2)) +
+                       ((S3 * S21) * S2)) +
+                      ((S2 * S3) * S17)) +
+                     ((S2 * S3) * S13)) +
+                    ((S2 * S3) * S1)) +
+                   4095) /
+                  (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) +
+                      ((S3 * S21) * S2)) +
+                     ((S2 * S3) * S17)) +
+                    ((S2 * S3) * S13)) +
+                   ((S2 * S3) * S1))) *
+                 S2) *
+                S3) *
+               S13)) +
+             (((((((((((((((((((((S3 * S5) * S2) * S0) +
+                               (((S3 * S9) * S2) * S0)) +
+                              (((S3 * S21) * S2) * S0)) +
+                             (((S2 * S3) * S17) * S0)) +
+                            (((S2 * S3) * S13) * S0)) +
+                           (((S2 * S3) * S1) * S0)) /
+                          4096) *
+                         4096) +
+                        ((S3 * S5) * S2)) +
+                       ((S3 * S9) * S2)) +
+                      ((S3 * S21) * S2)) +
+                     ((S2 * S3) * S17)) +
+                    ((S2 * S3) * S13)) +
+                   ((S2 * S3) * S1)) +
+                  4095) /
+                 (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) +
+                     ((S3 * S21) * S2)) +
+                    ((S2 * S3) * S17)) +
+                   ((S2 * S3) * S13)) +
+                  ((S2 * S3) * S1))) *
+                S2) *
+               S3) *
+              S1))) /
+           (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) + ((S3 * S21) * S2)) +
+              ((S2 * S3) * S17)) +
+             ((S2 * S3) * S13)) +
+            ((S2 * S3) * S1))) *
+          (((((S1 + S13) + S17) + S21) + S5) + S9)) *
+         S2) *
+        S3) +
+       ((((f * 1024) + tx) + (bx * 4096)) %
+        (((((((S3 * S5) * S2) + ((S3 * S9) * S2)) + ((S3 * S21) * S2)) +
+           ((S2 * S3) * S17)) +
+          ((S2 * S3) * S13)) +
+         ((S2 * S3) * S1))));
+
+  EXPECT_EQ(q.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+            (((((((S1 + S13) + S17) + S21) + S5) + S9) * S2) * S3));
+  EXPECT_EQ(q1.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+            (((((((S5 + S9) + S21) + S17) + S13) + S1) * S2) * S3));
+  EXPECT_EQ(
+      q2.as_index().Normalize(ir::IndexExpr::OptLevel::Level2),
+      ((((f * 1024) + tx) + (bx * 4096)) %
+       ((((((((((((((((S5 + S9) + S21) + S17) + S13) + S1) * S2) * S3) * S0) /
+               4096) *
+              4096) +
+             (((((((S5 + S9) + S21) + S17) + S13) + S1) * S2) * S3)) +
+            4095) /
+           (((((((S5 + S9) + S21) + S17) + S13) + S1) * S2) * S3)) *
+          S3) *
+         S2) *
+        (((((S5 + S9) + S21) + S17) + S13) + S1))));
+}
 }  // namespace common
 }  // namespace cinn
