feat: Data Structure update for Dynamo Registry

- Add custom class overriding default Dictionary class to access
converters from various registries
- Add new dictionary type `Dict[Target, Sequence[ConverterSupport]]` as
well as ConverterSupport class which stores a converter and its
validation implementation
- Add unified `DYNAMO_CONVERTERS` dictionary which coalesces both the FX
and Dynamo converter dictionaries and acts as a single unified
dictionary

Author: gs-olive

py/torch_tensorrt/dynamo/backend/lowering/_partition.py

@@ -124,7 +124,7 @@ def is_node_supported(
         )
 
         if (
-            node.target in CONVERTERS.keys()
+            CONVERTERS.contains_validated(node)
             and node_name not in self.torch_executed_ops
         ):
             # If node is a proper, supported computational node, store the operator

py/torch_tensorrt/dynamo/converter_registry.py

@@ -1,23 +1,253 @@
-from typing import Any, Callable, Dict
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
-from torch.fx.node import Target
+from torch.fx.node import Target, Node
 from torch_tensorrt.fx.converter_registry import CONVERTERS
 
-DYNAMO_CONVERTERS: Dict[Target, Any] = dict(CONVERTERS)
+
+@dataclass(frozen=True)
+class GenericNode:
+    """Class representing a generic Torch FX Node
+
+    Convenience class for when we do not have access to the node,
+    and instead only its target, args, and kwargs, as is the case
+    in the TRTInterpreter
+    """
+
+    target: Any
+    args: Sequence[Any]
+    kwargs: Dict[str, Any]
+
+
+# Defines the different types of valid nodes
+NodeKinds = Union[GenericNode, Node]
+
+
+@dataclass(frozen=True)
+class ConverterSupport:
+    """Class representing a converter implementation and support function
+
+    Args:
+        check_args: Function which takes in a NodeKind and returns a bool indicating
+            whether that node can be supported by its companion converter. Note that
+            this function must only access the .target, .args, and .kwargs fields
+            of the node, and cannot modify the node or its graph
+        converter_implementation: Function which converts said node to a TRT equivalent
+    """
+
+    check_args: Callable[[NodeKinds], bool] = field(default=lambda node: True)
+    converter_implementation: Callable
+
+
+# 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 dynamo_tensorrt_converter(
     key: Target,
     enabled: bool = True,
+    check_args: Optional[Callable[[NodeKinds], bool]] = None,
 ) -> Callable[[Any], Any]:
+    """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
+        check_args: Function which evaluates whether a node is valid for conversion
+            by the decorated converter. See ConverterSupport for more details.
+            Defaults to None, implying the check_args function is always true -
+            this means all nodes of "key" kind can be supported by this converter
+    Returns:
+        The converter being decorated
+    """
+
     def register_converter(converter):
-        DYNAMO_CONVERTERS[key] = converter
+        """Helper function to register the converter, then return it"""
+        assert callable(converter), "Converter function must be callable"
+
+        # If no check_args function is specified, use the default function - always return true
+        if check_args is None:
+            converter_support = ConverterSupport(converter_implementation=converter)
+        else:
+            assert callable(check_args), "Argument checking function must be callable"
+            converter_support = ConverterSupport(
+                check_args=check_args, converter_implementation=converter
+            )
+
+        # 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:
+            DYNAMO_ATEN_CONVERTERS[key].append(converter_support)
+        else:
+            DYNAMO_ATEN_CONVERTERS[key] = [converter_support]
+
         return converter
 
     def disable_converter(converter):
         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 two different signatures:
