opt longlong2int (#69522)

Author: liuruyan

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();

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

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);
   }
 

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 {