【SCU】【Paddle TensorRT No.49】Add pd_op.pad3d converter (#70556)

PolaKuma · web-flow · commit 6fbf9ceb5b73 · 2025-02-19T14:24:57.000+08:00
* add

* Update trt_op_marker_pass.cc

* Update test_converter_manipulation.py

* fix codestyle

* update

* update

* fix codestyle

* Update manipulation.py

* update

* update

* Update test_converter_manipulation.py

* Update manipulation.py

* Update trt_op_marker_pass.cc

* Update trt_op_marker_pass.cc

* fix

* fix codestyle
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -2333,6 +2333,42 @@ class OneHotOpPattern
   }
 };
 
+class Pad3dOpPattern : public pir::OpRewritePattern<paddle::dialect::Pad3dOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::Pad3dOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::Pad3dOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value paddings_tensor = op.operand_source(1);
+    if (!paddings_tensor) {
+      VLOG(3) << "pad3d should have paddings.";
+      return false;
+    }
+    if (op->HasAttribute("mode")) {
+      auto mode = op->attribute<pir::StrAttribute>("mode").AsString();
+      if (mode != "constant" && mode != "reflect" && mode != "replicate") {
+        VLOG(3) << "The pad3d layer of TRT only support "
+                   "constant/reflect/replicate mode.";
+        return false;
+      }
+    }
+    if (op->HasAttribute("data_format")) {
+      auto data_format =
+          op->attribute<pir::StrAttribute>("data_format").AsString();
+      if (data_format != "NCDHW") {
+        VLOG(3) << "The pad3d layer of TRT only support NCDHW data format.";
+        return false;
+      }
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TemporalShiftOpPattern
     : public pir::OpRewritePattern<paddle::dialect::TemporalShiftOp> {
  public:
@@ -2840,6 +2876,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<TemporalShiftOpPattern>(context));
     ps.Add(std::make_unique<IndexPutOpPattern>(context));
     ps.Add(std::make_unique<InstanceNormOpPattern>(context));
+    ps.Add(std::make_unique<Pad3dOpPattern>(context));
     ps.Add(std::make_unique<EinsumOpPattern>(context));
     ps.Add(std::make_unique<PNormOpPattern>(context));
     ps.Add(std::make_unique<AffineChannelOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import numpy as np
 import tensorrt as trt
 
 from paddle.tensorrt.converter_utils import (
@@ -961,6 +961,81 @@ def roll_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+@converter_registry.register("pd_op.pad3d", trt_version="8.x")
+def pad3d_converter(network, paddle_op, inputs):
+    input_tensor, paddings = inputs
+
+    value = paddle_op.attrs().get("pad_value", 0.0)
+    padding_mode = paddle_op.attrs().get("mode", "constant")
+    input_dim = len(input_tensor.shape)
+    pad_size = paddings.shape[0]
+    assert (
+        input_dim * 2 - 4 == pad_size
+    ), f"Expected paddings size is {input_dim * 2 - 4}, but received {pad_size}."
+
+    shuffle_index = [4, 2, 0, 5, 3, 1]
+    shuffle_inputs = [
+        get_shape_tensor_element(network, paddings, shuffle_index[i])
+        for i in range(pad_size)
+    ]
+    paddings = trt_concat(network, shuffle_inputs)
+
+    pre_zeros = add_1D_constant_layer(network, [0, 0])
+    start_slice1 = [0]
+    start_slice2 = [3]
+    size_slice = [3]
+    stride_slice = [1]
+    pre_pad = network.add_slice(
+        paddings, start_slice1, size_slice, stride_slice
+    ).get_output(0)
+    pre_pad = trt_concat(network, [pre_zeros, pre_pad])
+    post_pad = network.add_slice(
+        paddings, start_slice2, size_slice, stride_slice
+    ).get_output(0)
+    post_pad = trt_concat(network, [pre_zeros, post_pad])
+
+    zeros = add_1D_constant_layer(network, [0] * input_dim)
+
+    start = trt_sub(network, zeros, pre_pad)
+    total_padding = trt_sum(network, pre_pad, post_pad)
+    input_shape = trt_shape(network, input_tensor)
+    size = trt_sum(network, input_shape, total_padding)
+
+    # Add slice layer
+    stride = [1] * input_dim
+    dummy = stride
+    slice_layer = network.add_slice(input_tensor, dummy, dummy, stride)
+    slice_layer.set_input(1, start)
+    slice_layer.set_input(2, size)
+
+    # Set padding mode
+    if padding_mode == "constant":
+        slice_layer.mode = trt.SampleMode.FILL
+        if value != 0.0:
+            if input_tensor.dtype in (
+                trt.DataType.FLOAT,
+                trt.DataType.HALF,
+                trt.DataType.INT8,
+            ):
+                fill_value = add_1D_constant_layer(
+                    network, value, dtype=np.float32
+                )
+            else:
+                value_int = int(value)
+                fill_value = add_1D_constant_layer(
+                    network, value_int, dtype=np.int32
+                )
+            slice_layer.set_input(4, fill_value)
+    elif padding_mode == "reflect":
+        slice_layer.mode = trt.SampleMode.REFLECT
+    elif padding_mode == "replicate":
+        slice_layer.mode = trt.SampleMode.CLAMP
+    else:
+        raise ValueError(f"Unsupported padding mode: {padding_mode}")
+
+    return slice_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.numel", trt_version="8.x")
 def numel_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -743,6 +743,167 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+def wrapper_pad3d(x, paddings, mode, value, data_format):
+    pad3d = paddle.nn.Pad3D(
+        padding=[1, 0, 1, 2, 0, 0],
+        mode=mode,
+        value=value,
+        data_format=data_format,
+    )
+    return pad3d(x)
+
+
+def wrapper_pad3d_error2(x):
+    pad3d = paddle.nn.Pad3D(
+        padding=[1, 0, 1, 2, 0, 0],
+        mode="constant",
+        value=1.0,
+        data_format="NCDHW",
+    )
+    return pad3d(x)
+
+
+class TestPad3dCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = wrapper_pad3d
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 2, 3]).astype("float32"),
+            "paddings": np.array([1, 0, 1, 2, 0, 0], dtype="int32"),
+            "mode": "constant",
+            "value": 1.0,
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x", "paddings"]}
+        self.min_shape = {"x": [1, 1, 1, 2, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 2, 3]}
+        self.max_shape = {"x": [10, 1, 1, 2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPad3dCaseINTTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = wrapper_pad3d
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 2, 3]).astype("int32"),
+            "paddings": np.array([1, 0, 1, 2, 0, 0], dtype="int32"),
+            "mode": "constant",
+            "value": 1.0,
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x", "paddings"]}
+        self.min_shape = {"x": [1, 1, 1, 2, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 2, 3]}
+        self.max_shape = {"x": [10, 1, 1, 2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPad3dOtherformat1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = wrapper_pad3d
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 3, 3]).astype("float32"),
+            "paddings": np.array([1, 0, 1, 2, 0, 0], dtype="int32"),
+            "mode": "reflect",
+            "value": 1.0,
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x", "paddings"]}
+        self.min_shape = {"x": [1, 1, 1, 3, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 3, 3]}
+        self.max_shape = {"x": [10, 1, 1, 3, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPad3dOtherformat2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = wrapper_pad3d
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 2, 3]).astype("float32"),
+            "paddings": np.array([1, 0, 1, 2, 0, 0], dtype="int32"),
+            "mode": "replicate",
+            "value": 1.0,
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x", "paddings"]}
+        self.min_shape = {"x": [1, 1, 1, 2, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 2, 3]}
+        self.max_shape = {"x": [10, 1, 1, 2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPad3dNoPaddingTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = wrapper_pad3d_error2
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 2, 3]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 1, 1, 2, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 2, 3]}
+        self.max_shape = {"x": [10, 1, 1, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestPad3dErrorModeTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle._C_ops.pad3d
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 2, 3]).astype("float32"),
+            "paddings": np.array([1, 0, 1, 2, 0, 0], dtype="int32"),
+            "mode": "error",
+            "value": 1.0,
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x", "paddings"]}
+        self.min_shape = {"x": [1, 1, 1, 2, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 2, 3]}
+        self.max_shape = {"x": [10, 1, 1, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestPad3dErrorDataformatTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle._C_ops.pad3d
+        self.api_args = {
+            "x": np.random.random([1, 1, 1, 2, 3]).astype("float32"),
+            "paddings": np.array([1, 0, 1, 2, 0, 0], dtype="int32"),
+            "mode": "constant",
+            "value": 1.0,
+            "data_format": "error",
+        }
+        self.program_config = {"feed_list": ["x", "paddings"]}
+        self.min_shape = {"x": [1, 1, 1, 2, 3]}
+        self.opt_shape = {"x": [1, 1, 1, 2, 3]}
+        self.max_shape = {"x": [10, 1, 1, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
 class TestNumelTRTCase1Pattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.numel
