fix: Reorganize, add tests

Author: gs-olive

.circleci/config.yml

@@ -707,6 +707,22 @@ commands:
       - store_artifacts:
           path: /tmp/testlogs
 
+  test-dynamo-torch_compile:
+    description: "Test the Dynamo torch_compile path"
+    steps:
+      - run:
+          name: Run Dynamo torch_compile E2E tests
+          command: |
+            cd py/torch_tensorrt/dynamo/
+            pushd test/
+            pip install timm
+            pytest --junitxml=/tmp/artifacts/test_results/dynamo/test_results.xml --ir torch_compile
+            popd
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
 # Define a job to be invoked later in a workflow.
 # See: https://circleci.com/docs/2.0/configuration-reference/#jobs
 jobs:
@@ -883,6 +899,68 @@ jobs:
       - dump-test-env
       - test-fx-no-aten
 
+  test-py-dynamo-x86_64-linux:
+    parameters:
+      torch-build:
+        type: string
+      torch-build-index:
+        type: string
+      trt-version-long:
+        type: string
+    machine:
+      image: ubuntu-2004-cuda-11.4:202110-01
+    resource_class: gpu.nvidia.large
+    steps:
+      - checkout
+      - attach_workspace:
+          at: /tmp/dist/
+      - install-torch-from-index:
+          torch-build:  << parameters.torch-build >>
+          torch-build-index: << parameters.torch-build-index >>
+      - create-py-env:
+          trt-version-long: << parameters.trt-version-long >>
+      - install-cudnn
+      # - run:
+      #     name: "Set LD_LIBRARY_PATH path to include the installed CUDNN"
+      #     command: export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
+      - run:
+          name: "Install torch-tensorrt"
+          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-torch_compile
+
+  test-py-dynamo-x86_64-linux-no-aten:
+    parameters:
+      torch-build:
+        type: string
+      torch-build-index:
+        type: string
+      trt-version-long:
+        type: string
+    machine:
+      image: ubuntu-2004-cuda-11.4:202110-01
+    resource_class: gpu.nvidia.large
+    steps:
+      - checkout
+      - attach_workspace:
+          at: /tmp/dist/
+      - install-torch-from-index:
+          torch-build:  << parameters.torch-build >>
+          torch-build-index: << parameters.torch-build-index >>
+      - create-py-env:
+          trt-version-long: << parameters.trt-version-long >>
+      - install-cudnn
+      # - run:
+      #     name: "Set LD_LIBRARY_PATH path to include the installed CUDNN"
+      #     command: export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
+      - run:
+          name: "Install torch-tensorrt"
+          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-torch_compile
+
   package-x86_64-linux:
     parameters:
       enabled:
@@ -1261,6 +1339,13 @@ workflows:
           requires:
             - build-x86_64-linux
 
+      - test-py-dynamo-x86_64-linux:
+          torch-build: << pipeline.parameters.torch-build >>
+          torch-build-index: << pipeline.parameters.torch-build-index >>
+          trt-version-long: << pipeline.parameters.trt-version-long >>
+          requires:
+            - build-x86_64-linux
+
       - build-x86_64-linux:
           name: build-x86_64-linux-legacy
           torch-build: << pipeline.parameters.torch-build-legacy >>
@@ -1291,6 +1376,12 @@ workflows:
           requires:
             - build-x86_64-linux-legacy
 
+      - test-py-dynamo-x86_64-linux-no-aten:
+          torch-build: << pipeline.parameters.torch-build-legacy >>
+          torch-build-index: << pipeline.parameters.torch-build-index-legacy >>
+          trt-version-long: << pipeline.parameters.trt-version-long >>
+          requires:
+            - build-x86_64-linux-legacy
   release:
     when: << pipeline.parameters.enable-packaging >>
     jobs:
@@ -1328,6 +1419,13 @@ workflows:
           requires:
             - package-x86_64-linux
 
+      - test-py-dynamo-x86_64-linux:
+          torch-build: << pipeline.parameters.torch-build >>
+          torch-build-index: << pipeline.parameters.torch-build-index >>
+          trt-version-long: << pipeline.parameters.trt-version-long >>
+          requires:
+            - package-x86_64-linux
+
   on-push:
     jobs:
       - build-x86_64-linux:
@@ -1357,6 +1455,13 @@ workflows:
           requires:
             - build-x86_64-linux
 
+      - test-py-dynamo-x86_64-linux:
+          torch-build: << pipeline.parameters.torch-build >>
+          torch-build-index: << pipeline.parameters.torch-build-index >>
+          trt-version-long: << pipeline.parameters.trt-version-long >>
+          requires:
+            - build-x86_64-linux
+
       - build-x86_64-linux-cmake:
           torch-build: << pipeline.parameters.torch-build >>
           torch-build-index: << pipeline.parameters.torch-build-index >>

py/torch_tensorrt/__init__.py

@@ -93,6 +93,8 @@ def _find_lib(name, paths):
 from torch_tensorrt._TRTModuleNext import TRTModuleNext
 
 from torch_tensorrt import fx
+from torch_tensorrt import dynamo
+from torch_tensorrt.dynamo import torch_compile
 
 
 def _register_with_torch():

py/torch_tensorrt/_compile.py

@@ -15,6 +15,8 @@ class _IRType(Enum):
 
     ts = 0
     fx = 1
+    torch_compile = 2
+    fx_ts_compat_compile = 3
 
 
 class _ModuleType(Enum):
