From 8a9d4003aa86af35486a02536fd48189fd65bb43 Mon Sep 17 00:00:00 2001
From: Wangzheee <634486483@qq.com>
Date: Tue, 11 Oct 2022 20:39:51 +0800
Subject: [PATCH] [Paddle Inference]add_lookup_table_v2 (#46893)

* add_lookup_table_v2
---
 paddle/fluid/inference/tensorrt/convert/op_converter.h | 7 +++++++
 paddle/fluid/inference/tensorrt/op_teller.cc           | 6 ++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/op_converter.h b/paddle/fluid/inference/tensorrt/convert/op_converter.h
index da500014b8e..1605490ac45 100644
--- a/paddle/fluid/inference/tensorrt/convert/op_converter.h
+++ b/paddle/fluid/inference/tensorrt/convert/op_converter.h
@@ -151,6 +151,13 @@ class OpConverter {
               platform::errors::Unimplemented("no OpConverter for optype [%s]",
                                               op_desc.Type()));
         }
+        if (op_desc.Type() == "lookup_table_v2") {
+          it = Registry<OpConverter>::Global().Lookup("lookup_table");
+          PADDLE_ENFORCE_NOT_NULL(
+              it,
+              platform::errors::Unimplemented("no OpConverter for optype [%s]",
+                                              op_desc.Type()));
+        }
         if (!it) {
           it = Registry<OpConverter>::Global().Lookup(op_desc.Type());
         }
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
index d109fb051ec..904768d179d 100644
--- a/paddle/fluid/inference/tensorrt/op_teller.cc
+++ b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -2221,7 +2221,8 @@ struct SimpleOpTypeSetTeller : public Teller {
       "squeeze2",
       "unsqueeze2",
       "layernorm_shift_partition",
-      "lookup_table"};
+      "lookup_table",
+      "lookup_table_v2"};
   std::unordered_set<std::string> teller_set{
       "mul",
       "matmul",
@@ -2333,7 +2334,8 @@ struct SimpleOpTypeSetTeller : public Teller {
       "unsqueeze2",
       "fused_token_prune",
       "layernorm_shift_partition",
-      "lookup_table"};
+      "lookup_table",
+      "lookup_table_v2"};
 };
 
 struct GenericPluginTeller : public Teller {
-- 
GitLab