[Paddle TensorRT] add pd_op.max converter (#67956)

Layssy · web-flow · commit 95052cc8bb3a · 2024-09-06T15:11:41.000+08:00
* 添加maxop的marker

* 添加maxop的converter

* fix converter单测文件

* 删除max的marker 单测

* 将max的converter 移动 到math.py

* 添加maxconverter 单测时间限制

* fix max converter

* fix  CmakeList.txt

* rename  test/tensorrt/test_converter_max.py -&gt; test/tensorrt/test_converter_math.py

* fix CmakeLists.txt

* fix  CMakeLists.txt
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1088,6 +1088,32 @@ class RemainderOpPattern
     return true;
   }
 };
+class MaxOpPattern : public pir::OpRewritePattern<paddle::dialect::MaxOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::MaxOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::MaxOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!op->HasAttribute("keepdim")) {
+      VLOG(3) << "the max does not have attr keep_dim ";
+      return false;
+    }
+    pir::Value x = op.operand_source(0);
+    auto x_dtype = pir::GetDataTypeFromValue(x);
+    if (!(x_dtype.isa<pir::Float32Type>() || x_dtype.isa<pir::Float64Type>() ||
+          x_dtype.isa<pir::Int32Type>() || x_dtype.isa<pir::Int64Type>())) {
+      VLOG(3) << "max input data type must be int32 or int64 or "
+                 "float32 or "
+                 "float64";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -1153,6 +1179,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<MaximumOpPattern>(context));
     ps.Add(std::make_unique<FloorDivideOpPattern>(context));
     ps.Add(std::make_unique<RemainderOpPattern>(context));
+    ps.Add(std::make_unique<MaxOpPattern>(context));
     return ps;
   }
 };
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -32,6 +32,7 @@
 )
 from paddle.tensorrt.converter_utils import (
     broadcast,
+    get_axes_for_reduce_op,
 )
 
 
@@ -84,3 +85,28 @@ def scale_converter(network, paddle_op, inputs):
         power=power_weight,
     )
     return scale_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.max", trt_version="8.x")
+def max_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    axis = paddle_op.operands()[1].source().get_defining_op().attrs()["value"]
+    input_shape = paddle_op.operands()[0].source().shape
+    keepdim = paddle_op.attrs()["keepdim"]
+    if network.has_implicit_batch_dimension:
+        assert (
+            axis != 0
+        ), "can't reduce on axis == 0 when network has implicit batch dimension"
+    output_shape = []
+    if len(axis) == 0:
+        axis = list(range(len(input_shape)))
+    for i in range(len(axis)):
+        if axis[i] < 0:
+            axis[i] = len(input_shape) + axis[i]
+    layer = network.add_reduce(
+        input_tensor,
+        trt.ReduceOperation.MAX,
+        axes=get_axes_for_reduce_op(axis),
+        keep_dims=keepdim,
+    )
+    return layer.get_output(0)
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -14,4 +14,5 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_conv2d PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_add PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_batch_norm PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_math PROPERTIES TIMEOUT "100")
 endif()
diff --git a/test/tensorrt/test_converter_math.py b/test/tensorrt/test_converter_math.py
@@ -0,0 +1,41 @@
+# 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 unittest
+
+import numpy as np
+from tensorrt_test_base import TensorRTBaseTest
+
+import paddle
+
+
+class TestMaxTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.max
+        self.api_args = {
+            "x": np.array([[0.2, 0.3, 0.5, 0.9], [0.1, 0.2, 0.6, 0.7]]).astype(
+                "float32"
+            ),
+            "axis": [0, 1],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 4], "y": [1, 4]}
+        self.max_shape = {"x": [5, 4], "y": [5, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+if __name__ == '__main__':
+    unittest.main()
