【CINN】Simplify offset temporary version (#67554)

liuruyan · web-flow · commit 65c463b13f49 · 2024-08-29T11:23:59.000+08:00
* simple case in IndiceToOffset

* fix is_zero bug

* fix codestyle.

* use RampRelatedAdd and RampRelatedMul.

* Add SimplifyCast.

* fix error.

* fix error.

* add init.

* fix autosimplify bug.

* fix codestyle.
diff --git a/paddle/cinn/common/cas.cc b/paddle/cinn/common/cas.cc
@@ -1541,6 +1541,9 @@ Expr CasSimplifyMutator::SimplifySpecificSum(Expr tmp) {
                       2U,
                       ::common::errors::InvalidArgument(
                           "left_mul's operands size should be greater than 2"));
+
+    if (left_mul->operands().size() > 2) return tmp;
+
     Expr mul_left = left_mul->operand(0);
     Expr mul_right = left_mul->operand(1);
 
diff --git a/paddle/cinn/common/ir_util.cc b/paddle/cinn/common/ir_util.cc
@@ -192,18 +192,20 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
             "the current data type of shape[{}] is {}",
             i,
             shape[i].type()));
-    Expr indice_prod = indices[i];
-    optim::SimplifyCast(&indice_prod);
-    for (int j = i + 1; j < shape.size(); j++) {
-      indice_prod = RampRelatedMul(indice_prod, shape[j]);
-    }
+    Expr indice_cast = indices[i];
+    optim::SimplifyCast(&indice_cast);
     if (res.defined()) {
-      res = RampRelatedAdd(res, indice_prod);
+      res = RampRelatedAdd(RampRelatedMul(res, shape[i]), indice_cast);
     } else {
-      res = indice_prod;
+      res = indice_cast;
+    }
+
+    if (i > 0) {
+      res = cinn::common::AutoSimplify(res);
     }
   }
-  return cinn::common::AutoSimplify(res);
+
+  return res;
 }
 
 Expr IndiceToAbsOffset(const std::vector<int> &shape,
@@ -258,7 +260,7 @@ bool is_zero(Expr v) {
   auto *float_n = v.As<ir::FloatImm>();
 
   if (int_n) return int_n->value == 0;
-  if (float_n) return float_n->value = 0.f;
+  if (float_n) return float_n->value == 0.f;
   return false;
 }
 