@@ -45,11 +47,17 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
 
     ir_targets_torchscript = any([ir == opt for opt in ["torchscript", "ts"]])
     ir_targets_fx = ir == "fx"
+    ir_targets_torch_compile = ir == "torch_compile"
+    ir_targets_fx_ts_compat_compile = ir == "fx_ts_compat_compile"
 
     if module_is_tsable and ir_targets_torchscript:
         return _IRType.ts
     elif module_is_fxable and ir_targets_fx:
         return _IRType.fx
+    elif module_is_fxable and ir_targets_torch_compile:
+        return _IRType.torch_compile
+    elif module_is_fxable and ir_targets_fx_ts_compat_compile:
+        return _IRType.fx_ts_compat_compile
     else:
         if ir == "default":
             # Options are listed in order of preference
@@ -148,6 +156,10 @@ def compile(
             dynamic_batch=False,
             **kwargs,
         )
+    elif target_ir == _IRType.torch_compile:
+        return torch_tensorrt.dynamo.torch_compile(
+            module, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs
+        )
     else:
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
 

py/torch_tensorrt/dynamo/__init__.py

@@ -1,126 +1 @@
-import torch
-import logging
-import collections.abc
-import torch_tensorrt
-from functools import partial
-
-from typing import Any
-from torch_tensorrt import EngineCapability, Device
-from torch_tensorrt.fx.utils import LowerPrecision
-
-from torch_tensorrt.dynamo._settings import CompilationSettings
-from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
-from torch_tensorrt.dynamo.backends import tensorrt_backend
-from torch_tensorrt.dynamo._defaults import (
-    PRECISION,
-    DEBUG,
-    MAX_WORKSPACE_SIZE,
-    MAX_NUM_TRT_ENGINES,
-)
-
-
-logger = logging.getLogger(__name__)
-
-
-def compile(
-    gm: torch.nn.Module,
-    inputs: Any,
-    *,
-    device=Device._current_device(),
-    disable_tf32=False,
-    sparse_weights=False,
-    enabled_precisions=set(),
-    refit=False,
-    debug=DEBUG,
-    capability=EngineCapability.default,
-    num_avg_timing_iters=1,
-    workspace_size=MAX_WORKSPACE_SIZE,
-    dla_sram_size=1048576,
-    dla_local_dram_size=1073741824,
-    dla_global_dram_size=536870912,
-    calibrator=None,
-    truncate_long_and_double=False,
-    require_full_compilation=False,
-    min_block_size=3,
-    torch_executed_ops=[],
-    torch_executed_modules=[],
-    **kwargs,
-):
-
-    logger.warn(
-        "The Dynamo backend is an experimental feature, for which only the "
-        + "following arguments are supported: "
-        + "{enabled_precisions, debug, workspace_size, max_num_trt_engines}"
-    )
-
-    if not isinstance(inputs, collections.abc.Sequence):
-        inputs = [inputs]
-
-    inputs = prepare_inputs(inputs, prepare_device(device))
-
-    if (
-        torch.float16 in enabled_precisions
-        or torch_tensorrt.dtype.half in enabled_precisions
-    ):
-        lower_precision = LowerPrecision.FP16
-    elif (
-        torch.float32 in enabled_precisions
-        or torch_tensorrt.dtype.float in enabled_precisions
-    ):
-        lower_precision = LowerPrecision.FP32
-    elif len(enabled_precisions) == 0:
-        logger.info(f"No precision specified, defaulting to {PRECISION}")
-        lower_precision = PRECISION
-    else:
-        raise ValueError(
-            f"Precision {enabled_precisions} not supported in the Dynamo Path"
-        )
-
-    custom_backend = create_backend(
-        precision=lower_precision,
-        debug=debug,
-        workspace_size=workspace_size,
-        **kwargs,
-    )
-
-    model = torch.compile(gm, backend=custom_backend)
-
-    # Ensure compilation occurs by calling the function with provided inputs
-    model(*inputs)
-
-    return model
-
-
-from torch_tensorrt.fx.utils import LowerPrecision
-
-logger = logging.getLogger(__name__)
-
-
-def create_backend(
-    precision: LowerPrecision = PRECISION,
-    debug: bool = DEBUG,
-    workspace_size: int = MAX_WORKSPACE_SIZE,
-    max_num_trt_engines: int = MAX_NUM_TRT_ENGINES,
-    **kwargs,
-):
-    """Create torch.compile backend given specified arguments
-
-    Args:
-        precision:
-        debug: Whether to print out verbose debugging information
-        workspace_size: Maximum workspace TRT is allowed to use for the module
-        precision: Model Layer precision
-    Returns:
-        Backend for torch.compile
-    """
-    settings = CompilationSettings(
-        debug=debug,
-        precision=precision,
-        workspace_size=workspace_size,
-        max_num_trt_engines=max_num_trt_engines,
-    )
-
-    return partial(
-        tensorrt_backend,
-        settings=settings,
-    )
+from .torch_compile import compile as torch_compile

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -0,0 +1,18 @@
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--ir",
+        metavar="Internal Representation",
+        nargs=1,
+        type=str,
+        required=True,
+        help="IR to compile with",
+        choices=["torch_compile", "fx_ts_compat_compile"],
+    )
+
+
+@pytest.fixture
+def ir(request):
+    return request.config.getoption("--ir")[0]