[Paddle TensorRT] fix pd_op.pool2d (#69864)

lizexu123 · web-flow · commit 6043c7f38202 · 2024-12-03T19:05:14.000+08:00
* fix pool2d

* fix
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -331,6 +331,60 @@ class Pool2dOpPattern
         }
       }
     }
+
+    auto ceil_mode = op->attribute<pir::BoolAttribute>("ceil_mode").data();
+    auto global_pooling =
+        op->attribute<pir::BoolAttribute>("global_pooling").data();
+    std::string padding_algorithm =
+        op->attribute<pir::StrAttribute>("padding_algorithm").AsString();
+    // TODO(Lizexu): The general plugin approach for entering TensorRT has not
+    // been supported yet.
+    auto adaptive = op->attribute<pir::BoolAttribute>("adaptive").data();
+    if (adaptive) {
+      VLOG(3)
+          << "The adaptive is true pd_op.pool2d is not supported by trt now";
+      return false;
+    }
+    // TODO(Lizexu): This piece of code exists in the old IR-TRT implementation
+    // but is not covered by unit tests, raising suspicions about its
+    // correctness. In the PIR-TRT implementation, following the same approach
+    // causes precision issues. For now, we will exclude it from entering
+    // TensorRT.
+    pir::Value input = op.operand_source(0);
+    auto kernel_size_attr =
+        full_int_array_op->attribute<pir::ArrayAttribute>("value");
+    std::vector<int64_t> kernel_size;
+    for (const auto &attr : kernel_size_attr.AsVector()) {
+      kernel_size.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
+    }
+
+    auto input_type = input.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    auto input_dims = input_type.dims();
+    int g_post_pad_h = 0;
+    int g_post_pad_w = 0;
+    int input_height = input_dims[input_dims.size() - 2];
+    int input_width = input_dims[input_dims.size() - 1];
+    std::vector<int32_t> strides;
+    auto strides_attr = op->attribute<pir::ArrayAttribute>("strides");
+    for (const auto &attr : strides_attr.AsVector()) {
+      strides.push_back(attr.dyn_cast<pir::Int32Attribute>().data());
+    }
+    if (input_height > 0 &&
+        input_height - kernel_size[0] + 2 * paddings[0] < 0) {
+      g_post_pad_h = strides[0] - 1;
+    }
+    if (input_width > 0 && input_width - kernel_size[1] + 2 * paddings[1] < 0) {
+      g_post_pad_w = strides[1] - 1;
+    }
+    if (!adaptive && !global_pooling && !ceil_mode) {
+      if (padding_algorithm != "SAME" &&
+          ((g_post_pad_h > 0 && input_height > 0) ||
+           (g_post_pad_w > 0 && input_width > 0))) {
+        VLOG(3) << "The pool2d op meets the condition that may cause precision "
+                   "issues in TRT. Skip TRT conversion.";
+        return false;
+      }
+    }
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
diff --git a/python/paddle/tensorrt/impls/pooling.py b/python/paddle/tensorrt/impls/pooling.py
@@ -21,68 +21,139 @@
 @converter_registry.register("pd_op.pool2d", trt_version="8.x")
 def pool2d_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]
-    pooling_type = paddle_op.attrs().get("pooling_type", "max")
-    padding = paddle_op.attrs().get("paddings", [0, 0])
-    stride = paddle_op.attrs().get("strides", [1, 1])
-    ceil_mode = paddle_op.attrs().get("ceil_mode", False)
-    exclusive = paddle_op.attrs().get("exclusive")
-    adaptive = paddle_op.attrs().get("adaptive")
-    padding_algorithm = paddle_op.attrs().get("padding_algorithm")
 
-    input_shape = input_tensor.shape
+    input_shape = paddle_op.operands()[0].source().shape
+    input_dims = len(input_shape)
+
+    global_pooling = paddle_op.attrs().get("global_pooling", False)
+    pool_type = paddle_op.attrs().get("pooling_type")
+    strides = paddle_op.attrs().get("strides")
+    paddings = paddle_op.attrs().get("paddings")
+    exclusive = paddle_op.attrs().get("exclusive", True)
+    ceil_mode = paddle_op.attrs().get("ceil_mode", False)
+    adaptive = paddle_op.attrs().get("adaptive", False)
+    padding_algorithm = paddle_op.attrs().get("padding_algorithm", "EXPLICIT")
 
-    # TODO attention for these codes
     if not paddle_op.attrs().get("kernel_size") and len(inputs) == 2:
-        # the size of pool2d inputs is 2, means kernel size is the second input.
-        # kernel_size_tensor = inputs[1]
         full_int_op = paddle_op.operands()[1].source().get_defining_op()
         if full_int_op.name() == "pd_op.full_int_array":
             kernel_size = full_int_op.attrs().get("value")
         else:
             raise Exception(
-                "the defining op of kernel size must be pd_op.full_int_array"
+                "The defining op of kernel size must be pd_op.full_int_array"
             )
     else:
         kernel_size = paddle_op.attrs().get("kernel_size")
 
-    if len(stride) == 0 or stride[0] is None:
-        stride = kernel_size
+    nv_pool_type = trt.PoolingType.MAX
+    reduce_operation = trt.ReduceOperation.MAX
+    if pool_type == "max":
+        nv_pool_type = trt.PoolingType.MAX
+        reduce_operation = trt.ReduceOperation.MAX
+    elif pool_type == "avg":
+        nv_pool_type = trt.PoolingType.AVERAGE
+        reduce_operation = trt.ReduceOperation.AVG
 
