[PIR Inference] add map_op_to_another_pass (#61188)

bukejiyu · web-flow · commit 2e1ca29e952b · 2024-01-30T20:02:56.000+08:00
* add depthwise_conv_to_conv_pass

* update

* update
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -117,6 +117,7 @@
 #include "paddle/fluid/pir/transforms/fusion/multihead_matmul_fuse_pass.h"
 #include "paddle/fluid/pir/transforms/identity_op_clean_pass.h"
 #include "paddle/fluid/pir/transforms/inplace_pass.h"
+#include "paddle/fluid/pir/transforms/map_op_to_another_pass.h"
 #include "paddle/fluid/pir/transforms/params_sync_among_devices_pass.h"
 #include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
 #include "paddle/fluid/pir/transforms/replace_fetch_with_shadow_output_pass.h"
@@ -798,6 +799,7 @@ bool AnalysisPredictor::PrepareExecutor() {
         ::pir::PassManager gpu_pm(::pir::IrContext::Instance(), 2);
         //----------------------------------------------------------------------------------------------//
         // Functional pass
+        gpu_pm.AddPass(::pir::CreateMapOpToAnotherPass());
         gpu_pm.AddPass(::pir::CreateIdentityOpCleanPass());
         //----------------------------------------------------------------------------------------------//
 
diff --git a/paddle/fluid/pir/transforms/fusion/conv2d_add_act_fuse_pass.cc b/paddle/fluid/pir/transforms/fusion/conv2d_add_act_fuse_pass.cc
@@ -56,7 +56,7 @@ class Conv2dAddActFusePattern
     if (next_op->isa<paddle::dialect::ReluOp>()) {
       act_name = "relu";
     }
-#if CUDNN_VERSION >= 8000 && CUDNN_VERSION < 8700
+#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 8000 && CUDNN_VERSION < 8700
     if (next_op->isa<paddle::dialect::TanhOp>()) {
       act_name = "tanh";
     } else if (next_op->isa<paddle::dialect::SigmoidOp>()) {
@@ -136,7 +136,7 @@ class Conv2dAdd2ActFusePattern
     if (next_op->isa<paddle::dialect::ReluOp>()) {
       act_name = "relu";
     }
-#if CUDNN_VERSION >= 8000 && CUDNN_VERSION < 8700
+#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 8000 && CUDNN_VERSION < 8700
     if (next_op->isa<paddle::dialect::TanhOp>()) {
       act_name = "tanh";
     } else if (next_op->isa<paddle::dialect::SigmoidOp>()) {
diff --git a/paddle/fluid/pir/transforms/map_op_to_another_pass.cc b/paddle/fluid/pir/transforms/map_op_to_another_pass.cc
@@ -0,0 +1,84 @@
+// 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/fluid/pir/transforms/map_op_to_another_pass.h"
+
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
+
+#include "paddle/pir/pass/pass.h"
+#include "paddle/pir/pass/pass_registry.h"
+
+namespace {
+
+class DepthWiseConv2d2Conv2dPattern : public paddle::drr::DrrPatternBase {
+ public:
+  void operator()(paddle::drr::DrrPatternContext *ctx) const override {
+    paddle::drr::SourcePattern pat = ctx->SourcePattern();
+    const auto &depthwise_conv2d_op =
+        pat.Op(paddle::dialect::DepthwiseConv2dOp::name(),
+               {{"strides", pat.Attr("strides")},
+                {"paddings", pat.Attr("paddings")},
+                {"padding_algorithm", pat.Attr("padding_algorithm")},
+                {"dilations", pat.Attr("dilations")},
+                {"groups", pat.Attr("groups")},
+                {"data_format", pat.Attr("data_format")}});
+    depthwise_conv2d_op({&pat.Tensor("input"), &pat.Tensor("filter")},
+                        {&pat.Tensor("depthwise_conv2d_out")});
+    pat.RequireNativeCall(
+        [](const paddle::drr::MatchContext &match_ctx) -> bool {
+#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 8100
+          auto groups = match_ctx.Attr<int>("groups");
+          return groups > 1;
+#else
+          return false;
+#endif
+        });
+
+    paddle::drr::ResultPattern res = pat.ResultPattern();
+    const auto &conv2d =
+        res.Op(paddle::dialect::Conv2dOp::name(),
+               {{"strides", pat.Attr("strides")},
+                {"paddings", pat.Attr("paddings")},
+                {"padding_algorithm", pat.Attr("padding_algorithm")},
+                {"dilations", pat.Attr("dilations")},
+                {"groups", pat.Attr("groups")},
+                {"data_format", pat.Attr("data_format")}});
+    conv2d({&res.Tensor("input"), &res.Tensor("filter")},
+           {&res.Tensor("depthwise_conv2d_out")});
+  }
+  std::string name() const override { return "DepthWiseConv2d2Conv2dPattern"; }
+};
+
+class MapOpToAnotherPass : public pir::PatternRewritePass {
+ public:
+  MapOpToAnotherPass() : pir::PatternRewritePass("map_op_to_another_pass", 2) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add(DepthWiseConv2d2Conv2dPattern().Build(context));
+    return ps;
+  }
+};
+
+}  // namespace
+
+namespace pir {
+
+std::unique_ptr<Pass> CreateMapOpToAnotherPass() {
+  return std::make_unique<MapOpToAnotherPass>();
+}
+}  // namespace pir
+
+REGISTER_IR_PASS(map_op_to_another_pass, MapOpToAnotherPass);
diff --git a/paddle/fluid/pir/transforms/map_op_to_another_pass.h b/paddle/fluid/pir/transforms/map_op_to_another_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
+//
+//     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 <memory>
+#include "paddle/pir/core/dll_decl.h"
+
+namespace pir {
+
+class Pass;
+
+IR_API std::unique_ptr<Pass> CreateMapOpToAnotherPass();
+
+}  // namespace pir
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -57,6 +57,7 @@
 #include "paddle/fluid/pir/transforms/fusion/multihead_matmul_fuse_pass.h"
 #include "paddle/fluid/pir/transforms/identity_op_clean_pass.h"
 #include "paddle/fluid/pir/transforms/inplace_pass.h"
+#include "paddle/fluid/pir/transforms/map_op_to_another_pass.h"
 #include "paddle/fluid/pir/transforms/replace_fetch_with_shadow_output_pass.h"
 #include "paddle/fluid/pir/transforms/shape_optimization_pass.h"
 #include "paddle/fluid/pybind/eager_utils.h"
@@ -129,6 +130,7 @@ USE_PIR_PASS(fused_linear_param_grad_add_pass);
 USE_PIR_PASS(inplace_pass);
 USE_PIR_PASS(replace_fetch_with_shadow_output_pass);
 USE_PIR_PASS(identity_op_clean_pass);
+USE_PIR_PASS(map_op_to_another_pass);
 USE_PIR_PASS(matmul_scale_fuse_pass);
 USE_PIR_PASS(fc_fuse_pass);
 USE_PIR_PASS(fc_elementwise_layernorm_fuse_pass);
diff --git a/test/ir/pir/CMakeLists.txt b/test/ir/pir/CMakeLists.txt
@@ -28,6 +28,7 @@ foreach(target ${TEST_IR_SYSTEM_CASES})
 endforeach()
 
 set_tests_properties(test_pd_inplace_pass PROPERTIES TIMEOUT 60)
+set_tests_properties(test_map_op_another_pass PROPERTIES TIMEOUT 60)
 
 py_test_modules(
   test_subgraph_exporter
diff --git a/test/ir/pir/test_map_op_another_pass.py b/test/ir/pir/test_map_op_another_pass.py
@@ -0,0 +1,97 @@
+# 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.
+
+import os
+import re
+import unittest
+
+import numpy as np
+from fused_pass.pass_test import PassTest
+
+import paddle
+import paddle.nn.functional as F
+from paddle.base import core
+from paddle.pir.core import create_parameter
+
+paddle.enable_static()
+
+
+def get_cuda_version():
+    result = os.popen("nvcc --version").read()
+    regex = r'release (\S+),'
+    match = re.search(regex, result)
+    if match:
+        num = str(match.group(1))
+        integer, decimal = num.split('.')
+        return int(integer) * 1000 + int(float(decimal) * 10)
+    else:
+        return -1
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() or get_cuda_version() < 8100,
+    "DepthwiseConv2ConvPattern requires CUDA >= 8100",
+)
+class TestDepthwiseConv2ConvPattern(PassTest):
+    r""" """
+
+    def is_program_valid(self, program=None):
+        return True
+
+    def sample_program(self):
+        with paddle.pir_utils.IrGuard():
+            main_prog = paddle.static.Program()
+            start_prog = paddle.static.Program()
+            for x_shape in [[3, 32, 150, 150]]:
+                for conv2d_filter_shape in [[32, 1, 3, 3]]:
+                    with paddle.pir.core.program_guard(main_prog, start_prog):
+                        x = paddle.static.data(
+                            name='x', shape=x_shape, dtype='float32'
+                        )
+                        initializer = paddle.nn.initializer.Assign(
+                            np.random.rand(32, 1, 3, 3)
+                        )
+                        conv2d_filter = create_parameter(
+                            shape=conv2d_filter_shape,
+                            dtype='float32',
+                            initializer=initializer,
+                            padding="SAME",
+                            dilation=[1, 1],
+                            stride=[1, 1],
+                        )
+                        depthwise_conv2d_out = F.conv2d(
+                            x, conv2d_filter, groups=32, data_format="NCHW"
+                        )
+                        out = paddle.assign(depthwise_conv2d_out)
+                        self.pass_list = ['map_op_to_another_pass']
+                        self.feeds = {
+                            "x": np.random.random(x_shape).astype("float32"),
+                        }
+                        self.fetch_list = [out]
+                        self.valid_op_map = {
+                            "pd_op.depthwise_conv2d": 0,
+                            "pd_op.conv2d": 1,
+                        }
+                        yield [main_prog, start_prog], False
+
+    def setUp(self):
+        if core.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+
+if __name__ == "__main__":
+    unittest.main()
