[CINN+PIR]Refine into NoBufferRelateOps and Add BatchNorm.reserve_space (#60800)

Aurelius84 · web-flow · commit 47b2759bf4ce · 2024-01-16T17:24:02.000+08:00
* [CINN+PIR]Refine into NoBufferRelateOps and Add BatchNorm.reserve_space

* [CINN+PIR]Refine into NoBufferRelateOps and Add BatchNorm.reserve_space

* add dropout

* updaate
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -144,19 +144,26 @@ const std::unordered_set<std::string> SpecialLowerOps = {
     SelectInputOp::name(),
     "cinn_runtime.jit_kernel"};
 
-const std::unordered_map<std::string, uint32_t> XShapeRelatedOps = {
+const std::unordered_map<std::string, uint32_t> NoBufferRelatedOps = {
     {paddle::dialect::ReshapeOp::name(), /*xshape_idx*/ 1U},
+    {paddle::dialect::Reshape_Op::name(), /*xshape_idx*/ 1U},
     {paddle::dialect::SqueezeOp::name(), /*xshape_idx*/ 1U},
+    {paddle::dialect::Squeeze_Op::name(), /*xshape_idx*/ 1U},
+    {paddle::dialect::UnsqueezeOp::name(), /*xshape_idx*/ 1U},
+    {paddle::dialect::Unsqueeze_Op::name(), /*xshape_idx*/ 1U},
     {paddle::dialect::FlattenOp::name(), /*xshape_idx*/ 1U},
+    {paddle::dialect::Flatten_Op::name(), /*xshape_idx*/ 1U},
+    {paddle::dialect::BatchNormOp::name(), /*reserve_space*/ 5U},
+    {paddle::dialect::BatchNorm_Op::name(), /*reserve_space*/ 5U},
 };
 
 static bool NeedSkipPlaceTransfer(const pir::Operation* op) {
   bool need_skip = false;
   if (op->isa<paddle::dialect::FetchOp>()) {
     auto define_op_name = op->operand_source(0).defining_op()->name();
     uint32_t index = op->operand_source(0).dyn_cast<pir::OpResult>().index();
-    need_skip = XShapeRelatedOps.count(define_op_name) > 0 &&
-                (XShapeRelatedOps.at(define_op_name) == index);
+    need_skip = NoBufferRelatedOps.count(define_op_name) > 0 &&
+                (NoBufferRelatedOps.at(define_op_name) == index);
   }
   return need_skip;
 }
