PaddlePaddle
diff --git a/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+99-2Lines changed: 99 additions & 2 deletions b/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+99-2Lines changed: 99 additions & 2 deletions
diff --git a/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions b/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions
diff --git a/‎python/paddle/tensorrt/impls/creation.py
Copy file name to clipboard
+22-3Lines changed: 22 additions & 3 deletions b/‎python/paddle/tensorrt/impls/creation.py
Copy file name to clipboard
+22-3Lines changed: 22 additions & 3 deletions
diff --git a/‎python/paddle/tensorrt/impls/input.py
Copy file name to clipboard
+62Lines changed: 62 additions & 0 deletions b/‎python/paddle/tensorrt/impls/input.py
Copy file name to clipboard
+62Lines changed: 62 additions & 0 deletions
diff --git a/‎python/paddle/tensorrt/impls/search.py
Copy file name to clipboard
+7-2Lines changed: 7 additions & 2 deletions b/‎python/paddle/tensorrt/impls/search.py
Copy file name to clipboard
+7-2Lines changed: 7 additions & 2 deletions
diff --git a/‎python/paddle/tensorrt/register.py
Copy file name to clipboard
+46-15Lines changed: 46 additions & 15 deletions b/‎python/paddle/tensorrt/register.py
Copy file name to clipboard
+46-15Lines changed: 46 additions & 15 deletions
diff --git a/‎test/tensorrt/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎test/tensorrt/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎test/tensorrt/test_converter_attribute.py
Copy file name to clipboard
+15-1Lines changed: 15 additions & 1 deletion b/‎test/tensorrt/test_converter_attribute.py
Copy file name to clipboard
+15-1Lines changed: 15 additions & 1 deletion
@@ -78,7 +78,6 @@ DEFINE_GENERAL_PATTERN(Sqrt, paddle::dialect::SqrtOp)
 DEFINE_GENERAL_PATTERN(Hardsigmoid, paddle::dialect::HardsigmoidOp)
 DEFINE_GENERAL_PATTERN(Hardswish, paddle::dialect::HardswishOp)
 DEFINE_GENERAL_PATTERN(Assign, paddle::dialect::AssignOp)
-DEFINE_GENERAL_PATTERN(AssignValue_, paddle::dialect::AssignValue_Op)
 DEFINE_GENERAL_PATTERN(Tile, paddle::dialect::TileOp)
 DEFINE_GENERAL_PATTERN(Share_Data, paddle::dialect::ShareDataOp)
 DEFINE_GENERAL_PATTERN(AssignOut, paddle::dialect::AssignOut_Op)
@@ -1819,6 +1818,11 @@ class TopkOpPattern : public pir::OpRewritePattern<paddle::dialect::TopkOp> {
       VLOG(3) << "pd_op.topk must has axis attribute";
       return false;
     }
+    if (!pir::GetDefiningOpForInput(op, 1)->isa<paddle::dialect::FullOp>()) {
+      VLOG(3) << "The 'k' input of pd_op.topk must be an integer";
+      return false;
+    }
+
     if (op->HasAttribute("sorted")) {
       bool sorted = op->attribute<pir::BoolAttribute>("sorted").data();
       if (!sorted) {
@@ -1987,6 +1991,97 @@ class SetValueWithTensor_OpPattern
   }
 };
 
+class OneHotOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::OneHotOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::OneHotOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::OneHotOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8510)
+    VLOG(3) << "pd_op.one_hot is not supported when TensorRT<8.5.1";
+    return false;
+    pir::Value input = op.operand_source(0);
+    auto input_type = pir::GetDataTypeFromValue(input);
+    if (!input_type.isa<pir::Float32Type>() ||
+        !input_type.isa<pir::Int32Type>() ||
+        !input_type.isa<pir::Int64Type>()) {
+      VLOG(3) << "pd_op.one_hot only support int32,int64,float.";
+      return false;
+    }
+#endif
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
+bool CheckStaticShape(const pir::Operation *op) {
+  std::vector<int32_t> vec_shape;
+  auto shape_attr = op->attribute("shape").dyn_cast<pir::ArrayAttribute>();
+  for (const auto &attr : shape_attr.AsVector()) {
+    vec_shape.push_back(attr.dyn_cast<pir::Int32Attribute>().data());
+  }
+  for (int32_t dim : vec_shape) {
+    if (dim == -1) {
+      VLOG(3) << "pd_op.assign_value_ or pd_op.assign_value cannot support "
+                 "dynamic shape";
+      return false;
+    }
+  }
+  int shape_size = vec_shape.size();
+  int values_count =
+      op->attribute("values").dyn_cast<pir::ArrayAttribute>().size();
+  if (shape_size != values_count) {
+    VLOG(3) << "pd_op.assign_value or pd_op.assign_value shape size is not "
+               "equal to the values size";
+    return false;
+  }
+  return true;
+}
+
+class AssignValue_OpPattern
+    : public pir::OpRewritePattern<paddle::dialect::AssignValue_Op> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::AssignValue_Op>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::AssignValue_Op op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!CheckStaticShape(op)) {
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
+class AssignValueOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::AssignValueOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::AssignValueOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::AssignValueOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!CheckStaticShape(op)) {
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -2024,7 +2119,6 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Hardswish)
     ADD_PATTERN(AssignOut)
     ADD_PATTERN(Assign)
