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

@@ -2,5 +2,9 @@
 from packaging import version
 
 if version.parse(torch.__version__) >= version.parse("2.1.dev"):
+    from .converter_registry import (
+        DYNAMO_CONVERTERS,
+        dynamo_tensorrt_converter,
+    )
     from torch_tensorrt.dynamo import fx_ts_compat
     from .backend import compile

py/torch_tensorrt/dynamo/backend/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__)

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/fx_ts_compat/fx2trt.py

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