+    Dict[Target, Callable] and Dict[Target, Sequence[ConverterSupport]]
+    """
+
+    def __init__(self, registries: Dict[Target, Any]):
+        # Copy reference to each dictionary object into attribute list
+        self.registries = [registry for registry in registries]
+        self.validate_invariants()
+
+    def validate_invariants(self):
+        """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__(self, key: Target):
+        """Get the first-found converter in any registry
+
+        Searches all registries in order and returns the first converter encountered
+        """
+        self.validate_invariants()
+
+        # Iterate over all registries and return the first converter found
+        for registry in self.registries:
+            if key in registry:
+                converters = registry[key]
+
+                if isinstance(converters, (list, tuple)):
+                    return converters[0].converter_implementation
+                else:
+                    return converters
+
+        raise KeyError(f"None of the converter registries have an entry for {key}")
+
+    def __getitem_with_validation__(self, node: NodeKinds):
+        """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
+        """
+        self.validate_invariants()
+        key = node.target
+
+        # Iterate over all registries, validating the converter on the input node
+        # If no check_args function is found, assume full coverage
+        for registry in self.registries:
+            if key in registry:
+                converters = registry[key]
+
+                if isinstance(converters, (list, tuple)):
+                    for candidate in converters:
+                        if candidate.check_args(node):
+                            return candidate.converter_implementation
+                else:
+                    return converters
+
+        raise KeyError(
+            f"None of the converter registries have a validated entry for {key}, with node {node}"
+        )
+
+    def keys(self):
+        """Get all unique targets across all dictionaries"""
+        return self.unique_targets()
+
+    def get(self, key: Target, value=None):
+        """Get converter for input target with a default return"""
+        try:
+            return self.__getitem__(key)
+        except KeyError:
+            return value
+
+    def __contains__(self, key: Target):
+        """Check whether a converter for input target exists"""
+        return any(key in registry for registry in self.registries)
+
+    def get_validated(self, node: NodeKinds, value=None):
+        """Get validated converter for input node with a default return"""
+        try:
+            return self.__getitem_with_validation__(node)
+        except KeyError:
+            return value
+
+    def contains_validated(self, node: NodeKinds):
+        """Check whether a validated converter for input node exists"""
+        try:
+            self.__getitem_with_validation__(node)
+            return True
+        except KeyError:
+            return False
+
+    def get_all_converters_with_target(self, key: Target):
+        """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 = []
+
+        for registry in 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]
+                    )
+                else:
+                    converters_with_target.append(converters)
+
+        return converters_with_target
+
+    def __setitem__(self, key, value):
+        raise AssertionError(
+            f"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):
+        raise AssertionError(
+            f"Do not delete registry members directly through the ConverterRegistry object. "
+            + f"Attempted to delete {key} via direct del on ConverterRegistry."
+        )
+
+    def __len__(self):
+        """Returns the sum of lengths of all registries stored"""
+        return sum(len(registry) for registry in self.registries)
+
+    def unique_targets(self):
+        """Returns the set of unique converter targets stored across all registries"""
+        return set.union(*[set(registry.keys()) for registry in self.registries])
+
+
+# Initialize dynamo converter registry with the FX and Dynamo aten registries
+DYNAMO_CONVERTERS: ConverterRegistry = ConverterRegistry(
+    [CONVERTERS, DYNAMO_ATEN_CONVERTERS]
+)

py/torch_tensorrt/dynamo/fx_ts_compat/fx2trt.py

@@ -15,6 +15,7 @@
 from torch.fx.passes.shape_prop import TensorMetadata
 
 from torch_tensorrt.dynamo import DYNAMO_CONVERTERS as CONVERTERS
+from torch_tensorrt.dynamo.converter_registry import GenericNode
 from .input_tensor_spec import InputTensorSpec
 from torch_tensorrt.fx.observer import Observer
 from torch_tensorrt.fx.utils import (
@@ -141,9 +142,9 @@ def validate_conversion(self):
         missing_converter = set()
 
         for node in self.module.graph.nodes:
-            if node.op == "call_function" and not CONVERTERS.get(node.target):
+            if node.op == "call_function" and not CONVERTERS.get_validated(node):
                 missing_converter.add(f"{node.op} {_get_qualified_name(node.target)}")
-            elif node.op == "call_method" and not CONVERTERS.get(node.target):
+            elif node.op == "call_method" and not CONVERTERS.get_validated(node):
                 missing_converter.add(f"{node.op} torch.Tensor.{node.target}")
             elif node.op == "call_module":
                 submod = self.fetch_attr(node.target)
@@ -347,7 +348,7 @@ def call_module(self, target, args, kwargs):
         return converter(self.network, submod, args, kwargs, self._cur_node_name)
 
     def call_function(self, target, args, kwargs):
-        converter = CONVERTERS.get(target)
+        converter = CONVERTERS.get_validated(GenericNode(target, args, kwargs))
         if not converter:
             raise RuntimeError(
                 f"Conversion of function {torch.typename(target)} not currently supported!"
@@ -358,7 +359,7 @@ def call_function(self, target, args, kwargs):
 
     def call_method(self, target, args, kwargs):
         assert isinstance(target, str)
-        converter = CONVERTERS.get(target)
+        converter = CONVERTERS.get_validated(GenericNode(target, args, kwargs))
 
         if not converter:
             raise RuntimeError(