【CINN】ReAssignPipeline (#63859)

Vvsmile · web-flow · commit 151a374becf7 · 2024-05-29T10:40:07.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -34,6 +34,7 @@
 #include "paddle/cinn/lang/placeholder.h"
 #include "paddle/cinn/optim/eliminate_common_global_memory_read.h"
 #include "paddle/cinn/optim/if_fusion.h"
+#include "paddle/cinn/optim/rearrange_load_instruction.h"
 #include "paddle/cinn/optim/schedule_block_dce.h"
 #include "paddle/cinn/optim/transform_gpu_forloop.h"
 #include "paddle/common/ddim.h"
@@ -850,6 +851,7 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
     }
     // 4.Apply low level pass
     func = optim::Optimize(Expr(func), target_, false).as_lowered_func_ref();
+    optim::RearrangeLoadInstruction(&(func->body));
     lowered_funcs.push_back(std::move(func));
   }
 
diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
@@ -400,6 +400,38 @@ Expr Store::index() const {
   return res;
 }
 
+void Store::replace(Expr old_op, Expr new_op) {
+  if (value == old_op) {
+    value = new_op;
+  }
+  if (tensor == old_op) {
+    tensor = new_op;
+  }
+  for (int i = 0; i < indices.size(); i++) {
+    if (indices[i] == old_op) {
+      indices[i] = new_op;
+    }
+  }
+}
+
+void Select::replace(Expr old_op, Expr new_op) {
+  if (condition == old_op) {
+    condition = new_op;
+  }
+  if (true_value == old_op) {
+    true_value = new_op;
+  }
+  if (false_value == old_op) {
+    false_value = new_op;
+  }
+}
+
+void Cast::replace(Expr old_op, Expr new_op) {
+  if (v() == old_op) {
+    v() = new_op;
+  }
+}
+
 const std::string &Store::name() const {
   auto *t = tensor.As<ir::_Tensor_>();
   CHECK(t);
@@ -495,6 +527,20 @@ Expr Call::Make(Type type,
   node->attrs = attrs;
   return Expr(node);
 }
+
+void Call::replace(Expr old_op, Expr new_op) {
+  for (int i = 0; i < read_args.size(); i++) {
+    if (read_args[i] == old_op) {
+      read_args[i] = new_op;
+    }
+  }
+  for (int i = 0; i < write_args.size(); i++) {
+    if (read_args[i] == old_op) {
+      read_args[i] = new_op;
+    }
+  }
+}
+
 std::vector<Expr *> Call::expr_fields() {
   std::vector<Expr *> res;
   for (auto &x : read_args) res.push_back(&x);
diff --git a/paddle/cinn/ir/ir.h b/paddle/cinn/ir/ir.h
@@ -66,6 +66,8 @@ struct Cast : public ExprNode<Cast> {
   Expr& v() { return operand(0); }
   const Expr& v() const { return operand(0); }
 
+  void replace(Expr old_op, Expr new_op);
+
   void Verify() const override;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Cast;
@@ -123,6 +125,7 @@ struct Div : public BinaryOpNode<Div> {
 
   static Expr Make(Expr a, Expr b);
   void Verify() const override;
+
   static const IrNodeTy _node_type_ = IrNodeTy::Div;
 };
 
@@ -357,6 +360,8 @@ struct Call : public ExprNode<Call> {
 
   void Verify() const override;
 
+  void replace(Expr old_op, Expr new_op);
+
   inline size_t total_args_count() const {
     return read_args.size() + write_args.size();
   }
@@ -521,6 +526,7 @@ struct Select : public ExprNode<Select> {
     return {&condition, &true_value, &false_value};
   }
 
+  void replace(Expr old_op, Expr new_op);
   static const IrNodeTy _node_type_ = IrNodeTy::Select;
 };
 
@@ -572,6 +578,8 @@ struct Store : public ExprNode<Store>, public LoadStoreAddrMnger {
 
   const std::string& name() const;
 
+  void replace(Expr old_op, Expr new_op);
+
   Type type() const override;
 
   Expr index() const;
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/cinn/ir/ir_base.h"
 
+#include <sstream>
 #include "paddle/cinn/common/cinn_value.h"
 #include "paddle/cinn/common/common.h"
 #include "paddle/cinn/ir/buffer.h"
@@ -238,6 +239,15 @@ const Expr &IrNode::operand(int i) {
 
 void IrNode::set_type(Type type) { type_ = type; }
 
+void IrNode::replace(Expr old_op, Expr new_op) {
+  std::stringstream ss;
+  ss << "Not Implemented, The node:(" << node_type() << ") has an old_op: ("
+     << old_op.node_type() << ") should be replaced with new_op: ("
+     << new_op.node_type() << ") but not Implemented";
+
+  PADDLE_THROW(phi::errors::Unimplemented(ss.str()));
+}
+
 void IrNode::convert_int32_to_int64() {
   CHECK(type_ == Int(64) || type_ == Int(32) || type_.is_unk())
       << "Current only support convert int32_t to int64_t, but get type is "
diff --git a/paddle/cinn/ir/ir_base.h b/paddle/cinn/ir/ir_base.h
@@ -166,6 +166,7 @@ class IrNode : public cinn::common::Object {
   //! Elevate int32 to int64 if needed
   virtual void convert_int32_to_int64();
 
+  virtual void replace(Expr old_op, Expr new_op);
   //! Get i-th operand
   const Expr& operand(int i);
 
@@ -430,6 +431,13 @@ struct BinaryOpNode : public ExprNode<T> {
 
   Type type() const override { return a().type(); }
 
+  void replace(Expr old_op, Expr new_op) {
+    for (int i = 0; i < operands().size(); i++) {
+      if (operands()[i] == old_op) {
+        operands()[i] = new_op;
+      }
+    }
+  }
   std::vector<Expr*> expr_fields() override { return {&a(), &b()}; }
   std::vector<const Expr*> expr_fields() const override { return {&a(), &b()}; }
 
diff --git a/paddle/cinn/optim/CMakeLists.txt b/paddle/cinn/optim/CMakeLists.txt
@@ -32,7 +32,8 @@ gather_srcs(
   schedule_block_dce.cc
   eliminate_common_factor_of_local_index.cc
   if_fusion.cc
-  eliminate_common_global_memory_read.cc)
+  eliminate_common_global_memory_read.cc
+  rearrange_load_instruction.cc)
 
 if(WITH_CUDA)
   gather_srcs(cinnapi_src SRCS transform_gpu_forloop.cc)
diff --git a/paddle/cinn/optim/optimize.cc b/paddle/cinn/optim/optimize.cc
@@ -28,6 +28,7 @@
 #include "paddle/cinn/optim/lower_function_call_bind_vars.h"
 #include "paddle/cinn/optim/lower_intrin.h"
 #include "paddle/cinn/optim/map_extern_call.h"
+#include "paddle/cinn/optim/rearrange_load_instruction.h"
 #include "paddle/cinn/optim/remove_schedule_block.h"
 #include "paddle/cinn/optim/replace_const_param_to_integer.h"
 #include "paddle/cinn/optim/replace_cross_thread_reduction.h"
diff --git a/paddle/cinn/optim/rearrange_load_instruction.cc b/paddle/cinn/optim/rearrange_load_instruction.cc
@@ -0,0 +1,175 @@
+// Copyright (c) 2024 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
+//
+//     https://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/optim/rearrange_load_instruction.h"
+
+#include <stack>
+#include "paddle/cinn/adt/map_expr.h"
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_base.h"
+#include "paddle/cinn/ir/ir_mutator.h"
+#include "paddle/cinn/ir/ir_printer.h"
+#include "paddle/cinn/ir/utils/ir_compare.h"
+#include "paddle/cinn/optim/ir_simplify.h"
+
+#define VisitImpl(_TYPE)                                 \
+  void Visit(const ir::_TYPE *op, Expr *expr) override { \
+    auto old_last_op = last_op;                          \
+    auto *ncop = const_cast<ir::_TYPE *>(op);            \
+    last_op = Expr(ncop);                                \
+    ir::IRMutator<>::Visit(op, expr);                    \
+    last_op = old_last_op;                               \
+  }
+
+namespace cinn {
+namespace optim {
+
+namespace {
+
+struct RearrangeLoadInstructionMutator : public ir::IRMutator<Expr *> {
+  RearrangeLoadInstructionMutator() {
+    is_inner_store = false;
+    last_op = Expr(nullptr);
+  }
+  void operator()(Expr *expr) { Visit(expr, expr); }
+
+ private:
+  void Visit(const ir::Load *op, Expr *expr) override {
+    if (is_inner_store) {
+      if (op->tensor.as_tensor_ref()->buffer->memory_type ==
+              ir::MemoryType::GPULocal ||
+          op->tensor.as_tensor_ref()->buffer->memory_type ==
+              ir::MemoryType::GPUShared)
+        return;
+
+      auto local_var =
+          ir::_Var_::Make(common::UniqName("local_var"), op->type());
+      auto let_op = ir::Let::Make(local_var, const_cast<ir::Load *>(op));
+      let_list.push_back(let_op);
+      last_op->replace(Expr(const_cast<ir::Load *>(op)), local_var);
+    }
+  }
+
+  void Visit(const ir::Store *op, Expr *expr) override {
+    auto old_last_op = last_op;
+    last_op = Expr(const_cast<ir::Store *>(op));
+    is_inner_store = true;
+    ir::IRMutator<>::Visit(op, expr);
+    is_inner_store = false;
+    last_op = old_last_op;
+  }
+
+  void Visit(const ir::Block *op, Expr *expr) override {
+    auto old_last_op = last_op;
+    last_op = Expr(const_cast<ir::Block *>(op));
+    int old_let_size = let_list.size();
+    int old_stmts_size = stmts_list.size();
+
+    for (auto &stmt : op->stmts) {
+      IRVisitorRequireReImpl<void, Expr *>::Visit(
+          &stmt, const_cast<ir::Expr *>(&stmt));
+      stmts_list.push_back(stmt);
+    }
+
+    if (let_list.size() > old_let_size) {
+      replaceBlock(const_cast<ir::Block *>(op), old_let_size, old_stmts_size);
+    }
+    last_op = old_last_op;
+  }
+
+  void Visit(const ir::Expr *op, Expr *expr) override {
+    ir::IRMutator<>::Visit(op, expr);
+  }
+
+  void Visit(const ir::Select *op, Expr *expr) override {}
+  void Visit(const ir::Broadcast *op, Expr *expr) override {}
+
+  void replaceBlock(ir::Block *op, int old_let_size, int old_stmts_size) {
+    std::vector<Expr> new_stmts;
+
+    for (int i = old_let_size; i < let_list.size(); i++) {
+      new_stmts.push_back(let_list[i]);
+    }
+
+    for (int i = old_stmts_size; i < stmts_list.size(); i++) {
+      new_stmts.push_back(stmts_list[i]);
+    }
+
+    while (let_list.size() > old_let_size) {
+      let_list.pop_back();
+    }
+
+    while (stmts_list.size() > old_stmts_size) {
+      stmts_list.pop_back();
+    }
+
+    op->stmts = new_stmts;
+  }
+
+  std::unordered_map<std::string, ir::Expr> collection_name_map_expr;
+  std::vector<Expr> let_list;
+  std::vector<Expr> stmts_list;
+  bool is_inner_store;
+  Expr last_op;
+
+  VisitImpl(ScheduleBlock);
+  VisitImpl(For);
+  VisitImpl(Cast);
+  VisitImpl(PolyFor);
+  VisitImpl(Call);
+  VisitImpl(_Module_);
+  VisitImpl(_Var_);
+  VisitImpl(Alloc);
+  VisitImpl(Free);
+  VisitImpl(_Buffer_);
+  VisitImpl(_Tensor_);
+  VisitImpl(_LoweredFunc_);
+  VisitImpl(Let);
+  VisitImpl(Reduce);
+  VisitImpl(Ramp);
+  VisitImpl(FracOp);
+  VisitImpl(Product);
+  VisitImpl(Sum);
+  VisitImpl(PrimitiveNode);
+  VisitImpl(IntrinsicOp);
+  VisitImpl(_BufferRange_);
+  VisitImpl(_Dim_);
+
+  VisitImpl(Add);
+  VisitImpl(Sub);
+  VisitImpl(Mul);
+  VisitImpl(Div);
+  VisitImpl(Mod);
+  VisitImpl(EQ);
+  VisitImpl(NE);
+  VisitImpl(LT);
+  VisitImpl(LE);
+  VisitImpl(GT);
+  VisitImpl(GE);
+  VisitImpl(And);
+  VisitImpl(Or);
+  VisitImpl(Not);
+  VisitImpl(Min);
+  VisitImpl(Max);
+  VisitImpl(Minus)
+};
+}  // namespace
+
+void RearrangeLoadInstruction(Expr *expr) {
+  RearrangeLoadInstructionMutator collector;
+  collector(expr);
+}
+
+}  // namespace optim
+}  // namespace cinn
diff --git a/paddle/cinn/optim/rearrange_load_instruction.h b/paddle/cinn/optim/rearrange_load_instruction.h
@@ -0,0 +1,27 @@
+// Copyright (c) 2024 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
+//
+//     https://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 "paddle/cinn/ir/ir.h"
+
+namespace cinn {
+namespace optim {
+
+/*
+ * Do fusion with the adjaccnt if-block.
+ */
+void RearrangeLoadInstruction(Expr *expr);
+
+}  // namespace optim
+}  // namespace cinn

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@`
`34`	`34`	`#include "paddle/cinn/lang/placeholder.h"`
`35`	`35`	`#include "paddle/cinn/optim/eliminate_common_global_memory_read.h"`
`36`	`36`	`#include "paddle/cinn/optim/if_fusion.h"`
	`37`	`+#include "paddle/cinn/optim/rearrange_load_instruction.h"`
`37`	`38`	`#include "paddle/cinn/optim/schedule_block_dce.h"`
`38`	`39`	`#include "paddle/cinn/optim/transform_gpu_forloop.h"`
`39`	`40`	`#include "paddle/common/ddim.h"`
`@@ -850,6 +851,7 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(`
`850`	`851`	`}`
`851`	`852`	`// 4.Apply low level pass`
`852`	`853`	`func = optim::Optimize(Expr(func), target_, false).as_lowered_func_ref();`
	`854`	`+ optim::RearrangeLoadInstruction(&(func->body));`
`853`	`855`	`lowered_funcs.push_back(std::move(func));`
`854`	`856`	`}`
`855`	`857`