[Paddle TensorRT] add pd_op.clip converter (#68778)

gzy19990617 · web-flow · commit 6bb1bcfc2dde · 2024-11-15T14:06:07.000+08:00
* add clip op converter

* add clip op converter

* "fix code style"

* fix review

* fix review

* fix review

* fix

* fix

* merge

* fix

* fix

* Update trt_op_marker_pass.cc

* Update trt_op_marker_pass.cc
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1497,6 +1497,32 @@ class NearestInterV2Pattern
   }
 };
 
+class ClipPattern : public pir::OpRewritePattern<paddle::dialect::ClipOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::ClipOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::ClipOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    pir::Value x = op.operand_source(0);
+    auto x_shape = pir::GetShapeFromValue(x);
+    if (x_shape.size() == 0) {
+      VLOG(3) << " clip op does not support input's dim is 0 in tensorrt.";
+      return false;
+    }
+    auto min_tensor = op.operand_source(1);
+    if (!min_tensor) {
+      VLOG(3) << "clip op does not have input min tensor";
+      return false;
+    }
+    auto max_tensor = op.operand_source(2);
+    if (!max_tensor) {
+      VLOG(3) << "clip op does not have input max tensor";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class GridSampleOpPattern
     : public pir::OpRewritePattern<paddle::dialect::GridSampleOp> {
  public:
@@ -1528,7 +1554,6 @@ class GridSampleOpPattern
       return false;
     }
 #endif
-
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
@@ -1899,6 +1924,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<SumOpPattern>(context));
     ps.Add(std::make_unique<BilinearInterpV2Pattern>(context));
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
+    ps.Add(std::make_unique<ClipPattern>(context));
     ps.Add(std::make_unique<GridSampleOpPattern>(context));
     ps.Add(std::make_unique<StackOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -170,6 +170,53 @@ def add_1D_constant_layer(network, data, dtype=np.int32, is_scalar=False):
     return constant_layer.get_output(0)
 
 
+# Create an constant layer with shape_tensor and value
+def fill_constant_layer(network, shape_tensor, tensor_rank, data, trt_dtype):
+    fill_layer = network.add_fill(
+        trt.Dims([tensor_rank]), trt.FillOperation.LINSPACE
+    )
+    np_dtype = map_trt_dtype(trt_dtype)
+    fill_layer.set_input(0, shape_tensor)
+    fill_layer.set_input(
+        1, add_1D_constant_layer(network, data, np_dtype, is_scalar=True)
+    )
+    beta = [0] * tensor_rank
+    fill_layer.set_input(
+        2, add_1D_constant_layer(network, beta, np_dtype, is_scalar=False)
+    )
+    return fill_layer.get_output(0)
+
+
+def trt_expand(network, input, rank, shape_tensor, shape_rank):
+    if rank < shape_rank:
+        one_rank_tensor = add_1D_constant_layer(
+            network, [1] * (shape_rank - rank)
+        )
+        in_shape_tensor = trt_shape(network, input)
+        itensors = [one_rank_tensor, in_shape_tensor]
+        input_shape_tensor = trt_concat(network, itensors)
+    else:
+        input_shape_tensor = trt_shape(network, input)
+
+    new_input_tensor = trt_reshape(network, input, input_shape_tensor, "", True)
+
+    start = [0] * shape_rank
+    starts_tensor = add_1D_constant_layer(network, start)
+    one_tensor = add_1D_constant_layer(network, 1)
+    sizes_tensor = trt_max(network, input_shape_tensor, shape_tensor)
+    input_sub_tensor = trt_sub(network, input_shape_tensor, one_tensor)
+    strides_tensor = trt_min(network, one_tensor, input_sub_tensor)
+
+    slice_layer = network.add_slice(
+        new_input_tensor, start, [0] * len(start), [0] * len(start)
+    )
+    slice_layer.set_input(1, starts_tensor)
+    slice_layer.set_input(2, sizes_tensor)
+    slice_layer.set_input(3, strides_tensor)
+
+    return slice_layer.get_output(0)
+
+
 # Concat not make rank changed
 def trt_concat(network, inputs, axis=0):
     concat_layer = network.add_concatenation(inputs=inputs)
@@ -324,6 +371,21 @@ def build_size_tensor(
     return size_tensor
 
 
+# convert trt_dtype to numpy dtype
+def map_trt_dtype(trt_dtype):
+    dtype_map = {
+        trt.DataType.FLOAT: np.float32,
+        trt.DataType.HALF: np.float16,
+        trt.DataType.INT32: np.int32,
+        trt.DataType.INT8: np.int8,
+        trt.DataType.BOOL: np.bool,
+    }
+    if trt_dtype in dtype_map:
+        return dtype_map[trt_dtype]
+    else:
+        raise TypeError(f"Unsupported trt_dtype: {trt_dtype}")
+
+
 # Reduce the given tensor in the TensorRT network to a scalar
 def trt_reduce_to_scalar(network, tensor):
     if len(tensor.shape) == 0:
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -27,6 +27,7 @@
     get_shape_tensor_element,
     has_dynamic_shape,
     trt_concat,
+    trt_expand,
     trt_floor_div,
     trt_less,
     trt_max,
@@ -226,36 +227,6 @@ def squeeze_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
-def get_expand_output(network, input, rank, shape_tensor, shape_rank):
-    if rank < shape_rank:
-        one_rank_tensor = add_1D_constant_layer(
-            network, [1] * (shape_rank - rank)
-        )
-        in_shape_tensor = trt_shape(network, input)
-        itensors = [one_rank_tensor, in_shape_tensor]
-        input_shape_tensor = trt_concat(network, itensors)
-    else:
-        input_shape_tensor = trt_shape(network, input)
-
-    new_input_tensor = trt_reshape(network, input, input_shape_tensor, "", True)
-
-    start = [0] * shape_rank
-    starts_tensor = add_1D_constant_layer(network, start)
-    one_tensor = add_1D_constant_layer(network, 1)
-    sizes_tensor = trt_max(network, input_shape_tensor, shape_tensor)
-    input_sub_tensor = trt_sub(network, input_shape_tensor, one_tensor)
-    strides_tensor = trt_min(network, one_tensor, input_sub_tensor)
-
-    slice_layer = network.add_slice(
-        new_input_tensor, start, [0] * len(start), [0] * len(start)
-    )
-    slice_layer.set_input(1, starts_tensor)
-    slice_layer.set_input(2, sizes_tensor)
-    slice_layer.set_input(3, strides_tensor)
-
-    return slice_layer.get_output(0)
-
-
 @converter_registry.register("pd_op.expand", trt_version="8.x")
 def expand_converter(network, paddle_op, inputs):
     input = inputs[0]
@@ -275,7 +246,7 @@ def expand_converter(network, paddle_op, inputs):
     else:
         shape_tensor = inputs[1]
         shape_rank = shape_tensor.shape[0]
-    return get_expand_output(network, input, rank, shape_tensor, shape_rank)
+    return trt_expand(network, input, rank, shape_tensor, shape_rank)
 
 
 @converter_registry.register("pd_op.expand_as", trt_version="8.x")
@@ -293,7 +264,7 @@ def expand_as_converter(network, paddle_op, inputs):
         shape = paddle_op.attrs().get("target_shape")
         shape_tensor = add_1D_constant_layer(network, shape)
         shape_rank = len(shape)
-    return get_expand_output(network, input, rank, shape_tensor, shape_rank)
+    return trt_expand(network, input, rank, shape_tensor, shape_rank)
 
 
 @converter_registry.register("pd_op.cast", trt_version="8.x")
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -20,10 +20,14 @@
     add_elementwise_layer,
     add_reduce_layer,
     broadcast,
+    cast_tensor,
+    fill_constant_layer,
     get_axes_for_reduce_op,
     trt_cast,
     trt_div,
+    trt_expand,
     trt_floor_div,
+    trt_max,
     trt_mul,
     trt_sub,
 )
@@ -105,6 +109,51 @@ def multiply_converter(network, paddle_op, inputs):
     )
 
 
