【Paddle TensorRT No.22 No.39】add einsum and p_norm OP for pir-trt (#70816)

zhink · web-flow · commit 87343b483eb2 · 2025-01-16T11:00:37.000+08:00
* int8 cuda core must transpose y

* [PIR-TRT] add einsum for pir-trt

* [PIR-TRT] add einsum for pir-trt

* add p_norm for pir-trt
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -2235,6 +2235,68 @@ class InstanceNormOpPattern
   }
 };
 
+class EinsumOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::EinsumOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::EinsumOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::EinsumOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    std::string equation =
+        op->attribute<pir::StrAttribute>("equation").AsString();
+    if (equation.empty()) {
+      VLOG(3) << "Einsum equation is empty";
+      return false;
+    }
+
+    auto operands = op->operands();
+    if (operands.size() > 2) {
+      VLOG(3) << "TensorRT currently supports up to 2 input tensors"
+              << "to einsum but operation had" << operands.size()
+              << "input tensors !";
+      return false;
+    }
+
+    if (equation.find("...") != std::string::npos) {
+      VLOG(3) << "TensorRT currently does not support ellipses !";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
+class PNormOpPattern : public pir::OpRewritePattern<paddle::dialect::PNormOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::PNormOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::PNormOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!op->HasAttribute("asvector") || !op->HasAttribute("axis") ||
+        !op->HasAttribute("porder") || !op->HasAttribute("keepdim")) {
+      VLOG(3) << "p_norm op needs attributes: asvector, porder, axis, keepdim.";
+      return false;
+    }
+    bool asvector = op->attribute<pir::BoolAttribute>("asvector").data();
+    int axis = op->attribute<pir::Int32Attribute>("axis").data();
+    float porder = op->attribute<pir::FloatAttribute>("porder").data();
+
+    if (asvector || porder != 2.0f || axis != -1) {
+      VLOG(3) << "p_norm op only supports asvector=False, porder=2, axis=-1.";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class AffineChannelOpPattern
     : public pir::OpRewritePattern<paddle::dialect::AffineChannelOp> {
  public:
@@ -2411,6 +2473,8 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<TemporalShiftOpPattern>(context));
     ps.Add(std::make_unique<InstanceNormOpPattern>(context));
+    ps.Add(std::make_unique<EinsumOpPattern>(context));
+    ps.Add(std::make_unique<PNormOpPattern>(context));
     ps.Add(std::make_unique<AffineChannelOpPattern>(context));
     return ps;
   }
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -31,6 +31,7 @@
 from .impls.common import *  # noqa: F403
 from .impls.conv import *  # noqa: F403
 from .impls.creation import *  # noqa: F403
+from .impls.einsum import *  # noqa: F403
 from .impls.input import *  # noqa: F403
 from .impls.linalg import *  # noqa: F403
 from .impls.logic import *  # noqa: F403
diff --git a/python/paddle/tensorrt/impls/einsum.py b/python/paddle/tensorrt/impls/einsum.py
@@ -0,0 +1,25 @@
+# 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.
+
+
+from paddle.tensorrt.register import converter_registry
+
+
+@converter_registry.register("pd_op.einsum", trt_version="8.x")
+def convert_einsum(network, paddle_op, inputs):
+    equation = paddle_op.attrs().get("equation", "")
+
+    layer = network.add_einsum(inputs[0], equation)
+    output_tensor = layer.get_output(0)
+    return output_tensor
diff --git a/python/paddle/tensorrt/impls/linalg.py b/python/paddle/tensorrt/impls/linalg.py
@@ -116,3 +116,29 @@ def get_axis_length(axis_idx):
 
     identity_layer = network.add_identity(input_tensor)
     return identity_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.p_norm", trt_version="8.x")
+def p_norm_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    input_dims = input_tensor.shape
+
+    axis = paddle_op.attrs().get("axis", -1)
+    keepdim = paddle_op.attrs().get("keepdim", False)
+    axis = axis if axis >= 0 else axis + len(input_dims)
+    axis_mask = 1 << axis
+
+    prod_layer = network.add_elementwise(
+        input_tensor, input_tensor, trt.ElementWiseOperation.PROD
+    )
+    prod_tensor = prod_layer.get_output(0)
+
+    reduce_layer = network.add_reduce(
+        prod_tensor, trt.ReduceOperation.SUM, axis_mask, keepdim
+    )
+    reduced_tensor = reduce_layer.get_output(0)
+
+    sqrt_layer = network.add_unary(reduced_tensor, trt.UnaryOperation.SQRT)
+    output_tensor = sqrt_layer.get_output(0)
+
+    return output_tensor
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -26,6 +26,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_input PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_vision PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_linalg PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_einsum PROPERTIES TIMEOUT "200")
   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")
diff --git a/test/tensorrt/test_converter_einsum.py b/test/tensorrt/test_converter_einsum.py
@@ -0,0 +1,88 @@
+# 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
+
+
+def einsum_wrapper(equation, x):
+    if not isinstance(x, list):
+        x = [x]
+    out = paddle.einsum(equation, *x)
+    return out[0]
+
+
+class TestEinsumCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = einsum_wrapper
+        self.api_args = {
+            "equation": "ijk->ij",
+            "x": np.random.randn(2, 3, 4).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+
+        self.min_shape = {"x": [1, 3, 4]}
+        self.opt_shape = {"x": [2, 3, 4]}
+        self.max_shape = {"x": [4, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestEinsumCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = einsum_wrapper
+        self.api_args = {
+            "equation": "abcd,bcd->a",
+            "operands": [
+                np.random.randn(2, 3, 4, 5).astype("float32"),
+                np.random.randn(3, 4, 5).astype("float32"),
+            ],
+        }
+        self.program_config = {"feed_list": ["operands"]}
+
+        self.min_shape = {"operands_0": [1, 2, 3, 4], "operands_1": [2, 3, 4]}
+        self.opt_shape = {"operands_0": [2, 3, 4, 5], "operands_1": [3, 4, 5]}
+        self.max_shape = {"operands_0": [4, 6, 8, 9], "operands_1": [6, 8, 9]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestEinsumCaseTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = einsum_wrapper
+        self.api_args = {
+            "equation": "mij,jk->ki",
+            "operands": [
+                np.random.randn(2, 3, 4).astype("float16"),
+                np.random.randn(4, 3).astype("float16"),
+            ],
+        }
+        self.program_config = {"feed_list": ["operands"]}
+
+        self.min_shape = {"operands_0": [1, 3, 4], "operands_1": [1, 3]}
+        self.opt_shape = {"operands_0": [2, 3, 4], "operands_1": [4, 3]}
+        self.max_shape = {"operands_0": [4, 3, 4], "operands_1": [6, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/tensorrt/test_converter_linalg.py b/test/tensorrt/test_converter_linalg.py
@@ -134,5 +134,41 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestPNormTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.linalg.norm
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4).astype("float32"),
+            "p": 2,
+            "axis": -1,
+            "keepdim": False,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4]}
+        self.opt_shape = {"x": [2, 3, 4]}
+        self.max_shape = {"x": [4, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestPNormCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.linalg.norm
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float16"),
+            "p": 2,
+            "axis": -1,
+            "keepdim": False,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [4, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
 if __name__ == '__main__':
     unittest.main()