-    ADD_PATTERN(AssignValue_)
     ADD_PATTERN(Tile)
     ADD_PATTERN(Share_Data)
     ADD_PATTERN(Swish)
@@ -2095,6 +2189,9 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<SoftplusOpPatten>(context));
     ps.Add(std::make_unique<EqualOpPattern>(context));
     ps.Add(std::make_unique<NotEqualOpPattern>(context));
+    ps.Add(std::make_unique<OneHotOpPattern>(context));
+    ps.Add(std::make_unique<AssignValueOpPattern>(context));
+    ps.Add(std::make_unique<AssignValue_OpPattern>(context));
     return ps;
   }
 };
 
@@ -34,6 +34,7 @@
 from .impls.common import *  # noqa: F403
 from .impls.conv import *  # noqa: F403
 from .impls.creation import *  # noqa: F403
+from .impls.input import *  # noqa: F403
 from .impls.linalg import *  # noqa: F403
 from .impls.logic import *  # noqa: F403
 from .impls.manipulation import *  # noqa: F403
 
@@ -160,10 +160,9 @@ def __init__(
             min_subgraph_size (int, optional):
                 The minimum number of operations in a subgraph for TensorRT to optimize (default is 3).
             save_model_dir (str, optional):
-                The directory where the optimized model will be saved (default is None).
+                 The directory where the optimized model will be saved (default is None).
             disable_ops : (str|list, optional):
                 A string representing the names of operations that should not be entering by TensorRT (default is None).
-
         Returns:
             None
 
 
@@ -58,16 +58,35 @@ def assign_converter(network, paddle_op, inputs):
 
 
 @converter_registry.register("pd_op.assign_value_", trt_version="8.x")
+@converter_registry.register("pd_op.assign_value", trt_version="8.x")
 def assign_value_converter(network, paddle_op, inputs):
     attrs = paddle_op.attrs()
     shape = attrs['shape']
     dtype = attrs['dtype']
     values = attrs['values']
 
-    if dtype == paddle.float32:
-        dtype = np.float32
+    dtype_map = {
+        paddle.float32: np.float32,
+        paddle.int32: np.int32,
+    }
+    np_dtype = dtype_map.get(dtype)
+    if np_dtype is None:
+        raise NotImplementedError(
+            f"assign_value_ converter does not support dtype {dtype}"
+        )
+    # Initialize a NumPy array with zeros
+    np_values = np.zeros(shape, dtype=np_dtype)
+
+    # Flatten the NumPy array to a 1D array
+    flat_np_values = np_values.flatten()
+    # Assign values from the 'values' list to the flattened array
+    flat_np_values[: len(values)] = values
+    # Reshape the 1D array back to the original shape
+    np_values = flat_np_values.reshape(shape)
+
+    constant_layer = network.add_constant(shape=tuple(shape), weights=np_values)
+    constant_layer.name = paddle_op.name()
 
-    constant_layer = network.add_constant(shape, np.array(values, dtype=dtype))
     return constant_layer.get_output(0)
 
 
 
@@ -0,0 +1,62 @@
+# 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 numpy as np
+import tensorrt as trt
+
+from paddle.tensorrt.converter_utils import add_1D_constant_layer, cast_tensor
+from paddle.tensorrt.register import converter_registry
+
+
+@converter_registry.register(
+    "pd_op.one_hot", trt_version="trt_version_ge=8.5.1"
+)
+def one_hot_converter(network, paddle_op, inputs):
+    input_tensor, num_classes_tensor = inputs
+
+    input_type = input_tensor.dtype
+
+    trt_dtype_map = {
+        trt.DataType.INT32: trt.int32,
+    }
+    trt_dtype = trt_dtype_map.get(input_type, None)
+
+    trt_dtype = trt_dtype_map[input_type]
+
+    if trt_dtype == trt.int32:
+        values_data = [0, 1]
+        np_dtype = np.int32
+    # trt version>10  support int64
+    elif trt_dtype == trt.int64:
+        values_data = [0, 1]
+        np_dtype = np.int64
+    else:
+        raise ValueError(f"Unsupported trt_dtype for one_hot: {trt_dtype}")
+
+    values_tensor = add_1D_constant_layer(network, values_data, dtype=np_dtype)
+
+    reshape_layer = network.add_shuffle(num_classes_tensor)
+    reshape_layer.reshape_dims = ()
+    depth_tensor = reshape_layer.get_output(0)
+
+    depth_tensor = cast_tensor(network, depth_tensor, trt.int32)
+
+    one_hot_layer = network.add_one_hot(
+        input_tensor, values_tensor, depth_tensor, axis=-1
+    )
+    one_hot_layer.set_output_type(0, trt_dtype)
+    output_tensor = one_hot_layer.get_output(0)
+
+    return [output_tensor]
@@ -157,11 +157,15 @@ def topk_converter(network, paddle_op, inputs):
 
     input_shape = paddle_op.operands()[0].source().shape
 
-    k = paddle_op.attrs().get("k", 1)
     axis = paddle_op.attrs().get("axis", -1)
     largest = paddle_op.attrs().get("largest", True)
     flag = trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN
 
+    k_op = paddle_op.operands()[1].source().get_defining_op()
+    if k_op.name() == "pd_op.full":
+        k = k_op.attrs()["value"]
+    else:
+        raise NotImplementedError("Dynamic k is not supported in TensorRT.")
     input_rank = len(input_shape)
 
     expand_to_2d = input_rank == 1
@@ -174,7 +178,8 @@ def topk_converter(network, paddle_op, inputs):
 
     if axis < 0:
         axis += input_rank
-    layer = network.add_topk(input_tensor, flag, k, 1 << axis)
+
+    layer = network.add_topk(input_tensor, flag, int(k), 1 << axis)
     values = layer.get_output(0)
     indices = layer.get_output(1)
 
 
@@ -39,21 +39,52 @@ def get(self, op_name, trt_version=None):
         return self._registry.get(op_name)
 
     def _version_match(self, trt_version, version_range):
-        if version_range is None:
-            return True
-
-        trt_major, trt_minor = map(int, trt_version.split('.')[:2])
-        if version_range.startswith('trt_version_ge='):
-            min_version = float(version_range.split('=')[1])
-            return float(trt_major) + trt_minor / 10 >= min_version
-        elif version_range.startswith('trt_version_le='):
-            max_version = float(version_range.split('=')[1])
-            return float(trt_major) + trt_minor / 10 <= max_version
-        elif 'x' in version_range:
-            major_version = int(version_range.split('.')[0])
-            return trt_major == major_version
-        else:
-            return False
+        """
+        Check if a given TensorRT version matches the specified version range.
+
+        Args:
+            trt_version (str): The TensorRT version, e.g., "8.4.1".
+            version_range (str): The version range to check against, e.g.,
+                                "trt_version_ge=8.2", "trt_version_le=7.1", or "8.x".
+
+        Returns:
+            bool: True if the version matches the range, False otherwise.
+        """
+
+        def _normalize_version(version):
+            """
+            Normalize the version string into a 3-tuple for easy comparison.
+            If the version has fewer than 3 parts, it pads with zeros.
+
+            Args:
+                version (str): The version string, e.g., "8.4.1", "8.2", or "9".
+
+            Returns:
+                tuple: A tuple representing the version, e.g., (8, 4, 1).
+            """
+            return tuple(map(int, [*version.split('.'), '0', '0'][:3]))
+
+        # Convert the given TensorRT version to a normalized tuple
+        trt_version_tuple = _normalize_version(trt_version)
+        # Split the version range into comparator and reference version
+        if '=' in version_range:
+            comparator, ref_version = version_range.split('=')
+            # Normalize the reference version into a tuple
+            ref_version_tuple = _normalize_version(ref_version)
+            # Check the comparator and compare the versions
+            return (
+                comparator == 'trt_version_ge'
+                and trt_version_tuple >= ref_version_tuple
+            ) or (
+                comparator == 'trt_version_le'
+                and trt_version_tuple <= ref_version_tuple
+            )
+        # Check if the version range includes 'x' (e.g., "8.x")
+        if 'x' in version_range:
+            # Match only the major version (first part)
+            return trt_version_tuple[0] == int(version_range.split('.')[0])
+
+        return False
 
 
 converter_registry = ConverterOpRegistry()
@@ -23,7 +23,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
-  set_tests_properties(test_converter_clip PROPERTIES TIMEOUT "300")
+  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_search PROPERTIES TIMEOUT "300")
 
@@ -24,7 +24,21 @@ class TestShapeTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.shape
         self.api_args = {
-            "x": np.random.randn(2, 3).astype(np.float32),
+            "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 TestShapeTRTCase1Pattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.shape
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int64"),
         }
         self.program_config = {"feed_list": ["x"]}
         self.min_shape = {"x": [1, 3]}