[Inference] add add_shadow_output_after_dead_parameter_pass (#68476)

* add add_shadow_output_after_dead_parameter_pass

* fix

* fix

* fix ci

Author: yuanlehome

paddle/fluid/inference/api/paddle_pass_builder.cc

@@ -581,6 +581,7 @@ IpuPassStrategy::IpuPassStrategy() : PassStrategy({}) {
 
 const std::vector<std::string> kPirGpuPasses{
     // Functional pass
+    "add_shadow_output_after_dead_parameter_pass",
     "delete_quant_dequant_linear_op_pass",
     "delete_weight_dequant_linear_op_pass",
     "map_op_to_another_pass",
@@ -609,6 +610,7 @@ const std::vector<std::string> kPirGpuPasses{
 
 const std::vector<std::string> kPirXpuPasses{
     // Functional pass
+    "add_shadow_output_after_dead_parameter_pass",
     "delete_quant_dequant_linear_op_pass",
     "delete_weight_dequant_linear_op_pass",
     "map_op_to_another_pass",
@@ -622,7 +624,8 @@ const std::vector<std::string> kPirXpuPasses{
     "fc_xpu_fuse_pass"};
 
 const std::vector<std::string> kPirMkldnnPasses {
-  "delete_quant_dequant_linear_op_pass",          //
+  "add_shadow_output_after_dead_parameter_pass",
+      "delete_quant_dequant_linear_op_pass",      //
       "delete_weight_dequant_linear_op_pass",     //
       "depthwise_conv_onednn_pass",               //
       "squeeze_transpose_onednn_fuse_pass",       //
@@ -663,13 +666,15 @@ const std::vector<std::string> kPirMkldnnPasses {
 };
 
 const std::vector<std::string> kPirMkldnnBf16Passes{
+    "add_shadow_output_after_dead_parameter_pass",
     "cpu_bfloat16_placement_pass",
     "cpu_bfloat16_pass",
     "cpu_bfloat16_type_placement_pass",
     "cpu_bf16_quantize_squash_pass",
 };
 
 const std::vector<std::string> kPirCpuPasses{
+    "add_shadow_output_after_dead_parameter_pass",
     "delete_quant_dequant_linear_op_pass",
     "delete_weight_dequant_linear_op_pass"};
 

paddle/fluid/pir/transforms/general/add_shadow_output_after_dead_parameter_pass.cc

@@ -0,0 +1,70 @@
+// 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
+//
+//     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/fluid/pir/transforms/general/add_shadow_output_after_dead_parameter_pass.h"
+
+#include "paddle/pir/include/core/builtin_op.h"
+#include "paddle/pir/include/pass/pass.h"
+#include "paddle/pir/include/pass/pass_registry.h"
+
+namespace {
+
+class AddShadowOutputAfterDeadParameterPattern
+    : public pir::OpRewritePattern<pir::ParameterOp> {
+ public:
+  using pir::OpRewritePattern<pir::ParameterOp>::OpRewritePattern;
+  bool MatchAndRewrite(
+      pir::ParameterOp op,
+      pir::PatternRewriter& rewriter) const override {  // NOLINT
+    if (!op->use_empty()) {
+      return false;
+    }
+    rewriter.SetInsertionPointToBlockEnd(op->GetParent());
+    rewriter.Build<pir::ShadowOutputOp>(op->result(0), op.param_name());
+    return true;
+  }
+};
+
+class AddShadowOutputAfterDeadParameterPass : public pir::PatternRewritePass {
+ public:
+  AddShadowOutputAfterDeadParameterPass()
+      : pir::PatternRewritePass("add_shadow_output_after_dead_parameter_pass",
+                                0) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add<AddShadowOutputAfterDeadParameterPattern>(context);
+    return ps;
+  }
+
+  bool CanApplyOn(pir::Operation* op) const override {
+    return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
+  }
+
+ private:
+  pir::FrozenRewritePatternSet patterns_;
+};
+
+}  // namespace
+
+namespace pir {
+
+std::unique_ptr<pir::Pass> CreateAddShadowOutputAfterDeadParameterPass() {
+  return std::make_unique<AddShadowOutputAfterDeadParameterPass>();
+}
+
+}  // namespace pir
+
+REGISTER_IR_PASS(add_shadow_output_after_dead_parameter_pass,
+                 AddShadowOutputAfterDeadParameterPass);

paddle/fluid/pir/transforms/general/add_shadow_output_after_dead_parameter_pass.h

@@ -0,0 +1,26 @@
+// 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
+//
+//     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 <memory>
+#include "paddle/pir/include/core/dll_decl.h"
+
+namespace pir {
+
+class Pass;
+
+IR_API std::unique_ptr<Pass> CreateAddShadowOutputAfterDeadParameterPass();
+
+}  // namespace pir

paddle/fluid/pir/transforms/passes.h

@@ -48,6 +48,7 @@ USE_PIR_PASS(transfer_layout_pass);
 USE_PIR_PASS(fused_rotary_position_embedding_pass);
 USE_PIR_PASS(horizontal_fuse_pass);
 USE_PIR_PASS(common_subexpression_elimination_pass);
+USE_PIR_PASS(add_shadow_output_after_dead_parameter_pass);
 
 #ifdef PADDLE_WITH_DNNL
 USE_PIR_PASS(depthwise_conv_onednn_pass);

test/ir/inference/test_inference_predictor_run_pir.py

@@ -27,6 +27,7 @@ def __init__(self):
         super().__init__()
         self.fc1 = paddle.nn.Linear(4, 4)
         self.fc2 = paddle.nn.Linear(4, 4)
+        self.register_buffer("buffer", paddle.randn([5, 1]))
 
     def forward(self, x1, x2):
         y1 = self.fc1(x1)
@@ -122,6 +123,9 @@ def init_pir_predictor(self):
         # config.enable_memory_optim()
         config.enable_new_executor()
         config.enable_new_ir()
+        config.switch_ir_debug(
+            True, ['add_shadow_output_after_dead_parameter_pass']
+        )
         predictor = create_predictor(config)
         return predictor