from __future__ import annotations import functools import logging from dataclasses import dataclass, field from enum import Enum, auto from typing import ( Any, Callable, Collection, Dict, List, Optional, Sequence, Set, Tuple, Union, cast, ) import tensorrt as trt import torch import torch_tensorrt from torch import SymBool, SymFloat, SymInt from torch._ops import OpOverloadPacket from torch.fx.node import Argument, Node, Target, _get_qualified_name from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext logger = logging.getLogger(__name__) LegacyConverterImplSignature = Callable[ [ trt.INetworkDefinition, Target, Tuple[Argument, ...], Dict[str, Argument], str, ], Union[trt.ITensor, Sequence[trt.ITensor]], ] DynamoConverterImplSignature = Callable[ [ ConversionContext, Target, Tuple[Argument, ...], Dict[str, Argument], str, ], Union[trt.ITensor, Sequence[trt.ITensor]], ] ConverterImplSignature = Union[ LegacyConverterImplSignature, DynamoConverterImplSignature ] class ConverterPriority(Enum): """Enum to set a converter's priority in the registry""" STANDARD = auto() HIGH = auto() class CallingConvention(Enum): """Enum representing a converter's calling convention""" LEGACY = auto() # Legacy FX converters CTX = auto() # New Dynamo converters @dataclass(frozen=True) class ConverterSupport: """Class representing a converter implementation and support function Args: converter_implementation: Function which converts said node to a TRT equivalent capability_validator: Function which takes in a Node and returns a bool indicating whether that node can be supported by its companion converter. Note that this function must not modify the node or its graph supports_dynamic_shapes: Boolean flag indicating if the converter has support for dynamic inputs. requires_output_allocator: Boolean flag indicating if the converter creates operators which require an Output Allocator to run (e.g. data dependent operators). """ converter_implementation: ConverterImplSignature capability_validator: Callable[[Node, CompilationSettings], bool] = field( default=lambda node, compilation_settings: True ) supports_dynamic_shapes: bool = False requires_output_allocator: bool = False # Dictionary representing Dynamo aten-only converters # Each converter maps to a sequence of at least one ConverterSupport object(s) DYNAMO_ATEN_CONVERTERS: Dict[Target, Sequence[ConverterSupport]] = {} def has_static_shapes(node: torch.fx.Node) -> bool: """Returns True if a node has static args, kwargs, or outputs""" return not _has_dynamic_shapes(node=node) def node_has_dynamic_shapes(node: torch.fx.Node) -> bool: """Returns True if a node has dynamic args, kwargs, or outputs""" return _has_dynamic_shapes(node=node) def has_dynamic_shapes_in_args( arg_positions_to_check: Optional[List[int]] = None, ) -> Callable[[torch.fx.Node], bool]: """Returns True if a node has dynamic inputs in node.args at specified positions""" return functools.partial( _has_dynamic_shapes, arg_positions_to_check=arg_positions_to_check ) def has_static_shapes_in_args( arg_positions_to_check: Optional[List[int]] = None, ) -> Callable[[torch.fx.Node, CompilationSettings], bool]: """Returns True if a node has static inputs in node.args at specified positions""" _has_static_shapes = lambda node, compilation_settings, arg_positions_to_check: not _has_dynamic_shapes( node, compilation_settings, arg_positions_to_check ) return functools.partial( _has_static_shapes, arg_positions_to_check=arg_positions_to_check ) def _has_dynamic_shapes( node: torch.fx.Node, compilation_settings: CompilationSettings = None, 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""" _has_symbolic_sizes_strides, is_shape_dynamic = False, False if "val" in subnode.meta: _has_symbolic_sizes_strides = getattr( subnode.meta["val"], "_has_symbolic_sizes_strides", False ) meta_val = subnode.meta["val"] if isinstance(meta_val, (list, tuple)): for val in meta_val: shape = val.size() if any( isinstance(dim, (SymFloat, SymInt, SymBool)) for dim in shape ): is_shape_dynamic = True break elif isinstance(meta_val, (SymFloat, SymInt, SymBool)): is_shape_dynamic = True else: shape = subnode.meta["val"].size() is_shape_dynamic = any( isinstance(dim, (SymFloat, SymInt, SymBool)) for dim in shape ) return _has_symbolic_sizes_strides or is_shape_dynamic # Check node value itself if arg_positions_to_check is None and _is_subnode_dynamic(node): return True # Check node arguments individually 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 True # 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 True # Check node keyword arguments individually 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) ): return True return False def dynamo_tensorrt_converter( key: Target, *, enabled: bool = True, capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None, priority: ConverterPriority = ConverterPriority.STANDARD, supports_dynamic_shapes: bool = False, requires_output_allocator: bool = False, ) -> Callable[[ConverterImplSignature], ConverterImplSignature]: """Decorator for Dynamo TensorRT Converter Registers the decorated function in the DYNAMO_ATEN_CONVERTERS registry Args: key: Node target for which the converter is implemented for (for example, torch.ops.add.Tensor) enabled: Whether the converter should be enabled/cached or not capability_validator: Function which evaluates whether a node is valid for conversion by the decorated converter. See ConverterSupport for more details. Defaults to None, implying the capability_validator function is always true - this means all nodes of "key" kind can be supported by this converter priority: Converter's level of priority relative to other converters with the same target supports_dynamic_shapes: Boolean flag indicating if the converter has support for dynamic shapes. requires_output_allocator: Boolean flag indicating if the converter creates operators which require an Output Allocator to run (e.g. data dependent operators). Returns: The converter being decorated """ def register_converter(converter: ConverterImplSignature) -> ConverterImplSignature: """Helper function to register the converter, then return it""" assert callable(converter), "Converter function must be callable" # If no capability_validator function is specified, use the default function - always return true if capability_validator is None: converter_support = ConverterSupport( converter_implementation=converter, supports_dynamic_shapes=supports_dynamic_shapes, requires_output_allocator=requires_output_allocator, ) else: assert callable( capability_validator ), "Argument checking function must be callable" converter_support = ConverterSupport( converter_implementation=converter, capability_validator=capability_validator, supports_dynamic_shapes=supports_dynamic_shapes, requires_output_allocator=requires_output_allocator, ) # OpOverloadPackets are only valid if they have a single overload, or # only the ["default", "out"] overloads, due to PyTorch conventions if isinstance(key, OpOverloadPacket) and ( len(key.overloads()) >= 3 or (len(key.overloads()) == 2 and "out" not in key.overloads()) ): raise AssertionError( f"Detected converter for OpOverloadPacket {key}. " "We do not support OpOverloadPacket-keyed converters with multiple overloads. " "Make sure to explicitly specify each converter overload. For instance " "aten.mean is not a valid key, but aten.mean.default is." ) # If a converter for this operator already exists, append the new converter to the list # Otherwise, start a new list if key in DYNAMO_ATEN_CONVERTERS: # High priority converters are inserted at the front of the list, # so they can be checked first by the registry if priority is ConverterPriority.HIGH: DYNAMO_ATEN_CONVERTERS[key].insert(0, converter_support) else: DYNAMO_ATEN_CONVERTERS[key].append(converter_support) else: DYNAMO_ATEN_CONVERTERS[key] = [converter_support] logger.debug( f"Converter for {key} added to Dynamo ATen Converter Registry with priority: {priority}" ) return converter def disable_converter(converter: ConverterImplSignature) -> ConverterImplSignature: return converter # Select whether to cache/enable the converter if enabled: return register_converter else: return disable_converter class ConverterRegistry: """Registry for storing multiple converter dictionaries Capable of storing dictionaries with the following signature: Dict[Target, Union[Callable, Sequence[ConverterSupport]]] Also able to validate converter implementations against user-provided argument-checking functions Args: registries: List of dictionaries representing converter registries. The order of the provided dictionaries is the order in which they will be traversed. This is only significant when using non-validated methods registry_names: Optional list of names for each registry registry_calling_conventions: Optional list of calling conventions for each registry """ def __init__( self, registries: Sequence[ Dict[Target, Union[Callable[..., Any], Sequence[ConverterSupport]]] ], registry_names: Optional[Sequence[str]] = None, registry_calling_conventions: Optional[Sequence[CallingConvention]] = None, ): # Copy reference to each dictionary object into attribute list self.registries = list(registries) if registry_names is not None: assert len(self.registries) == len(registry_names) self.registry_names = list(registry_names) else: self.registry_names = [ f"Registry {i + 1}" for i in range(len(self.registries)) ] if registry_calling_conventions is not None: assert len(self.registries) == len(registry_calling_conventions) self.registry_calling_conventions = list(registry_calling_conventions) else: self.registry_calling_conventions = [ CallingConvention.CTX for _ in range(len(self.registries)) ] self.compilation_settings: CompilationSettings = None self.disallowed_targets: Collection[Target] = set() self.validate_invariants() def set_compilation_settings( self, compilation_settings: CompilationSettings ) -> None: self.compilation_settings = compilation_settings # set torch executed ops as disallowed targets self.set_disallowed_targets(compilation_settings.torch_executed_ops) def set_disallowed_targets(self, torch_executed_ops: Collection[Target]) -> None: self.disallowed_targets = torch_executed_ops def get_disallowed_targets(self, torch_executed_ops: Collection[Target]) -> None: self.disallowed_targets = torch_executed_ops def validate_invariants(self) -> None: """Validates the invariants required of the dictionaries in the registries Raises AssertionError if any invariants have been violated """ # All registries must be dictionaries assert all(isinstance(elt, dict) for elt in self.registries) # Every dictionary in the registry must have one of two signatures: # Dict[Target, Callable] or Dict[Target, Sequence[ConverterSupport]] # Where, for the latter, the sequence must be non-empty for registry in self.registries: for converters in registry.values(): if isinstance(converters, (list, tuple)): assert ( all(isinstance(c, ConverterSupport) for c in converters) and len(converters) > 0 ) else: assert callable(converters), "Converter function must be callable" def __getitem_without_validation__( self, key: Target ) -> Tuple[ Any, CallingConvention ]: # TODO: Narrow to ConverterImplSignature this when we can remove FX converters """Get the first-found converter in any registry Searches all registries in order and returns the first converter encountered, along with the calling convention of the registry the converter was sourced from """ if isinstance(key, Node): raise KeyError( "Unvalidated accesses to the Converter registry can only be " + "made with node targets. Try accessing the registry with node.target" ) self.validate_invariants() if ( key in self.disallowed_targets or self.qualified_name_or_str(key) in self.disallowed_targets ): raise KeyError( f"A converter exists for {key}, but it was " "explicitly disallowed" ) # Iterate over all registries and return the first converter found for registry, calling_convention in zip( self.registries, self.registry_calling_conventions ): if key in registry: converters = registry[key] if isinstance(converters, (list, tuple)): return converters[0].converter_implementation, calling_convention else: return converters, calling_convention raise KeyError(f"None of the converter registries have an entry for {key}") def __getitem__( self, node: Node ) -> Tuple[ Any, CallingConvention, Dict[str, bool] ]: # TODO: Narrow to ConverterImplSignature this when we can remove FX converters """Get the first-found validated converter in any registry Searches all registries in order and returns the first converter which passes validation on the input node, along with the calling convention of the registry the converter was sourced from """ if not isinstance(node, Node): raise KeyError( "Validated accesses to the Converter registry can only be " + "made with node inputs. Try accessing the registry with a node " + "or use get_unvalidated to access without node validation." ) self.validate_invariants() key = node.target assume_dynamic_shape_support = False if self.compilation_settings: assume_dynamic_shape_support = ( self.compilation_settings.assume_dynamic_shape_support ) if ( key in self.disallowed_targets or self.qualified_name_or_str(key) in self.disallowed_targets ): raise KeyError( f"A converter exists for {key}, but it was " "explicitly disallowed" ) # Iterate over all registries, validating the converter on the input node # If no capability_validator function is found, assume full coverage for registry, calling_convention in zip( self.registries, self.registry_calling_conventions ): if key in registry: converters = registry[key] if isinstance(converters, (list, tuple)): logger.debug(f"Converter options for {key}: {len(converters)}") for i, candidate in enumerate(converters): # We enable the converter under 4 conditions # 1) capability validator is True # 2) Assume dynamic_shape support is True # 3) Node only has static shaped inputs # 4) Node has dynamic inputs and the converter has supports_dynamic_shapes=True if is_valid := candidate.capability_validator( node, self.compilation_settings ) and ( assume_dynamic_shape_support or not node_has_dynamic_shapes(node) or candidate.supports_dynamic_shapes ): logger.debug( f"Selecting converter option {i} for converting {key}" ) return ( candidate.converter_implementation, calling_convention, { "supports_dynamic_shapes": candidate.supports_dynamic_shapes, "requires_output_allocator": candidate.requires_output_allocator, }, ) else: logger.debug( f"Skipping option {i} for {key}: (validator: {is_valid}, supports dynamic shapes: {candidate.supports_dynamic_shapes})" ) continue else: # Assuming FX converters don't have dynamic shapes supported if not node_has_dynamic_shapes(node): return ( converters, calling_convention, { "supports_dynamic_shapes": False, "requires_output_allocator": False, }, ) raise KeyError( f"None of the converter registries have a validated entry for {key}, with node {node}" ) def keys(self) -> Set[Target]: """Get all unique targets across all dictionaries""" return self.unique_targets() def get_unvalidated( self, key: Target, value: Optional[ConverterImplSignature] = None ) -> Union[ Any, Tuple[Any, CallingConvention] ]: # TODO: Narrow to ConverterImplSignature this when we can remove FX converters """Get unvalidated converter for input target with a default return""" try: return self.__getitem_without_validation__(key) except KeyError: return value def get( self, node: Node, value: Optional[ConverterImplSignature] = None ) -> Union[ Any, Tuple[Any, CallingConvention, Dict[str, bool]] ]: # TODO: Narrow to ConverterImplSignature this when we can remove FX converters """Get validated converter for input node with a default return""" try: return self.__getitem__(node) except KeyError: return value def __contains__(self, key: Target | Node) -> bool: """Check whether a converter for an input node or target exists""" try: # Attempt to access the item in the registry if isinstance(key, Node): self.__getitem__(key) else: self.__getitem_without_validation__(key) return True except KeyError: return False def get_all_converters_with_target( self, key: Target, return_registry_info: bool = False ) -> Tuple[ Union[List[Any], Dict[str, int], None] ]: # TODO: Narrow to ConverterImplSignature this when we can remove FX converters """Get all converters across all registries for the target Returns a list of all converterts having the specified target """ self.validate_invariants() converters_with_target = [] # Store count of number of registered converters per registry if return_registry_info: registry_data = {name: 0 for name in self.registry_names} for index, registry in enumerate(self.registries): if key in registry: converters = registry[key] if isinstance(converters, (list, tuple)): converters_with_target.extend( [c.converter_implementation for c in converters] ) # Add converter count to registry name storage if return_registry_info: registry_data[self.registry_names[index]] += len(converters) else: converters_with_target.append(converters) # Add converter count to registry name storage if return_registry_info: registry_data[self.registry_names[index]] += 1 if return_registry_info: return converters_with_target, registry_data else: return converters_with_target, None def __setitem__(self, key: Any, value: Any) -> None: raise AssertionError( "Do not set registry members directly through the ConverterRegistry object. " + f"Attempted to set {key}: {value} via direct assignment to ConverterRegistry." ) def __delitem__(self, key: Any) -> None: raise AssertionError( "Do not delete registry members directly through the ConverterRegistry object. " + f"Attempted to delete {key} via direct del on ConverterRegistry." ) def __len__(self) -> int: """Returns the sum of lengths of all registries stored""" return sum(len(registry) for registry in self.registries) def unique_targets(self) -> Set[Target]: """Returns the set of unique converter targets stored across all registries""" return set.union(*[set(registry.keys()) for registry in self.registries]) @staticmethod def qualified_name_or_str(target: Target) -> str: """Returns string representation of an FX Node target""" if isinstance(target, str): return target else: return cast(str, _get_qualified_name(target)) def get_converter_support_info(self) -> Dict[str, Optional[Dict[str, int]]]: """Returns a dictionary of targets backed by at least one converter""" available_converters = {} for target in sorted( self.unique_targets(), key=lambda target: self.qualified_name_or_str(target) ): _, registry_data = self.get_all_converters_with_target( target, return_registry_info=True ) available_converters[self.qualified_name_or_str(target)] = registry_data return available_converters def display_all_available_converters(self) -> str: """Returns a string with all converters and their source, separated by newlines""" available_converters = "Available converters in ATen registries with counts:\n" support_info = self.get_converter_support_info() for target, registry_data in support_info.items(): available_converters += f"Node: {self.qualified_name_or_str(target)} - Registry Presence Counts: {registry_data}\n" return available_converters # Initialize dynamo converter registry with the FX and Dynamo aten registries # Note the Dynamo registry is listed first, for precedence registries = [ DYNAMO_ATEN_CONVERTERS, ] registry_names = ["Dynamo ATen Converters Registry"] registry_calling_conventions = [ CallingConvention.CTX, ] if torch_tensorrt.ENABLED_FEATURES.fx_frontend: from torch_tensorrt.fx.converter_registry import CONVERTERS as FX_CONVERTERS registries.append(FX_CONVERTERS) registry_names.append("FX Legacy ATen Converters Registry") registry_calling_conventions.append(CallingConvention.LEGACY) DYNAMO_CONVERTERS: ConverterRegistry = ConverterRegistry( registries, registry_names, registry_calling_conventions, )