【CINN】Change sequence of Div and Mod (#70500)

liuruyan · web-flow · commit 937e6f443950 · 2024-12-30T21:13:39.000+08:00
* change seq

* fix bug
diff --git a/paddle/cinn/common/ir_util.cc b/paddle/cinn/common/ir_util.cc
@@ -19,6 +19,7 @@
 #include <unordered_set>
 
 #include "paddle/cinn/common/cas.h"
+#include "paddle/cinn/common/const_fold.h"
 #include "paddle/cinn/common/simplify_special_pattern.h"
 #include "paddle/cinn/ir/ir_mutator.h"
 #include "paddle/cinn/ir/ir_printer.h"
@@ -689,5 +690,63 @@ IndexType VerifyIndex(const ir::Expr &expr) {
   }
   return IndexType::kInvalid;
 }
+
+ir::IndexExpr ConstructIndexExprByNodeType(const ir::IrNodeTy &ty,
+                                           const ir::IndexExpr &lhs,
+                                           const ir::IndexExpr &rhs,
+                                           bool simplify_flag) {
+  switch (ty) {
+    case ir::IrNodeTy::Add:
+      return simplify_flag ? lhs + rhs : ir::Add::Make(lhs, rhs);
+    case ir::IrNodeTy::Sub:
+      return simplify_flag ? lhs - rhs : ir::Sub::Make(lhs, rhs);
+    case ir::IrNodeTy::Mul:
+      return simplify_flag ? lhs * rhs : ir::Mul::Make(lhs, rhs);
+    case ir::IrNodeTy::Div:
+      return simplify_flag ? lhs / rhs : ir::Div::Make(lhs, rhs);
+    case ir::IrNodeTy::Mod:
+      return simplify_flag ? lhs % rhs : ir::Mod::Make(lhs, rhs);
+    case ir::IrNodeTy::Min:
+      return ir::Min::Make(lhs, rhs);
+    case ir::IrNodeTy::Max:
+      return ir::Max::Make(lhs, rhs);
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type in Constructir::IndexExprByNodeType, which is: %s",
+          ty));
+  }
+}
+
+ir::IndexExpr ChangeSeqOfDivMod(const ir::IndexExpr &expr) {
+  switch (expr.node_type()) {
+    case ir::IrNodeTy::IntImm:
+    case ir::IrNodeTy::_Var_: {
+      return expr;
+    }
+    case ir::IrNodeTy::Add:
+    case ir::IrNodeTy::Sub:
+    case ir::IrNodeTy::Mul:
+    case ir::IrNodeTy::Div: {
+      auto lhs = ChangeSeqOfDivMod(expr.operand(0));
+      auto rhs = ChangeSeqOfDivMod(expr.operand(1));
+      return ConstructIndexExprByNodeType(expr.node_type(), lhs, rhs, false);
+    }
+    case ir::IrNodeTy::Mod: {
+      if (expr.operand(0).node_type() == ir::IrNodeTy::Div) {
+        auto div_lhs = ChangeSeqOfDivMod(expr.operand(0).operand(0));
+        auto div_rhs = ChangeSeqOfDivMod(expr.operand(0).operand(1));
+        auto mod_rhs = ChangeSeqOfDivMod(expr.operand(1));
+        return div_lhs % (div_rhs * mod_rhs) / div_rhs;
+      } else {
+        auto lhs = ChangeSeqOfDivMod(expr.operand(0));
+        auto rhs = ChangeSeqOfDivMod(expr.operand(1));
+        return ConstructIndexExprByNodeType(expr.node_type(), lhs, rhs, false);
+      }
+    }
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type of expr in ChangeSeqOfDivMod which is: %s", expr));
+  }
+}
 }  // namespace common
 }  // namespace cinn
diff --git a/paddle/cinn/common/ir_util.h b/paddle/cinn/common/ir_util.h
@@ -339,6 +339,31 @@ bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
 bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
                         ir::IndexExpr &expr);  // NOLINT
 
