fix bug and change pass place in d2s. (#68759)

liuruyan · web-flow · commit 2275419dda0f · 2024-10-17T16:44:04.000+08:00
diff --git a/paddle/fluid/pir/transforms/general/auto_layout_pass.cc b/paddle/fluid/pir/transforms/general/auto_layout_pass.cc
@@ -39,7 +39,7 @@ namespace {
 
 class AutoLayoutPass : public pir::Pass {
  public:
-  AutoLayoutPass() : pir::Pass("auto_layout_pass", 3) {}
+  AutoLayoutPass() : pir::Pass("auto_layout_pass", 2) {}
   void Run(pir::Operation* op) override {
     for (size_t i = 0; i < op->num_regions(); ++i) {
       auto& region = op->region(i);
@@ -138,8 +138,10 @@ class AutoLayoutPass : public pir::Pass {
       if (op->HasTrait<pir::ImmutableLayoutTrait>()) continue;
       if (op->operands().size() == 0) continue;
 
-      // NHWC ops branch, Only support conv2d now, it will add white list later.
-      if (op->isa<paddle::dialect::Conv2dOp>()) {
+      // NHWC ops branch, Only support conv2d and fused_conv2d_add_act now, it
+      // will add white list later.
+      if (op->isa<paddle::dialect::Conv2dOp>() ||
+          op->isa<paddle::dialect::FusedConv2dAddActOp>()) {
         if (op->HasAttribute("data_format") &&
             op->attribute<pir::StrAttribute>("data_format").AsString() ==
                 "NCHW") {
@@ -160,14 +162,8 @@ class AutoLayoutPass : public pir::Pass {
   // Skip the operand which is not dense tensor or not 4-D tensor, they don't
   // need transpose.
   bool JudgeValue(const pir::Value& value) {
-    if (!value) {
-      PADDLE_THROW(common::errors::Fatal(
-          "value is null, please check the input tensor."));
-    }
-    if (!value.type()) {
-      PADDLE_THROW(common::errors::Fatal(
-          "value type is null, please check the input tensor type."));
-    }
+    if (!value) return false;
+    if (!value.type()) return false;
     if (auto type = value.type().dyn_cast<paddle::dialect::DenseTensorType>()) {
       return type.dims().size() == 4;
     }
@@ -204,7 +200,6 @@ class AutoLayoutPass : public pir::Pass {
                            pir::Builder& builder) {  // NOLINT
     builder.SetInsertionPointAfter(op);
     for (auto& result : op->results()) {
-      if (result.use_empty()) continue;
       if (!JudgeValue(result)) continue;
       auto transpose_op =
           builder.Build<paddle::dialect::TransposeOp>(result, NHWC2NCHW_);
diff --git a/paddle/fluid/pir/transforms/general/auto_layout_simplify_pass.cc b/paddle/fluid/pir/transforms/general/auto_layout_simplify_pass.cc
@@ -85,7 +85,7 @@ class RedundantTransposePattern
 class AutoLayoutSimplifyPass : public pir::PatternRewritePass {
  public:
   AutoLayoutSimplifyPass()
-      : pir::PatternRewritePass("auto_layout_simplify_pass", 3) {}
+      : pir::PatternRewritePass("auto_layout_simplify_pass", 2) {}
   pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
     pir::RewritePatternSet ps(context);
     ps.Add<RedundantTransposePattern>(context);
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -674,11 +674,6 @@ def _get_scope(self, program_id=None, use_scope_cache=False):
     # whole
     @switch_to_static_graph
     def _create_program(self, is_infer_mode=False):
-        if auto_layout_is_enabled():
-            pm = paddle.pir.PassManager(3)
-            pm.add_pass("auto_layout_pass", {})
-            pm.add_pass("auto_layout_simplify_pass", {})
-            pm.run(self._origin_main_program)
 
         if is_infer_mode:
 
@@ -704,6 +699,11 @@ def pass_fn(forward_program, backward_program, program_name_attr):
 
             # TODO(xiongkun) who to transfer the pruning program?
             infer_program = self.origin_runnable_program.clone()
+            if auto_layout_is_enabled():
+                pm = paddle.pir.PassManager(2)
+                pm.add_pass("auto_layout_pass", {})
+                pm.add_pass("auto_layout_simplify_pass", {})
+                pm.run(infer_program.program)
             for hooker in self._hookers:
                 hooker.after_infer(infer_program)
             infer_program.apply_pir_program_pass(pass_fn)
@@ -712,6 +712,12 @@ def pass_fn(forward_program, backward_program, program_name_attr):
             train_program: RunnableProgram = (
                 self.origin_runnable_program.clone()
             )
+            # Author(liujinnan): auto_layout_pass should be applied to the original_program, before append backward. So we put it here.
+            if auto_layout_is_enabled():
+                pm = paddle.pir.PassManager(2)
+                pm.add_pass("auto_layout_pass", {})
+                pm.add_pass("auto_layout_simplify_pass", {})
+                pm.run(train_program.program)
             train_program = self._append_backward_desc(train_program)
             # Note: Only set grad type once after initializing train program. So we put it here.
             self._set_grad_type(self._params, train_program)
