[Paddle TensorRT No.42、45] Add (pd_op.isnan、pd_op.embedding) converter

paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc

@@ -63,6 +63,7 @@ DEFINE_GENERAL_PATTERN(Fused_gemm_epilogue,
                        paddle::dialect::FusedGemmEpilogueOp)
 DEFINE_GENERAL_PATTERN(Layer_norm, paddle::dialect::LayerNormOp)
 DEFINE_GENERAL_PATTERN(Add, paddle::dialect::AddOp)
+DEFINE_GENERAL_PATTERN(Isnan, paddle::dialect::IsnanOp)
 DEFINE_GENERAL_PATTERN(Full, paddle::dialect::FullOp)
 DEFINE_GENERAL_PATTERN(Silu, paddle::dialect::SiluOp)
 DEFINE_GENERAL_PATTERN(Conv2d, paddle::dialect::Conv2dOp)
@@ -1409,6 +1410,29 @@ class ArgsortOpPattern
     return true;
   }
 };
+
+class EmbeddingOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::EmbeddingOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::EmbeddingOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::EmbeddingOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    // if (pir::GetDefiningOpForInput(op, 1)->name() == "builtin.parameter") {
+    //   // trt.Weights don't have the shape info.
+    //   VLOG(3) << "Skip to convert into TRT while found weight is a parameter
+    //   "
+    //              "in pd_op.embedding.";
+    //   return false;
+    // }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class BilinearInterpV2Pattern
     : public pir::OpRewritePattern<paddle::dialect::BilinearInterpOp> {
  public:
@@ -2192,6 +2216,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(AssignOut)
     ADD_PATTERN(Assign)
     ADD_PATTERN(Tile)
+    ADD_PATTERN(Isnan)
     ADD_PATTERN(Share_Data)
     ADD_PATTERN(Swish)
     ADD_PATTERN(Log)
@@ -2230,6 +2255,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<FlattenOpPattern>(context));
     ps.Add(std::make_unique<CastOpPattern>(context));
     ps.Add(std::make_unique<SplitOpPattern>(context));
+    ps.Add(std::make_unique<EmbeddingOpPattern>(context));
     ps.Add(std::make_unique<SplitWithNumOpPattern>(context));
     ps.Add(std::make_unique<GreaterEqualOpPattern>(context));
     ps.Add(std::make_unique<GreaterThanOpPattern>(context));

python/paddle/tensorrt/impls/input.py

@@ -60,3 +60,10 @@ def one_hot_converter(network, paddle_op, inputs):
     output_tensor = one_hot_layer.get_output(0)
 
     return [output_tensor]
+
+
+@converter_registry.register("pd_op.embedding", trt_version="8.x")
+def embedding_converter(network, paddle_op, inputs):
+    x_tensor, weight = inputs
+    layer = network.add_gather(weight, x_tensor, 0)
+    return layer.get_output(0)

python/paddle/tensorrt/impls/math.py

@@ -25,6 +25,7 @@
     fill_constant_layer,
     get_axes_for_reduce_op,
     trt_cast,
+    trt_equal,
     trt_expand,
     trt_max,
 )
@@ -305,3 +306,13 @@ def maximum_converter(network, paddle_op, inputs):
         network, paddle_op, inputs, trt.ElementWiseOperation.MAX
     )
     return max_layer
+
+
+@converter_registry.register("pd_op.isnan", trt_version="8.x")
+def isnan_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    equal_tensor = trt_equal(network, input_tensor, input_tensor)
+    layer = network.add_unary(equal_tensor, trt.UnaryOperation.NOT)
+    cast_layer = network.add_identity(layer.get_output(0))
+    cast_layer.set_output_type(0, trt.bool)
+    return cast_layer.get_output(0)

test/tensorrt/test_converter_input.py

@@ -64,5 +64,49 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+def embedding_warpper_func(x):
+    layer = paddle.nn.Embedding(64, 4)
+    return layer(x)
+
+
+class TestEmbeddingCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.embedding
+        self.api_args = {
+            "x": np.array([[3, 16, 24], [6, 4, 47]]).astype("int64"),
+            "weight": np.random.uniform(-1, 1, [64, 4]).astype('float32'),
+        }
+        self.program_config = {"feed_list": ["x", "weight"]}
+        self.dynamic_shape_data = {
+            "x": lambda shape: np.random.randint(1, 64, size=shape).astype(
+                np.int64
+            ),
+        }
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestEmbeddingCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = embedding_warpper_func
+        self.api_args = {
+            "x": np.array([[3, 16, 24], [6, 4, 47]]).astype("int64"),
+        }
+        self.dynamic_shape_data = {
+            "x": lambda shape: np.random.randint(1, 64, size=shape).astype(
+                np.int64
+            ),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()

test/tensorrt/test_converter_math.py

@@ -450,5 +450,33 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestIsnanFloatTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.isnan
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestIsnanIntTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.isnan
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()