chore: Cannot support dynamic k in topk

Author: keehyuna

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -19,7 +19,6 @@
     get_positive_dim,
     is_only_operator_on_placeholder,
 )
-from torch_tensorrt.dynamo.utils import TRT_TOPK_MAX_ELEMENT
 from torch_tensorrt.fx.types import TRTTensor
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -2649,21 +2648,22 @@ def topk_validator(node: Node) -> bool:
 
 
 def sort_validator(node: Node) -> bool:
-    # if meta data is not available(e.g. dynamic shape), validate k value during runtime.
-    if not node.args[0].meta:
-        return True
-
-    shape = node.args[0].meta.get("tensor_meta").shape
+    meta_data = node.args[0].meta.get("tensor_meta")
+    if meta_data is None:
+        return False
+    shape = meta_data.shape
     dim = node.args[1]
     dim = get_positive_dim(dim, len(shape))
     k = shape[dim]
+    if not isinstance(k, int):
+        return False
     return topk_sort_validator(k)
 
 
 def topk_sort_validator(k: int) -> bool:
-    if k > TRT_TOPK_MAX_ELEMENT:
+    if k > 3840:
         _LOGGER.debug(
-            f"Currently only topk values up to {TRT_TOPK_MAX_ELEMENT} are supported, got k={k}."
+            f"Currently only topk values up to 3840 are supported, got k={k}."
         )
         return False
     return True

py/torch_tensorrt/dynamo/conversion/impl/topk.py

@@ -12,10 +12,8 @@
     get_positive_dim,
     set_layer_name,
 )
-from torch_tensorrt.dynamo.conversion.impl.elementwise.ops import le
-from torch_tensorrt.dynamo.conversion.impl.shape import get_shape_with_dynamic_shape
 from torch_tensorrt.dynamo.types import TRTTensor
-from torch_tensorrt.dynamo.utils import DYNAMIC_DIM, TRT_TOPK_MAX_ELEMENT
+from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
 
 
 def argmax_argmin(
@@ -158,40 +156,14 @@ def topk(
             k,
             get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape))),
         )
-    if k == DYNAMIC_DIM:
-        output_shape = get_shape_with_dynamic_shape(
-            ctx, target, source_ir, name, input.shape, input
-        )
-        layer = ctx.net.add_slice(
-            output_shape,
-            start=[dim],
-            shape=[1],
-            stride=[1],
-        )
-        set_layer_name(layer, target, name)
 
-        # Get scalar tensor from 1d tensor
-        shuffle_layer = ctx.net.add_shuffle(layer.get_output(0))
-        shuffle_layer.reshape_dims = trt.Dims()
-        set_layer_name(shuffle_layer, target, name, source_ir)
+    # topk layer supports dynamic k value but we cannot dertermin supported dynamic topk value at
+    # compile time.
+    assert k != DYNAMIC_DIM, "k value cannot be dynamic!"
 
-        cond = le(
-            ctx,
-            target,
-            source_ir,
-            f"{name}_k_cond",
-            shuffle_layer.get_output(0),
-            TRT_TOPK_MAX_ELEMENT,
-        )
-        ctx.net.add_assertion(
-            cond,
-            message=f"Currently only topk values up to {TRT_TOPK_MAX_ELEMENT} are supported",
-        )
-
-        topk_layer.set_input(1, shuffle_layer.get_output(0))
     # TensorRT ITopKLayer does not have a sorted flag, it is always returning the sorted topk elements
     # so here no matter sorted is True or False the returned the topk Tensor object is always sorted
-    set_layer_name(topk_layer, target, name, source_ir)
+    set_layer_name(topk_layer, target, f"{name}_topk", source_ir)
 
     if return_indices:
         return topk_layer.get_output(0), topk_layer.get_output(1)

py/torch_tensorrt/dynamo/utils.py

@@ -6,14 +6,14 @@
 from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 import numpy as np
-import tensorrt as trt
 import torch
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo import _defaults
 from torch_tensorrt.dynamo._settings import CompilationSettings
 
+import tensorrt as trt
 from packaging import version
 
 from .types import TRTDataType
@@ -22,7 +22,6 @@
 
 COSINE_THRESHOLD = 0.99
 DYNAMIC_DIM = -1
-TRT_TOPK_MAX_ELEMENT = 3840
 
 
 class Frameworks(Enum):

tests/py/dynamo/conversion/test_sort_aten.py

@@ -38,24 +38,24 @@ class TestSortConverterDynamic(DispatchTestCase):
         [
             (
                 "3d_dynamic_descending",
-                (2, 2, 1),
-                (2, 2, 1),
+                (2, 1, 4),
                 (3, 2, 4),
-                0,
+                (3, 3, 4),
+                2,
                 True,
             ),
             (
                 "4d_dynamic_ascending",
-                (2, 2, 1, 1),
-                (2, 2, 1, 2),
+                (2, 2, 1, 4),
+                (2, 2, 2, 4),
                 (3, 3, 2, 4),
                 3,
                 False,
             ),
             (
                 "4d_dynamic_descending_neg_dim",
-                (2, 2, 1, 1),
-                (2, 2, 1, 2),
+                (1, 3, 1, 1),
+                (2, 3, 2, 2),
                 (3, 3, 2, 4),
                 -3,
                 True,
@@ -79,6 +79,7 @@ def forward(self, x):
             Sort(),
             input_specs,
             output_dtypes=[torch.float, torch.int64],
+            use_dynamo_tracer=True,
         )