+@converter_registry.register("pd_op.clip", trt_version="8.x")
+def clip_converter(network, paddle_op, inputs):
+    def _get_constant_or_expand_tensor(
+        op, constant_inputs, input_shape_tensor, rank
+    ):
+        if op.name() == "pd_op.full":
+            value = op.attrs()["value"]
+            return fill_constant_layer(
+                network, input_shape_tensor, rank, value, input_tensor.dtype
+            )
+        else:
+            expanded_tensor = trt_expand(
+                network, constant_inputs, 1, input_shape_tensor, rank
+            )
+            if expanded_tensor.dtype != input_tensor.dtype:
+                expanded_tensor = cast_tensor(
+                    network, expanded_tensor, input_tensor.dtype
+                )
+            return expanded_tensor
+
+    input_tensor = inputs[0]
+    input_shape = paddle_op.operands()[0].source().shape
+    rank = len(input_shape)
+    input_shape_tensor = network.add_shape(input_tensor).get_output(0)
+
+    # handle min operation
+    min_op = paddle_op.operands()[1].source().get_defining_op()
+    alpha_t = _get_constant_or_expand_tensor(
+        min_op, inputs[1], input_shape_tensor, rank
+    )
+
+    # handle max operation
+    max_op = paddle_op.operands()[2].source().get_defining_op()
+    beta_t = _get_constant_or_expand_tensor(
+        max_op, inputs[2], input_shape_tensor, rank
+    )
+
+    # run the clip operation
+    lower_clip = trt_max(network, input_tensor, alpha_t)
+    layer = network.add_elementwise(
+        lower_clip, beta_t, trt.ElementWiseOperation.MIN
+    )
+    return layer.get_output(0)
+
+
 @converter_registry.register("pd_op.remainder", trt_version="8.x")
 @converter_registry.register("pd_op.remainder_", trt_version="8.x")
 def remainder_converter(network, paddle_op, inputs):
