opt longlong2int (#69522)

liuruyan · web-flow · commit 21446e93ad9d · 2024-11-20T14:16:49.000+08:00
diff --git a/paddle/cinn/optim/longlong2int.cc b/paddle/cinn/optim/longlong2int.cc
@@ -69,7 +69,7 @@ class CastLonglong2Int : public ir::IRMutator<> {
     std::for_each(node->shape.begin(),
                   node->shape.end(),
                   [&](cinn::ir::Expr& e) { e->convert_int64_to_int32(); });
-    CastBuffer(node->buffer);
+    CastBufferMeta(node->buffer);
   }
   void Visit(const ir::Load* op, Expr* expr) override {
     auto node = expr->As<ir::Load>();
@@ -107,7 +107,7 @@ class CastLonglong2Int : public ir::IRMutator<> {
                       range->ranges.end(),
                       [&](cinn::ir::Var& v) { CastVarWithBound(v); });
         auto bf = range->buffer.as_buffer_ref();
-        CastBuffer(bf);
+        CastBufferMeta(bf);
       }
     }
 
@@ -117,7 +117,7 @@ class CastLonglong2Int : public ir::IRMutator<> {
                       range->ranges.end(),
                       [&](cinn::ir::Var& v) { CastVarWithBound(v); });
         auto bf = range->buffer.as_buffer_ref();
-        CastBuffer(bf);
+        CastBufferMeta(bf);
       }
     }
     ir::IRMutator<>::Visit(&(node->body), &(node->body));
@@ -140,7 +140,7 @@ class CastLonglong2Int : public ir::IRMutator<> {
     if (lb.defined()) lb->convert_int64_to_int32();
     if (ub.defined()) ub->convert_int64_to_int32();
   }
-  void CastBuffer(cinn::ir::Buffer& bf) {  // NOLINT
+  void CastBufferMeta(cinn::ir::Buffer& bf) {  // NOLINT
     if (!bf.defined()) return;
     std::for_each(bf->shape.begin(), bf->shape.end(), [&](cinn::ir::Expr& e) {
       e->convert_int64_to_int32();
diff --git a/paddle/cinn/optim/replace_var_with_expr.cc b/paddle/cinn/optim/replace_var_with_expr.cc
@@ -155,13 +155,7 @@ struct CollectTensorIndexMutator : public ir::IRMutator<> {
 std::vector<std::vector<Expr>> CollectTensorIndex(
     Expr* source, const std::string& tensor_name) {
   CollectTensorIndexMutator mutator(tensor_name);
-  std::vector<std::vector<Expr>> result = mutator(source);
-  for (auto& i : result) {
-    for (auto& j : i) {
-      j = cinn::common::AutoSimplify(j);
-    }
-  }
-  return result;
+  return mutator(source);
 }
 
 }  // namespace optim
diff --git a/paddle/cinn/optim/transform_gpu_forloop.cc b/paddle/cinn/optim/transform_gpu_forloop.cc
@@ -40,7 +40,7 @@
 #include "paddle/cinn/utils/string.h"
 #include "paddle/common/enforce.h"
 
-PD_DECLARE_bool(cinn_longlong2int_for_integer);
+PD_DECLARE_bool(cinn_longlong2int);
 namespace cinn {
 namespace optim {
 
@@ -487,7 +487,7 @@ void OptimizeExprGPU(Expr *expr) {
   ReplaceVarToZero replace_var_to_zero;
   replace_var_to_zero(expr);
 
-  if (FLAGS_cinn_longlong2int_for_integer) {
+  if (FLAGS_cinn_longlong2int) {
     TryCastLonglong2Int(expr);
   }
 
diff --git a/paddle/cinn/runtime/flags.cc b/paddle/cinn/runtime/flags.cc
@@ -296,8 +296,8 @@ PD_DEFINE_bool(cinn_check_tensor_buffer_map,
                BoolFromEnv("FLAGS_cinn_check_tensor_buffer_map", false),
                "Whether to check tensor buffer mapping in cinn ir.");
 
-PD_DEFINE_bool(cinn_longlong2int_for_integer,
-               BoolFromEnv("FLAGS_cinn_longlong2int_for_integer", true),
+PD_DEFINE_bool(cinn_longlong2int,
+               BoolFromEnv("FLAGS_cinn_longlong2int", true),
                "Whether to cast long long to int for integer.");
 
 namespace cinn {

Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ class CastLonglong2Int : public ir::IRMutator<> {`
`69`	`69`	`std::for_each(node->shape.begin(),`
`70`	`70`	`node->shape.end(),`
`71`	`71`	`[&](cinn::ir::Expr& e) { e->convert_int64_to_int32(); });`
`72`		`- CastBuffer(node->buffer);`
	`72`	`+ CastBufferMeta(node->buffer);`
`73`	`73`	`}`
`74`	`74`	`void Visit(const ir::Load* op, Expr* expr) override {`
`75`	`75`	`auto node = expr->As<ir::Load>();`
`@@ -107,7 +107,7 @@ class CastLonglong2Int : public ir::IRMutator<> {`
`107`	`107`	`range->ranges.end(),`
`108`	`108`	`[&](cinn::ir::Var& v) { CastVarWithBound(v); });`
`109`	`109`	`auto bf = range->buffer.as_buffer_ref();`
`110`		`- CastBuffer(bf);`
	`110`	`+ CastBufferMeta(bf);`
`111`	`111`	`}`
`112`	`112`	`}`
`113`	`113`
`@@ -117,7 +117,7 @@ class CastLonglong2Int : public ir::IRMutator<> {`
`117`	`117`	`range->ranges.end(),`
`118`	`118`	`[&](cinn::ir::Var& v) { CastVarWithBound(v); });`
`119`	`119`	`auto bf = range->buffer.as_buffer_ref();`
`120`		`- CastBuffer(bf);`
	`120`	`+ CastBufferMeta(bf);`
`121`	`121`	`}`
`122`	`122`	`}`
`123`	`123`	`ir::IRMutator<>::Visit(&(node->body), &(node->body));`
`@@ -140,7 +140,7 @@ class CastLonglong2Int : public ir::IRMutator<> {`
`140`	`140`	`if (lb.defined()) lb->convert_int64_to_int32();`
`141`	`141`	`if (ub.defined()) ub->convert_int64_to_int32();`
`142`	`142`	`}`
`143`		`- void CastBuffer(cinn::ir::Buffer& bf) { // NOLINT`
	`143`	`+ void CastBufferMeta(cinn::ir::Buffer& bf) { // NOLINT`
`144`	`144`	`if (!bf.defined()) return;`
`145`	`145`	`std::for_each(bf->shape.begin(), bf->shape.end(), [&](cinn::ir::Expr& e) {`
`146`	`146`	`e->convert_int64_to_int32();`