fix/feat: Add Dynamo-only converter registry

- Add Dynamo converter registry which functions as a superset of the
standard FX converter registry
- For use with new + experimental converters
- Uses custom decorator `dynamo_tensorrt_converter`
- Update references within Dynamo functions to use the converter
registry `DYNAMO_CONVERTERS`

Author: gs-olive

py/torch_tensorrt/dynamo/__init__.py

@@ -5,3 +5,7 @@
     from ._settings import *
     from .compile import compile
     from .aten_tracer import trace
+    from .converter_registry import (
+        DYNAMO_CONVERTERS,
+        dynamo_tensorrt_converter,
+    )

py/torch_tensorrt/dynamo/conversion/trt_interpreter.py

@@ -14,7 +14,7 @@
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes.shape_prop import TensorMetadata
 
-from torch_tensorrt.fx import CONVERTERS
+from torch_tensorrt.dynamo import DYNAMO_CONVERTERS as CONVERTERS
 from torch_tensorrt import Input
 from torch_tensorrt.fx.observer import Observer
 from torch_tensorrt.fx.utils import (

py/torch_tensorrt/dynamo/converter_registry.py

@@ -0,0 +1,23 @@
+from typing import Any, Callable, Dict
+
+from torch.fx.node import Target
+from torch_tensorrt.fx.converter_registry import CONVERTERS
+
+DYNAMO_CONVERTERS: Dict[Target, Any] = dict(CONVERTERS)
+
+
+def dynamo_tensorrt_converter(
+    key: Target,
+    enabled: bool = True,
+) -> Callable[[Any], Any]:
+    def register_converter(converter):
+        DYNAMO_CONVERTERS[key] = converter
+        return converter
+
+    def disable_converter(converter):
+        return converter
+
+    if enabled:
+        return register_converter
+    else:
+        return disable_converter

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -10,7 +10,7 @@
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes.operator_support import OperatorSupport
 
-from torch_tensorrt.fx.converter_registry import CONVERTERS
+from torch_tensorrt.dynamo import DYNAMO_CONVERTERS as CONVERTERS
 
 
 logger = logging.getLogger(__name__)