feat/fix: Update dynamic unsupported implementation

- Add support for selecting individual argument positions to check and
expand checking to include symbolic types, which are sometimes passed in
as arguments

Author: gs-olive

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -349,36 +349,25 @@ def aten_ops_softmax(
     return impl.normalization.softmax(
         network, target, SourceIR.ATEN, name, args[0], args[1]
     )
-    
-def dynamic_unsupported_split(node: torch.fx.Node) -> bool:
-    # Validate that none of the inputs to the node have Dynamic shapes
-    assert isinstance(
-        node, torch.fx.Node
-    ), "Inputs to validator functions must be FX Nodes"
-
-    if isinstance(node.args[1], torch.fx.Node):
-        if getattr(node.args[1].meta["val"], "_has_symbolic_sizes_strides", True):
-            return False
-    return True
 
 
 @dynamo_tensorrt_converter(
-    torch.ops.aten.split.Tensor, capability_validator=dynamic_unsupported_split
+    torch.ops.aten.split.Tensor, capability_validator=dynamic_unsupported_with_args([1])
 )
 @dynamo_tensorrt_converter(
-    torch.ops.aten.split.sizes, capability_validator=dynamic_unsupported_split
+    torch.ops.aten.split.sizes, capability_validator=dynamic_unsupported_with_args([1])
 )
 @dynamo_tensorrt_converter(
     torch.ops.aten.split_with_sizes.default,
-    capability_validator=dynamic_unsupported_split,
+    capability_validator=dynamic_unsupported_with_args([1]),
 )
 def aten_ops_split(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]: 
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
     return impl.split.split(
         network,
         target,

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -6,6 +6,7 @@
 import numpy as np
 import tensorrt as trt
 import torch
+from torch import SymBool, SymFloat, SymInt
 from torch.fx.node import Target
 from torch_tensorrt.fx.converters.converter_utils import (
     Frameworks,
@@ -46,26 +47,51 @@ def get_node_name(node: torch.fx.Node) -> str:
 
 
 def dynamic_unsupported(node: torch.fx.Node) -> bool:
+    """Validates that a node has no dynamic args, kwargs, or outputs"""
+    return _dynamic_unsupported(node=node)
+
+
+def dynamic_unsupported_with_args(
+    arg_positions_to_check: Optional[List[int]] = None,
+) -> Callable[[torch.fx.Node], bool]:
+    """Returns a validator that a node has no dynamic args at specific positions"""
+    return functools.partial(_dynamic_unsupported, arg_positions_to_check=arg_positions_to_check)
+
+
+def _dynamic_unsupported(
+    node: torch.fx.Node, arg_positions_to_check: Optional[List[int]] = None
+) -> bool:
     # Validate that none of the inputs to the node have Dynamic shapes
     assert isinstance(
         node, torch.fx.Node
     ), "Inputs to validator functions must be FX Nodes"
 
+    def _is_subnode_dynamic(subnode: torch.fx.Node) -> bool:
+        """Checks if a node itself has Dynamic properties"""
+        return getattr(
+            subnode.meta["val"], "_has_symbolic_sizes_strides", False
+        ) or isinstance(subnode.meta["val"], (SymFloat, SymInt, SymBool))
+
     # Check node value itself
-    if getattr(node.meta["val"], "_has_symbolic_sizes_strides", False):
+    if arg_positions_to_check is None and _is_subnode_dynamic(node):
         return False
 
     # Check node arguments individually
-    if any(
-        getattr(arg.meta["val"], "_has_symbolic_sizes_strides", False)
-        for arg in node.args
-        if isinstance(arg, torch.fx.Node)
+    if arg_positions_to_check is None and any(
+        _is_subnode_dynamic(arg) for arg in node.args if isinstance(arg, torch.fx.Node)
+    ):
+        return False
+    # Check specific arg positions if the caller has specified positions to check
+    elif arg_positions_to_check is not None and any(
+        _is_subnode_dynamic(node.args[i])
+        for i in arg_positions_to_check
+        if isinstance(node.args[i], torch.fx.Node)
     ):
         return False
 
     # Check node keyword arguments individually
-    if any(
-        getattr(kwarg.meta["val"], "_has_symbolic_sizes_strides", False)
+    if arg_positions_to_check is None and any(
+        _is_subnode_dynamic(kwarg)
         for kwarg in node.kwargs.values()
         if isinstance(kwarg, torch.fx.Node)
     ):