diff --git a/python/paddle/tensorrt/register.py b/python/paddle/tensorrt/register.py
@@ -32,6 +32,10 @@ def get(self, op_name, trt_version=None):
         for version_range, func in self._registry[op_name]:
             if self._version_match(trt_version, version_range):
                 return func
+            else:
+                raise ValueError(
+                    f"Requested TensorRT version : {trt_version} does not match the range of pip installed tensorrt versions : {version_range}"
+                )
         return self._registry.get(op_name)
 
     def _version_match(self, trt_version, version_range):
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -23,6 +23,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_clip PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_vision PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_linalg PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_search PROPERTIES TIMEOUT "300")
diff --git a/test/tensorrt/test_converter_clip.py b/test/tensorrt/test_converter_clip.py
@@ -0,0 +1,99 @@
+# 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 TestClipTRTPatternCase1(TensorRTBaseTest):
+    '''min/max is attr, and x/min/max is float'''
+
+    def setUp(self):
+        self.python_api = paddle.clip
+        self.api_args = {
+            "x": np.array([[2, 0.3, 0.5, 0.9], [0.1, 0.2, 6, 7]]).astype(
+                "float32"
+            ),
+            "min": 2.2,
+            "max": 5.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 4]}
+        self.max_shape = {"x": [5, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestClipTRTPatternCase2(TensorRTBaseTest):
+    def setUp(self):
+        '''min/max is attr, and x is int, min/max is float'''
+        self.python_api = paddle.clip
+        self.api_args = {
+            "x": np.array([[2, 3, 5, 9], [1, 2, 6, 7]]).astype("int64"),
+            "min": 2.2,
+            "max": 5.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 4]}
+        self.max_shape = {"x": [5, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestClipTRTPatternCase3(TensorRTBaseTest):
+    '''min/max is input, and x/min/max is float'''
+
+    def setUp(self):
+        self.python_api = paddle.clip
+        self.api_args = {
+            "x": np.array([[2, 0.3, 0.5, 0.9], [0.1, 0.2, 6, 7]]).astype(
+                "float32"
+            ),
+            "min": np.array([2.2]).astype("float32"),
+            "max": np.array([5.2]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x", "min", "max"]}
+        self.min_shape = {"x": [1, 4]}
+        self.max_shape = {"x": [5, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestClipTRTPatternCase4(TensorRTBaseTest):
+    '''min/max is input, and x is int, min/max is float'''
+
+    def setUp(self):
+        self.python_api = paddle.clip
+        self.api_args = {
+            "x": np.array([[2, 3, 5, 9], [1, 2, 6, 7]]).astype("int64"),
+            "min": np.array([2]).astype("float32"),
+            "max": np.array([5]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x", "min", "max"]}
+        self.min_shape = {"x": [1, 4]}
+        self.max_shape = {"x": [5, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+if __name__ == "__main__":
+    unittest.main()
