[CINN] Fix bug of remove append var (#71927)

huangjiyi · web-flow · commit 977f82a90d91 · 2025-03-27T14:13:13.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.cc b/paddle/cinn/hlir/framework/pir/trivial_op_util.cc
@@ -128,6 +128,14 @@ void MappingTargetExprToDestExprMutator::Visit(const ir::Reduce* reduce,
   }
 }
 
+void MappingTargetExprToDestExprMutator::Visit(const ir::_Var_* var, Expr* op) {
+  if (source_.is_var() && source_.as_var_ref()->name == var->name) {
+    *op = dest_;
+  } else {
+    IRMutator::Visit(var, op);
+  }
+}
+
 bool CheckIterEq(const std::vector<ir::Var>& up_iter,
                  const std::vector<ir::Var>& down_iter) {
   if (up_iter.size() != down_iter.size()) return false;
@@ -722,6 +730,14 @@ ExprTransformer RemoveForTransformer(int axis) {
         (ExprSetFinderUtils::ChildScheduleBlockRealizes *
          ExprSetFinderUtils::ScheduleBlockRealizeIsInit),
         RemoveVarInScheduleBlockRealize(iters[axis], ir::Expr(0)));
+    // Remove var in append if
+    auto loop_var = target_for.As<ir::For>()->loop_var;
+    auto realizes = (ExprSetFinderUtils::ChildScheduleBlockRealizes *
+                     ExprSetFinderUtils::ScheduleBlockRealizeNotRoot)(copied);
+    for (auto& realize : realizes) {
+      ComposeUtils::MappingTargetExprToDestExprMutator(loop_var,
+                                                       ir::Expr(0))(&realize);
+    }
     return copied;
   };
   return ExprTransformer(f);
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.h b/paddle/cinn/hlir/framework/pir/trivial_op_util.h
@@ -79,6 +79,7 @@ struct MappingTargetExprToDestExprMutator : public ir::IRMutator<> {
   void Visit(const ir::For* for_node, Expr* op) override;
   void Visit(const ir::Block* block_node, Expr* op) override;
   void Visit(const ir::ScheduleBlockRealize* realize, Expr* op) override;
+  void Visit(const ir::_Var_* var, Expr* op) override;
 
  private:
   ir::Expr source_;