-    if pooling_type == "max":
-        pooling_type = trt.PoolingType.MAX
-    elif pooling_type == "avg":
-        pooling_type = trt.PoolingType.AVERAGE
-    else:
-        raise ValueError(f"Unsupported pooling type: {pooling_type}")
+    if global_pooling or adaptive:
+        paddings = [0] * len(paddings)
 
     if padding_algorithm == "VALID":
-        padding = [0, 0]
-
-    if adaptive:
-        output_size = kernel_size
-        stride = tuple(input_shape[-2 + i] // output_size[i] for i in range(2))
-        kernel_size = tuple(
-            input_shape[-2 + i] - (output_size[i] - 1) * stride[i]
-            for i in range(2)
+        paddings = [0] * len(paddings)
+
+    nv_paddings = trt.DimsHW(paddings[0], paddings[1])
+    nv_ksize = trt.DimsHW(kernel_size[0], kernel_size[1])
+    nv_strides = trt.DimsHW(strides[0], strides[1])
+
+    layer = None
+    g_pre_pad = trt.DimsHW(0, 0)
+    g_post_pad = trt.DimsHW(0, 0)
+
+    if (
+        input_shape[input_dims - 2] > 0
+        and input_shape[input_dims - 2] - kernel_size[0] + 2 * paddings[0] < 0
+    ):
+        g_post_pad.h = strides[0] - 1
+    if (
+        input_shape[input_dims - 1] > 0
+        and input_shape[input_dims - 1] - kernel_size[1] + 2 * paddings[1] < 0
+    ):
+        g_post_pad.w = strides[1] - 1
+
+    real_paddings = paddings.copy()
+    for i in range(2):
+        copy_pad = paddings[i]
+        real_paddings.insert(2 * i + 1, copy_pad)
+
+    if padding_algorithm == "SAME":
+        for i in range(2):
+            copy_pad = paddings[2 * i]
+            paddings.insert(2 * i + 1, copy_pad)
+
+        for i in range(2):
+            out_size = (input_shape[2 + i] + strides[i] - 1) // strides[i]
+            pad_sum = max(
+                (out_size - 1) * strides[i]
+                + kernel_size[i]
+                - input_shape[2 + i],
+                0,
+            )
+            pad_0 = pad_sum // 2
+            pad_1 = pad_sum - pad_0
+            paddings[2 * i] = pad_0
+            paddings[2 * i + 1] = pad_1
+        real_paddings = paddings.copy()
+
+    paddings = [paddings[i] for i in range(len(paddings)) if i % 2 == 0]
+
+    if padding_algorithm == "VALID":
+        read_paddings = [0] * len(real_paddings)
+
+    if not adaptive and not global_pooling and not ceil_mode:
+        if padding_algorithm != "SAME" and (
+            (g_post_pad.h > 0 and input_shape[input_dims - 2] > 0)
+            or (g_post_pad.w > 0 and input_shape[input_dims - 1] > 0)
+        ):
+            pad_layer = network.add_padding_nd(
+                input=input_tensor,
+                pre_padding=tuple(g_pre_pad),
+                post_padding=tuple(g_post_pad),
+            )
+            input_tensor = pad_layer.get_output(0)
+        pooling_layer = network.add_pooling_nd(
+            input=input_tensor, type=nv_pool_type, window_size=nv_ksize
         )
+        pooling_layer.stride_nd = nv_strides
+        pooling_layer.padding_nd = nv_paddings
+        pooling_layer.average_count_excludes_padding = exclusive
+        if padding_algorithm == "SAME":
+            pooling_layer.padding_mode = trt.PaddingMode.SAME_UPPER
 
-        pool_layer = network.add_pooling_nd(
-            input_tensor, pooling_type, window_size=kernel_size
+        layer = pooling_layer
+    elif not adaptive and not global_pooling and ceil_mode:
+        pooling_layer = network.add_pooling_nd(
+            input=input_tensor, type=nv_pool_type, window_size=nv_ksize
+        )
+        pooling_layer.stride_nd = nv_strides
+        pooling_layer.padding_nd = nv_paddings
+        pooling_layer.average_count_excludes_padding = exclusive
+        if padding_algorithm == "SAME":
+            pooling_layer.padding_mode = trt.PaddingMode.SAME_UPPER
+        else:
+            pooling_layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP
+        layer = pooling_layer
+    elif global_pooling and not adaptive:
+        reduce_axes = (1 << (input_dims - 2)) | (1 << (input_dims - 1))
+        reduce_layer = network.add_reduce(
+            input=input_tensor,
+            op=reduce_operation,
+            axes=reduce_axes,
+            keep_dims=True,
         )
-        pool_layer.stride_nd = stride
-        if pooling_type == "max":
-            pool_layer.padding_nd = padding
+        layer = reduce_layer
     else:
-        pool_layer = network.add_pooling(
-            input_tensor, pooling_type, window_size=kernel_size
+        raise NotImplementedError(
+            "The combination of attributes is not supported yet."
         )
-        pool_layer.stride = stride
-        pool_layer.padding = padding
-        if exclusive:
-            pool_layer.average_count_excludes_padding = True
-        else:
-            pool_layer.average_count_excludes_padding = False
-        if ceil_mode:
-            pool_layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP
 
-    return pool_layer.get_output(0)
+    output_tensor = layer.get_output(0)
+    return output_tensor
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -28,5 +28,5 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_linalg PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_search PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_logic PROPERTIES TIMEOUT "300")
-
+  set_tests_properties(test_converter_pooling PROPERTIES TIMEOUT "300")
 endif()
diff --git a/test/tensorrt/test_converter_pooling.py b/test/tensorrt/test_converter_pooling.py
