[CINN]backend support stmt to expr visitor and run on func body

paddle/cinn/ir/CMakeLists.txt

@@ -21,6 +21,7 @@ gather_srcs(
   schedule_block_graph.cc
   stmt.cc
   stmt_visitors.cc
+  expr_visitors.cc
   dim.cc)
 
 add_subdirectory(ir_analyzer)

paddle/cinn/ir/expr_visitors.cc

@@ -0,0 +1,218 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "paddle/cinn/ir/expr_visitors.h"
+
+namespace cinn {
+namespace ir {
+
+void VisitExpr(const stmt::Let &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &symbol = stmt->symbol();
+  const auto &body = stmt->body();
+  callback(symbol);
+  if (body.defined()) {
+    callback(body);
+  }
+}
+
+void VisitExpr(const stmt::Store &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &value = stmt->value();
+  const auto &tensor = stmt->tensor();
+  const auto &indices = stmt->indices();
+  callback(value);
+  callback(tensor);
+  for (const auto &indice : indices) {
+    callback(indice);
+  }
+}
+
+void VisitExpr(const stmt::Alloc &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &destination = stmt->destination();
+  const auto &extents = stmt->extents();
+  const auto &condition = stmt->condition();
+  const auto &body = stmt->body();
+  callback(destination);
+  for (const auto &extent : extents) {
+    callback(extent);
+  }
+  if (condition.defined()) {
+    callback(condition);
+  }
+  if (body.defined()) {
+    callback(body);
+  }
+}
+
+void VisitExpr(const stmt::Free &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &destination = stmt->destination();
+  callback(destination);
+}
+
+void VisitExpr(const stmt::IfThenElse &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &condition = stmt->condition();
+  callback(condition);
+}
+
+void VisitExpr(const stmt::For &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &min = stmt->min();
+  const auto &extent = stmt->extent();
+  callback(min);
+  callback(extent);
+}
+
+void VisitExpr(const stmt::Schedule &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &iter_vars = stmt->iter_vars();
+  const auto &iter_values = stmt->iter_values();
+  const auto &read_buffers = stmt->read_buffers();
+  const auto &write_buffers = stmt->write_buffers();
+
+  for (const auto &iter_var : iter_vars) {
+    if (iter_var->lower_bound.defined()) {
+      callback(iter_var->lower_bound);
+    }
+    if (iter_var->upper_bound.defined()) {
+      callback(iter_var->upper_bound);
+    }
+  }
+  for (const auto &iter_value : iter_values) {
+    callback(iter_value);
+  }
+  for (const auto &read_buffer : read_buffers) {
+    callback(read_buffer);
+  }
+  for (const auto &write_buffer : write_buffers) {
+    callback(write_buffer);
+  }
+}
+
+void VisitExpr(const stmt::Evaluate &stmt,
+               const std::function<void(const Expr &)> &callback) {
+  const auto &value = stmt->value();
+  callback(value);
+}
+
+void MutateExpr(stmt::Let stmt, const std::function<void(Expr *)> &callback) {
+  ir::Expr symbol = stmt->symbol();
+  ir::Expr body = stmt->body();
+  callback(&symbol);
+  if (body.defined()) {
+    callback(&body);
+  }
+  stmt->set_symbol(symbol);
+  stmt->set_body(body);
+}
+
+void MutateExpr(stmt::Store stmt, const std::function<void(Expr *)> &callback) {
+  ir::Expr value = stmt->value();
+  ir::Expr tensor = stmt->tensor();
+  std::vector<ir::Expr> indices = stmt->indices();
+  callback(&value);
+  callback(&tensor);
+  for (ir::Expr &indice : indices) {
+    callback(&indice);
+  }
+  stmt->set_value(value);
+  stmt->set_tensor(tensor);
+  stmt->set_indices(indices);
+}
+
+void MutateExpr(stmt::Alloc stmt, const std::function<void(Expr *)> &callback) {
+  ir::Expr destination = stmt->destination();
+  std::vector<ir::Expr> extents = stmt->extents();
+  ir::Expr condition = stmt->condition();
+  ir::Expr body = stmt->body();
+  callback(&destination);
+  for (ir::Expr &extent : extents) {
+    callback(&extent);
+  }
+  if (condition.defined()) {
+    callback(&condition);
+  }
+  if (body.defined()) {
+    callback(&body);
+  }
+  stmt->set_destination(destination);
+  stmt->set_extents(extents);
+  stmt->set_condition(condition);
+  stmt->set_body(body);
+}
+
+void MutateExpr(stmt::Free stmt, const std::function<void(Expr *)> &callback) {
+  ir::Expr destination = stmt->destination();
+  callback(&destination);
+  stmt->set_destination(destination);
+}
+
+void MutateExpr(stmt::IfThenElse stmt,
+                const std::function<void(Expr *)> &callback) {
+  ir::Expr condition = stmt->condition();
+  callback(&condition);
+  stmt->set_condition(condition);
+}
+
+void MutateExpr(stmt::For stmt, const std::function<void(Expr *)> &callback) {
+  ir::Expr min = stmt->min();
+  ir::Expr extent = stmt->extent();
+  callback(&min);
+  callback(&extent);
+  stmt->set_min(min);
+  stmt->set_extent(extent);
+}
+
+void MutateExpr(stmt::Schedule stmt,
+                const std::function<void(Expr *)> &callback) {
+  std::vector<ir::Var> iter_vars = stmt->iter_vars();
+  std::vector<ir::Expr> iter_values = stmt->iter_values();
+  std::vector<ir::Expr> read_buffers = stmt->read_buffers();
+  std::vector<ir::Expr> write_buffers = stmt->write_buffers();
+
+  for (ir::Var iter_var : iter_vars) {
+    if (iter_var->lower_bound.defined()) {
+      callback(&(iter_var->lower_bound));
+    }
+    if (iter_var->upper_bound.defined()) {
+      callback(&(iter_var->upper_bound));
+    }
+  }
+  for (ir::Expr &iter_value : iter_values) {
+    callback(&iter_value);
+  }
+  for (ir::Expr &read_buffer : read_buffers) {
+    callback(&read_buffer);
+  }
+  for (ir::Expr &write_buffer : write_buffers) {
+    callback(&write_buffer);
+  }
+
+  stmt->set_iter_vars(iter_vars);
+  stmt->set_iter_values(iter_values);
+  stmt->set_read_buffers(read_buffers);
+  stmt->set_write_buffers(write_buffers);
+}
+
+void MutateExpr(stmt::Evaluate stmt,
+                const std::function<void(Expr *)> &callback) {
+  ir::Expr value = stmt->value();
+  callback(&value);
+  stmt->set_value(value);
+}
+
+}  // namespace ir
+}  // namespace cinn

