rewrite SimplifyForLoop (#56350)

paddle/cinn/hlir/op/reduction.cc

@@ -182,6 +182,7 @@ std::shared_ptr<OpStrategy> StrategyForReduce(
         // support the length-1 for loop yet. So we simplify here. The todo
         // is that remove SimplifyForLoops below and change reduction schedule
         optim::SimplifyForLoops(&temp);
+        optim::SimplifyBlocks(&temp);
         vec_ast.emplace_back(temp);
       } else if (arg_pack[i].is_tensor()) {
         Expr temp = arg_pack[i];

paddle/cinn/optim/ir_simplify.cc

@@ -339,12 +339,9 @@ struct SimplifyForLoopsMutator : public ir::IRMutator<> {
       std::string var_name = node->loop_var->name;
       var_intervals.emplace(
           var_name, common::CasInterval{min_i->value, extent_i->value - 1});
-      if (node->body.As<ir::Block>() &&
-          node->body.As<ir::Block>()->stmts.size() == 1) {
-        *expr = node->body.As<ir::Block>()->stmts[0];
-      } else {
-        *expr = node->body;
-      }
+
+      *expr = node->body;
+
       Visit(expr, expr);
       var_intervals.erase(var_name);
     } else {