【CINN】Change int64 to int32 in IndexExpr simplify (#70461)

liuruyan · web-flow · commit 19c63a2b3eec · 2024-12-27T19:03:24.000+08:00
* change int64 to int32 in simplify

* fix overflow bug
diff --git a/paddle/cinn/ir/ir_base.h b/paddle/cinn/ir/ir_base.h
@@ -497,6 +497,9 @@ struct IndexExpr : public IrNodeRef {
   explicit IndexExpr(int32_t x) : IrNodeRef(new IntImm(Int(32), x)) {}
   explicit IndexExpr(int64_t x) : IrNodeRef(new IntImm(Int(64), x)) {}
 
+  explicit IndexExpr(Type t, int64_t x)
+      : IrNodeRef(new IntImm(x > INT32_MAX ? Int(64) : t, x)) {}
+
   bool is_var() const;
   _Var_* as_var();
   const _Var_* as_var() const;
diff --git a/paddle/cinn/ir/op/ir_operators.cc b/paddle/cinn/ir/op/ir_operators.cc
@@ -377,14 +377,16 @@ static IndexExpr SimplifyAdd(const IndexExpr &lhs, const IndexExpr &rhs) {
   auto lhsAdd = lhs.As<Add>();
   if (lhsAdd && rhsConst) {
     if (auto lrhs = lhsAdd->b().as_index().As<IntImm>()) {
-      return lhsAdd->a().as_index() + (lrhs->value + rhsConst->value);
+      return lhsAdd->a().as_index() +
+             IndexExpr(lrhs->type(), lrhs->value + rhsConst->value);
     }
   }
 
   // (d0 + 2) + d1 ===> d0 + d1 + 2.
   if (lhsAdd) {
     if (auto lrhs = lhsAdd->b().as_index().As<IntImm>()) {
-      return lhsAdd->a().as_index() + rhs + lrhs->value;
+      return lhsAdd->a().as_index() + rhs +
+             IndexExpr(lrhs->type(), lrhs->value);
     }
   }
   // expr * c1 + expr * c2 ===> expr * (c1 + c2)
@@ -409,12 +411,14 @@ static IndexExpr SimplifyAdd(const IndexExpr &lhs, const IndexExpr &rhs) {
   }
 
   if (first == second) {
-    return first * (lconst + rconst);
+    return first * IndexExpr(lhs->type(), lconst + rconst);
   }
 
   if (lconst != 1 && rconst != 1) {
-    if (lconst == rconst) return (first + second) * lconst;
-    if (lconst == -rconst) return (first - second) * lconst;
+    if (lconst == rconst)
+      return (first + second) * IndexExpr(lhs->type(), lconst);
+    if (lconst == -rconst)
+      return (first - second) * IndexExpr(lhs->type(), lconst);
   }
 
   // deal corner case!
@@ -460,22 +464,25 @@ static IndexExpr SimplifyMul(const IndexExpr &lhs, const IndexExpr &rhs) {
   auto lhsMul = lhs.As<Mul>();
   if (lhsMul && rhsConst) {
     if (auto lrhs = lhsMul->b().as_index().As<IntImm>()) {
-      return lhsMul->a().as_index() * (lrhs->value * rhsConst->value);
+      return lhsMul->a().as_index() *
+             IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);
     }
   }
 
   // (d0 + 3) * 5 ===> d0 * 5 + 15.
   auto lhsAdd = lhs.As<Add>();
   if (lhsAdd && rhsConst) {
     if (auto lrhs = lhsAdd->b().as_index().As<IntImm>()) {
-      return lhsAdd->a().as_index() * rhs + (lrhs->value * rhsConst->value);
+      return lhsAdd->a().as_index() * rhs +
+             IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);
     }
   }
 
   // (d0 * 2) * d1 ===> d0 * d1 * 2.
   if (lhsMul) {
     if (auto lrhs = lhsMul->b().as_index().As<IntImm>()) {
-      return lhsMul->a().as_index() * rhs * lrhs->value;
+      return lhsMul->a().as_index() * rhs *
+             IndexExpr(lrhs->type(), lrhs->value);
     }
   }
 
@@ -513,24 +520,25 @@ static IndexExpr SimplifyDiv(const IndexExpr &lhs, const IndexExpr &rhs) {
       int64_t lrhsFactor = lhsAdd->b().as_index().GetLargestMultiplyPart();
       if (llhsFactor % rhsConst->value == 0 &&
           lrhsFactor % rhsConst->value == 0) {
-        return lhsAdd->a().as_index() / rhsConst->value +
-               lhsAdd->b().as_index() / rhsConst->value;
+        return lhsAdd->a().as_index() / rhs + lhsAdd->b().as_index() / rhs;
       }
     }
 
     // expr1 * (c1 * c2) / c1 ===> expr1 * c2.
     if (lhsMul) {
       if (auto lrhs = lhsMul->b().as_index().As<IntImm>()) {
         if (lrhs->value % rhsConst->value == 0) {
-          return lhsMul->a().as_index() * (lrhs->value / rhsConst->value);
+          return lhsMul->a().as_index() *
+                 IndexExpr(lrhs->type(), lrhs->value / rhsConst->value);
         }
       }
     }
 
     // S0 / 2 / 5 ===> S0 / 10.
     if (lhsDiv) {
       if (auto lrhs = lhsDiv->b().as_index().As<IntImm>()) {
-        return lhsDiv->a().as_index() / (lrhs->value * rhsConst->value);
+        return lhsDiv->a().as_index() /
+               IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);
       }
     }
   } else {
@@ -568,9 +576,9 @@ static IndexExpr SimplifyMod(const IndexExpr &lhs, const IndexExpr &rhs) {
       int64_t llhsFactor = lhsAdd->a().as_index().GetLargestMultiplyPart();
       int64_t lrhsFactor = lhsAdd->b().as_index().GetLargestMultiplyPart();
       if (llhsFactor % rhsConst->value == 0)
-        return lhsAdd->b().as_index() % rhsConst->value;
+        return lhsAdd->b().as_index() % rhs;
       if (lrhsFactor % rhsConst->value == 0)
-        return lhsAdd->a().as_index() % rhsConst->value;
+        return lhsAdd->a().as_index() % rhs;
     }
 
     // expr1 * (c1 * c2) % c1 ===> 0.
@@ -581,7 +589,7 @@ static IndexExpr SimplifyMod(const IndexExpr &lhs, const IndexExpr &rhs) {
     if (lhsMod) {
       int64_t llhsFactor = lhsMod->b().as_index().GetLargestMultiplyPart();
       if (llhsFactor % rhsConst->value == 0)
-        return lhsMod->a().as_index() % rhsConst->value;
+        return lhsMod->a().as_index() % rhs;
     }
   } else {
     // dynamic branch!

Original file line number	Diff line number	Diff line change
`@@ -377,14 +377,16 @@ static IndexExpr SimplifyAdd(const IndexExpr &lhs, const IndexExpr &rhs) {`
`377`	`377`	`auto lhsAdd = lhs.As<Add>();`
`378`	`378`	`if (lhsAdd && rhsConst) {`
`379`	`379`	`if (auto lrhs = lhsAdd->b().as_index().As<IntImm>()) {`
`380`		`- return lhsAdd->a().as_index() + (lrhs->value + rhsConst->value);`
	`380`	`+ return lhsAdd->a().as_index() +`
	`381`	`+ IndexExpr(lrhs->type(), lrhs->value + rhsConst->value);`
`381`	`382`	`}`
`382`	`383`	`}`
`383`	`384`
`384`	`385`	`// (d0 + 2) + d1 ===> d0 + d1 + 2.`
`385`	`386`	`if (lhsAdd) {`
`386`	`387`	`if (auto lrhs = lhsAdd->b().as_index().As<IntImm>()) {`
`387`		`- return lhsAdd->a().as_index() + rhs + lrhs->value;`
	`388`	`+ return lhsAdd->a().as_index() + rhs +`
	`389`	`+ IndexExpr(lrhs->type(), lrhs->value);`
`388`	`390`	`}`
`389`	`391`	`}`
`390`	`392`	`// expr * c1 + expr * c2 ===> expr * (c1 + c2)`
`@@ -409,12 +411,14 @@ static IndexExpr SimplifyAdd(const IndexExpr &lhs, const IndexExpr &rhs) {`
`409`	`411`	`}`
`410`	`412`
`411`	`413`	`if (first == second) {`
`412`		`- return first * (lconst + rconst);`
	`414`	`+ return first * IndexExpr(lhs->type(), lconst + rconst);`
`413`	`415`	`}`
`414`	`416`
`415`	`417`	`if (lconst != 1 && rconst != 1) {`
`416`		`- if (lconst == rconst) return (first + second) * lconst;`
`417`		`- if (lconst == -rconst) return (first - second) * lconst;`
	`418`	`+ if (lconst == rconst)`
	`419`	`+ return (first + second) * IndexExpr(lhs->type(), lconst);`
	`420`	`+ if (lconst == -rconst)`
	`421`	`+ return (first - second) * IndexExpr(lhs->type(), lconst);`
`418`	`422`	`}`
`419`	`423`
`420`	`424`	`// deal corner case!`
`@@ -460,22 +464,25 @@ static IndexExpr SimplifyMul(const IndexExpr &lhs, const IndexExpr &rhs) {`
`460`	`464`	`auto lhsMul = lhs.As<Mul>();`
`461`	`465`	`if (lhsMul && rhsConst) {`
`462`	`466`	`if (auto lrhs = lhsMul->b().as_index().As<IntImm>()) {`
`463`		`- return lhsMul->a().as_index() * (lrhs->value * rhsConst->value);`
	`467`	`+ return lhsMul->a().as_index() *`
	`468`	`+ IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);`
`464`	`469`	`}`
`465`	`470`	`}`
`466`	`471`
`467`	`472`	`// (d0 + 3) * 5 ===> d0 * 5 + 15.`
`468`	`473`	`auto lhsAdd = lhs.As<Add>();`
`469`	`474`	`if (lhsAdd && rhsConst) {`
`470`	`475`	`if (auto lrhs = lhsAdd->b().as_index().As<IntImm>()) {`
`471`		`- return lhsAdd->a().as_index() * rhs + (lrhs->value * rhsConst->value);`
	`476`	`+ return lhsAdd->a().as_index() * rhs +`
	`477`	`+ IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);`
`472`	`478`	`}`
`473`	`479`	`}`
`474`	`480`
`475`	`481`	`// (d0 * 2) * d1 ===> d0 * d1 * 2.`
`476`	`482`	`if (lhsMul) {`
`477`	`483`	`if (auto lrhs = lhsMul->b().as_index().As<IntImm>()) {`
`478`		`- return lhsMul->a().as_index() * rhs * lrhs->value;`
	`484`	`+ return lhsMul->a().as_index() * rhs *`
	`485`	`+ IndexExpr(lrhs->type(), lrhs->value);`
`479`	`486`	`}`
`480`	`487`	`}`
`481`	`488`
`@@ -513,24 +520,25 @@ static IndexExpr SimplifyDiv(const IndexExpr &lhs, const IndexExpr &rhs) {`
`513`	`520`	`int64_t lrhsFactor = lhsAdd->b().as_index().GetLargestMultiplyPart();`
`514`	`521`	`if (llhsFactor % rhsConst->value == 0 &&`
`515`	`522`	`lrhsFactor % rhsConst->value == 0) {`
`516`		`- return lhsAdd->a().as_index() / rhsConst->value +`
`517`		`- lhsAdd->b().as_index() / rhsConst->value;`
	`523`	`+ return lhsAdd->a().as_index() / rhs + lhsAdd->b().as_index() / rhs;`
`518`	`524`	`}`
`519`	`525`	`}`
`520`	`526`
`521`	`527`	`// expr1 * (c1 * c2) / c1 ===> expr1 * c2.`
`522`	`528`	`if (lhsMul) {`
`523`	`529`	`if (auto lrhs = lhsMul->b().as_index().As<IntImm>()) {`
`524`	`530`	`if (lrhs->value % rhsConst->value == 0) {`
`525`		`- return lhsMul->a().as_index() * (lrhs->value / rhsConst->value);`
	`531`	`+ return lhsMul->a().as_index() *`
	`532`	`+ IndexExpr(lrhs->type(), lrhs->value / rhsConst->value);`
`526`	`533`	`}`
`527`	`534`	`}`
`528`	`535`	`}`
`529`	`536`
`530`	`537`	`// S0 / 2 / 5 ===> S0 / 10.`
`531`	`538`	`if (lhsDiv) {`
`532`	`539`	`if (auto lrhs = lhsDiv->b().as_index().As<IntImm>()) {`
`533`		`- return lhsDiv->a().as_index() / (lrhs->value * rhsConst->value);`
	`540`	`+ return lhsDiv->a().as_index() /`
	`541`	`+ IndexExpr(lrhs->type(), lrhs->value * rhsConst->value);`
`534`	`542`	`}`
`535`	`543`	`}`
`536`	`544`	`} else {`
`@@ -568,9 +576,9 @@ static IndexExpr SimplifyMod(const IndexExpr &lhs, const IndexExpr &rhs) {`
`568`	`576`	`int64_t llhsFactor = lhsAdd->a().as_index().GetLargestMultiplyPart();`
`569`	`577`	`int64_t lrhsFactor = lhsAdd->b().as_index().GetLargestMultiplyPart();`
`570`	`578`	`if (llhsFactor % rhsConst->value == 0)`
`571`		`- return lhsAdd->b().as_index() % rhsConst->value;`
	`579`	`+ return lhsAdd->b().as_index() % rhs;`
`572`	`580`	`if (lrhsFactor % rhsConst->value == 0)`
`573`		`- return lhsAdd->a().as_index() % rhsConst->value;`
	`581`	`+ return lhsAdd->a().as_index() % rhs;`
`574`	`582`	`}`
`575`	`583`
`576`	`584`	`// expr1 * (c1 * c2) % c1 ===> 0.`
`@@ -581,7 +589,7 @@ static IndexExpr SimplifyMod(const IndexExpr &lhs, const IndexExpr &rhs) {`
`581`	`589`	`if (lhsMod) {`
`582`	`590`	`int64_t llhsFactor = lhsMod->b().as_index().GetLargestMultiplyPart();`
`583`	`591`	`if (llhsFactor % rhsConst->value == 0)`
`584`		`- return lhsMod->a().as_index() % rhsConst->value;`
	`592`	`+ return lhsMod->a().as_index() % rhs;`
`585`	`593`	`}`
`586`	`594`	`} else {`
`587`	`595`	`// dynamic branch!`