paddle/cinn/ir/expr_visitors.h

@@ -0,0 +1,70 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <functional>
+
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/stmt.h"
+
+namespace cinn {
+namespace ir {
+
+// Defines utilities for walking exprs (no walking into nest block)
+void VisitExpr(const stmt::Let &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::Store &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::Alloc &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::Free &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::IfThenElse &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::For &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::Schedule &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void VisitExpr(const stmt::Evaluate &stmt,
+               const std::function<void(const Expr &)> &callback);
+
+void MutateExpr(stmt::Let stmt, const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::Store stmt, const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::Alloc stmt, const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::Free stmt, const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::IfThenElse stmt,
+                const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::For stmt, const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::Schedule stmt,
+                const std::function<void(Expr *)> &callback);
+
+void MutateExpr(stmt::Evaluate stmt,
+                const std::function<void(Expr *)> &callback);
+
+}  // namespace ir
+}  // namespace cinn

paddle/cinn/optim/longlong2int_pass.cc

@@ -144,7 +144,7 @@ class LongLong2IntStmtPass : public StmtPass {
 class LongLong2IntExprPass : public ExprPass {
  public:
   LongLong2IntExprPass() : ExprPass("longlong2int_expr") {}
-  LogicalResult Run(ir::Expr expr) override;
+  LogicalResult Run(ir::Expr* expr) override;
 };
 }  // namespace
 
@@ -233,9 +233,9 @@ LogicalResult LongLong2IntStmtPass::Run(ir::stmt::StmtRef stmt) {
   return LogicalResult::success();
 }
 
-LogicalResult LongLong2IntExprPass::Run(ir::Expr expr) {
+LogicalResult LongLong2IntExprPass::Run(ir::Expr* expr) {
   CastLonglong2IntMutator narrow;
-  narrow(&expr);
+  narrow(expr);
   return LogicalResult::success();
 }
 std::unique_ptr<StmtPass> CreateLongLong2IntStmtPass() {

paddle/cinn/pass/pass.h

@@ -58,7 +58,14 @@ class FuncPass : public Pass<ir::LoweredFunc> {
  public:
   explicit FuncPass(const std::string& name) : Pass(PassKind::PK_FUNC, name) {}
 
+  // Run on the whole function.
   virtual LogicalResult Run(ir::LoweredFunc f) = 0;
+  // Only run on function body.
+  virtual LogicalResult Run(ir::stmt::BlockRef block) {
+    LOG(WARNING) << name()
+                 << "should run on the whole function, not on the body block.";
+    return LogicalResult::failure();
+  }
 };
 
 class BlockPass : public Pass<ir::stmt::BlockRef> {
@@ -74,10 +81,10 @@ class StmtPass : public Pass<ir::stmt::StmtRef> {
   virtual LogicalResult Run(ir::stmt::StmtRef stmt) = 0;
 };
 
-class ExprPass : public Pass<ir::Expr> {
+class ExprPass : public Pass<ir::Expr*> {
  public:
   explicit ExprPass(const std::string& name) : Pass(PassKind::PK_STMT, name) {}
-  virtual LogicalResult Run(ir::Expr expr) = 0;
+  virtual LogicalResult Run(ir::Expr* expr) = 0;
 };
 
 }  // namespace optim