[Paddle TensorRT] add remainder op converter (#68793)

gzy19990617 · web-flow · commit cf99b8f7f5fa · 2024-10-29T10:27:15.000+08:00
* add_remainder

* fix review

* fix review

* fix review

* fix review

* fix review

* fix review

* fix review
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -21,38 +21,20 @@
     add_reduce_layer,
     broadcast,
     get_axes_for_reduce_op,
+    trt_div,
+    trt_floor_div,
+    trt_mul,
+    trt_sub,
 )
 from paddle.tensorrt.register import converter_registry
 
 
 @converter_registry.register("pd_op.add", trt_version="8.x")
 @converter_registry.register("pd_op.add_", trt_version="8.x")
 def add_converter(network, paddle_op, inputs):
-    weight_shape = paddle_op.operands()[1].source().shape
-    input_shape = paddle_op.operands()[0].source().shape
-
-    weight_tensor = inputs[1]
-    input_tensor = inputs[0]
-    if type(inputs[1]) == trt.Weights:
-        weight_tensor = network.add_constant(
-            weight_shape, inputs[1]
-        ).get_output(0)
-    if type(inputs[0]) == trt.Weights:
-        input_tensor = network.add_constant(input_shape, inputs[0]).get_output(
-            0
-        )
-    lhs_val, rhs_val = broadcast(
-        network,
-        input_tensor,
-        weight_tensor,
-        input_tensor.name,
-        weight_tensor.name,
-    )
-
-    out = network.add_elementwise(
-        lhs_val, rhs_val, trt.ElementWiseOperation.SUM
+    return add_elementwise_layer(
+        network, paddle_op, inputs, trt.ElementWiseOperation.SUM
     )
-    return out.get_output(0)
 
 
 @converter_registry.register("pd_op.scale", trt_version="8.x")
@@ -122,6 +104,50 @@ def multiply_converter(network, paddle_op, inputs):
     )
 
 
+@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):
+    weight_shape = paddle_op.operands()[1].source().shape
+    input_shape = paddle_op.operands()[0].source().shape
+
+    weight_tensor = inputs[1]
+    input_tensor = inputs[0]
+    if type(inputs[1]) == trt.Weights:
+        weight_tensor = network.add_constant(
+            weight_shape, inputs[1]
+        ).get_output(0)
+    if type(inputs[0]) == trt.Weights:
+        input_tensor = network.add_constant(input_shape, inputs[0]).get_output(
+            0
+        )
+
+    lhs_val, rhs_val = broadcast(
+        network,
+        input_tensor,
+        weight_tensor,
+        input_tensor.name,
+        weight_tensor.name,
+    )
+
+    # Check if floor division is needed
+    is_floor_div = input_tensor.dtype != trt.DataType.INT32
+
+    # Floor division
+    quotient = (
+        trt_floor_div(network, lhs_val, rhs_val)
+        if is_floor_div
+        else trt_div(network, lhs_val, rhs_val)
+    )
+
+    # Multiply rhs by the quotient
+    product = trt_mul(network, rhs_val, quotient)
+
+    # Subtract the product from lhs to get the remainder
+    remainder = trt_sub(network, lhs_val, product)
+
+    return remainder
+
+
 @converter_registry.register("pd_op.min", trt_version="8.x")
 def min_converter(network, paddle_op, inputs):
     return add_reduce_layer(network, paddle_op, inputs, trt.ReduceOperation.MIN)
diff --git a/test/tensorrt/tensorrt_test_base.py b/test/tensorrt/tensorrt_test_base.py
@@ -36,6 +36,7 @@ def __init__(self, methodName='runTest'):
         self.min_shape = None
         self.max_shape = None
         self.target_marker_op = ""
+        self.dynamic_shape_data = {}
 
     def create_fake_program(self):
         if self.python_api is None:
@@ -210,12 +211,20 @@ def check_trt_result(self, rtol=1e-5, atol=1e-5):
                         max_shape_data[feed_name] = self.api_args[feed_name]
                         continue
                     else:
-                        min_shape_data[feed_name] = np.random.randn(
-                            *self.min_shape[feed_name]
-                        ).astype(self.api_args[feed_name].dtype)
-                        max_shape_data[feed_name] = np.random.randn(
-                            *self.max_shape[feed_name]
-                        ).astype(self.api_args[feed_name].dtype)
+                        if self.dynamic_shape_data:
+                            min_shape_data[feed_name] = self.dynamic_shape_data[
+                                feed_name
+                            ](self.min_shape[feed_name])
+                            max_shape_data[feed_name] = self.dynamic_shape_data[
+                                feed_name
+                            ](self.max_shape[feed_name])
+                        else:
+                            min_shape_data[feed_name] = np.random.randn(
+                                *self.min_shape[feed_name]
+                            ).astype(self.api_args[feed_name].dtype)
+                            max_shape_data[feed_name] = np.random.randn(
+                                *self.max_shape[feed_name]
+                            ).astype(self.api_args[feed_name].dtype)
 
             scope = paddle.static.global_scope()
             main_program = warmup_shape_infer(
diff --git a/test/tensorrt/test_converter_math.py b/test/tensorrt/test_converter_math.py
@@ -95,6 +95,54 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestRemainderFloatTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.remainder
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype(np.float32),
+            "y": np.random.uniform(low=0.1, high=1, size=(2, 3)).astype(
+                np.float32
+            ),  # Ensure y is non-zero
+        }
+        self.dynamic_shape_data = {
+            "x": lambda shape: np.random.randn(*shape).astype(np.float32),
+            "y": lambda shape: np.random.uniform(
+                low=0.1, high=1, size=shape
+            ).astype(np.float32),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [1, 3]}
+        self.max_shape = {"x": [5, 3], "y": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestRemainderIntTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.remainder
+        self.api_args = {
+            "x": np.random.randint(1, 10, size=(2, 3)).astype(np.int32),
+            "y": np.random.randint(1, 10, size=(2, 3)).astype(
+                np.int32
+            ),  # Ensure y is non-zero
+        }
+        self.dynamic_shape_data = {
+            "x": lambda shape: np.random.randint(1, 10, size=shape).astype(
+                np.int32
+            ),
+            "y": lambda shape: np.random.randint(1, 10, size=shape).astype(
+                np.int32
+            ),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [1, 3]}
+        self.max_shape = {"x": [5, 3], "y": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestMinTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.min