+/*!
+ * \brief Construct index expression by node type with or without simplify.
+ * \param ty The node type of index expression.
+ * \param lhs left operand.
+ * \param rhs right operand.
+ * \param simplify_flag Whether to simplify the result.
+ * \return The constructed index expression.
+ */
+ir::IndexExpr ConstructIndexExprByNodeType(const ir::IrNodeTy &ty,
+                                           const ir::IndexExpr &lhs,
+                                           const ir::IndexExpr &rhs,
+                                           bool simplify_flag = true);
+
+/*!
+ * \brief Change the sequence of `Div` and `Mod` in index expression.
+ * Mathematical formula: `(a / b) % c = (a % (b * c)) / b`
+ * For example:
+ * 1. i / 4 % 8 => i % 32 / 4
+ * 2. i / S0 % S1 => i % (S0 * S1) / S0
+ * 3. (i * 32 + j) / 4 % 8 => (i * 32 + j) % 32 / 4
+ *
+ * \param expr The `IndexExpr` to be change
+ * \return `IndexExpr` after change.
+ */
+ir::IndexExpr ChangeSeqOfDivMod(const ir::IndexExpr &expr);
 enum IndexType {
   kInvalid = 0,  // invalid expr
   kValid = 1,    // valid expr
diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
@@ -146,6 +146,12 @@ Expr Sub::Make(Expr a, Expr b) {
   return Expr(node);
 }
 
+IndexExpr Sub::Make(IndexExpr a, IndexExpr b) {
+  auto node = make_shared<Sub>(a, b);
+  node->set_index(true);
+  return IndexExpr(node);
+}
+
 void Sub::Verify() const { BinaryNodeVerify(a(), b(), "Sub"); }
 
 Expr Mul::Make(Expr a, Expr b) {
diff --git a/paddle/cinn/ir/ir.h b/paddle/cinn/ir/ir.h
@@ -100,6 +100,7 @@ struct Sub : public BinaryOpNode<Sub> {
   Sub(Expr a, Expr b) : BinaryOpNode<Sub>(a.type(), a, b) {}
 
   static Expr Make(Expr a, Expr b);
+  static IndexExpr Make(IndexExpr a, IndexExpr b);
 
   void Verify() const override;
 
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
@@ -488,31 +488,6 @@ bool IndexExpr::IsDynamic() const {
   }
 }
 
-IndexExpr ConstructIndexExprByNodeType(const IrNodeTy &ty,
-                                       const IndexExpr &lhs,
-                                       const IndexExpr &rhs) {
-  switch (ty) {
-    case IrNodeTy::Add:
-      return lhs + rhs;
-    case IrNodeTy::Sub:
-      return lhs - rhs;
-    case IrNodeTy::Mul:
-      return lhs * rhs;
-    case IrNodeTy::Div:
-      return lhs / rhs;
-    case IrNodeTy::Mod:
-      return lhs % rhs;
-    case IrNodeTy::Min:
-      return Min::Make(lhs, rhs);
-    case IrNodeTy::Max:
-      return Max::Make(lhs, rhs);
-    default:
-      PADDLE_THROW(::common::errors::InvalidArgument(
-          "Unsupported type in ConstructIndexExprByNodeType, which is: %s",
-          ty));
-  }
-}
-
 IndexExpr Simplify(const IndexExpr &expr, IndexExpr::OptLevel level) {
   switch (expr.node_type()) {
     case ir::IrNodeTy::IntImm:
diff --git a/test/cpp/pir/cinn/adt/index_expr_test.cc b/test/cpp/pir/cinn/adt/index_expr_test.cc
@@ -168,5 +168,41 @@ TEST_F(TestIndexExpr, IndexExpr_3) {
             ir::IndexExpr((S4 * 256 + S5 + S6 * 1024)) % 25088);
   EXPECT_EQ(q16.as_index().Normalize(), ir::IndexExpr(S4 * 256 + S5));
 }
+
+TEST_F(TestIndexExpr, Change_Seq_Of_Div_Mod) {
+  ir::Expr q1 = S4 / S5;
+  ir::Expr q2 = S4 % S5;
+  ir::Expr q3 = S4 / S5 % S6;
+  ir::Expr q4 = S4 / S5 % S6;
+
+  EXPECT_EQ(ChangeSeqOfDivMod(q1.as_index()), q1);
+  EXPECT_EQ(ChangeSeqOfDivMod(q2.as_index()), q2);
+  EXPECT_EQ(ChangeSeqOfDivMod(q3.as_index()), S4 % (S5 * S6) / S5);
+}
+
+TEST_F(TestIndexExpr, Test_ConstructIndexExprByNodeType) {
+  ir::Expr result_add = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Add, S4.as_index(), S5.as_index(), true);
+  ir::Expr result_sub = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Sub, S4.as_index(), S5.as_index(), false);
+  ir::Expr result_mul = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Mul, S4.as_index(), S5.as_index(), true);
+  ir::Expr result_div = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Div, S4.as_index(), S5.as_index(), true);
+  ir::Expr result_mod = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Mod, S4.as_index(), S5.as_index(), true);
+  ir::Expr result_min = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Min, S4.as_index(), S5.as_index(), false);
+  ir::Expr result_max = ConstructIndexExprByNodeType(
+      ir::IrNodeTy::Max, S4.as_index(), S5.as_index(), false);
+
+  EXPECT_EQ(result_add, S4 + S5);
+  EXPECT_EQ(result_sub, S4 - S5);
+  EXPECT_EQ(result_mul, S4 * S5);
+  EXPECT_EQ(result_div, S4 / S5);
+  EXPECT_EQ(result_mod, S4 % S5);
+  EXPECT_EQ(result_min, ir::Min::Make(S4, S5));
+  EXPECT_EQ(result_max, ir::Max::Make(S4, S5));
+}
 }  // namespace common
 }  // namespace cinn
