Merge pull request #2070 from pytorch/dynamo_latest

fix: Move all `aten` PRs to Dynamo converter registry

Author: narendasan

.circleci/config.yml

@@ -780,6 +780,21 @@ commands:
       - store_artifacts:
           path: /tmp/testlogs
 
+  test-dynamo-converters:
+    description: "Test the Dynamo aten converters"
+    steps:
+      - run:
+          name: Run Dynamo converter tests
+          command: |
+            cd tests/py/dynamo/converters
+            TESTS_TO_RUN=$(circleci tests glob "test_*.py" | circleci tests split --split-by=timings)
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/converters/test_results.xml $TESTS_TO_RUN
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
 # =================== Dynamo tests end ======================== #
 
 # Define a job to be invoked later in a workflow.
@@ -1036,6 +1051,7 @@ jobs:
           command: pip3 install --pre /tmp/dist/x86_64-linux/*cp39-cp39*.whl
       # We install torch after torch-trt because pip automatically enforces the version constraint otherwise
       - dump-test-env
+      - test-dynamo-converters
       - test-dynamo-torch_compile
       - test-dynamo-models_torch_compile
       - test-dynamo-models_torch_export

py/torch_tensorrt/_Input.py

@@ -68,6 +68,17 @@ def __init__(self, *args, **kwargs):
             - Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)
             - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW
         """
+        # Compatibility code for switching over from InputTensorSpec
+        if "shape" in kwargs and "shape_ranges" in kwargs:
+            assert (
+                len(kwargs["shape_ranges"]) == 1 and len(kwargs["shape_ranges"][0]) == 3
+            )
+            del kwargs["shape"]
+
+            kwargs["min_shape"] = kwargs["shape_ranges"][0][0]
+            kwargs["opt_shape"] = kwargs["shape_ranges"][0][1]
+            kwargs["max_shape"] = kwargs["shape_ranges"][0][2]
+
         if len(args) == 1:
             if not Input._supported_input_size_type(args[0]):
                 raise TypeError(

py/torch_tensorrt/dynamo/__init__.py

@@ -3,8 +3,9 @@
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from ._settings import *
+    from .conversion import *
     from .aten_tracer import trace
-    from .converter_registry import (
+    from .conversion.converter_registry import (
         DYNAMO_CONVERTERS,
         dynamo_tensorrt_converter,
     )

py/torch_tensorrt/dynamo/conversion/SourceIR.py

@@ -0,0 +1,24 @@
+from enum import Enum, auto
+
+
+class SourceIR(Enum):
+    NN = auto()
+    ACC = auto()
+    ATEN = auto()
+    PRIM = auto()
+    TORCHTRT_LOWERED = auto()
+    UNKNOWN = auto()
+
+    def __str__(self):
+        if self == SourceIR.NN:
+            return "nn"
+        elif self == SourceIR.ACC:
+            return "acc"
+        elif self == SourceIR.ATEN:
+            return "aten"
+        elif self == SourceIR.PRIM:
+            return "prim"
+        elif self == SourceIR.TORCHTRT_LOWERED:
+            return "torchtrt_lowered"
+        else:
+            return "unknown_ir"

py/torch_tensorrt/dynamo/conversion/__init__.py

@@ -1,3 +1,5 @@
+from .SourceIR import SourceIR
+from .aten_ops_converters import *
 from .trt_interpreter import *
 from .conversion import *
 from .truncate_long_and_double import repair_long_or_double_inputs