Merge remote-tracking branch 'origin/dynamo_tensor_freeze' into export_prototype

Author: peri044

core/runtime/execute_engine.cpp

@@ -136,7 +136,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       TORCHTRT_CHECK(
           inputs[i].dtype() == expected_type,
           "Expected input tensors to have type " << expected_type << ", found type " << inputs[i].dtype());
-      auto dims = core::util::toDimsPad(inputs[i].sizes(), 1);
+      auto dims = core::util::toDims(inputs[i].sizes());
       auto shape = core::util::toVec(dims);
       LOG_DEBUG("Input Name: " << name << " Shape: " << dims);
       compiled_engine->exec_ctx->setInputShape(name.c_str(), dims);

docker/Dockerfile

@@ -106,8 +106,17 @@ COPY --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/dist/ .
 
 RUN cp /opt/torch_tensorrt/docker/WORKSPACE.docker /opt/torch_tensorrt/WORKSPACE
 RUN pip install -r /opt/torch_tensorrt/py/requirements.txt
+# Install all dependency wheel files and user-specified TensorRT
+RUN pip install *.whl
 RUN pip install tensorrt==${TENSORRT_VERSION}.*
-RUN pip install *.whl && rm -fr /workspace/torch_tensorrt/dist/* *.whl
+
+# Add the Torch-TensorRT wheel file to the dist directory and delete all other .whl files
+RUN rm -fr /workspace/torch_tensorrt/dist/*
+RUN mkdir -p /opt/torch_tensorrt/dist/ && mv torch_tensorrt*.whl /opt/torch_tensorrt/dist/
+RUN rm -fr *.whl
+
+# Remove other cache files if present
+RUN pip cache purge && rm -rf /opt/torch_tensorrt/.mypy_cache
 
 WORKDIR /opt/torch_tensorrt
 

py/torch_tensorrt/_Device.py

@@ -8,11 +8,10 @@
 
 import warnings
 
-import torch
-from torch_tensorrt import logging
-
 # from torch_tensorrt import _enums
 import tensorrt as trt
+import torch
+from torch_tensorrt import logging
 
 try:
     from torch_tensorrt import _C
@@ -120,6 +119,9 @@ def __str__(self) -> str:
             )
         )
 
+    def __repr__(self) -> str:
+        return self.__str__()
+
     def _to_internal(self) -> _C.Device:
         internal_dev = _C.Device()
         if self.device_type == trt.DeviceType.GPU:

py/torch_tensorrt/_Input.py

@@ -339,18 +339,18 @@ def from_tensor(
             A Input object.
         """
         if not (
-            t.is_contiguous(memory_format=torch.contiguous_format)
+            disable_memory_format_check
+            or t.is_contiguous(memory_format=torch.contiguous_format)
             or t.is_contiguous(memory_format=torch.channels_last)
-            or disable_memory_format_check
         ):
             raise ValueError(
                 "Tensor does not have a supported memory format, supported formats are contiguous or channel_last"
             )
         frmt = (
             torch.contiguous_format
             if (
-                t.is_contiguous(memory_format=torch.contiguous_format)
-                or disable_memory_format_check
+                disable_memory_format_check
+                or t.is_contiguous(memory_format=torch.contiguous_format)
             )
             else torch.channels_last
         )

py/torch_tensorrt/_compile.py

@@ -209,12 +209,12 @@ def compile(
         import collections.abc
 
         from torch_tensorrt import Device
-        from torch_tensorrt.dynamo.utils import prepare_device, prepare_inputs
+        from torch_tensorrt.dynamo.utils import prepare_inputs, to_torch_device
 
         if not isinstance(inputs, collections.abc.Sequence):
             inputs = [inputs]
         device = kwargs.get("device", Device._current_device())
-        torchtrt_inputs, torch_inputs = prepare_inputs(inputs, prepare_device(device))
+        torchtrt_inputs, torch_inputs = prepare_inputs(inputs, to_torch_device(device))
         module = torch_tensorrt.dynamo.trace(module, torch_inputs, **kwargs)
         compiled_aten_module: torch.fx.GraphModule = dynamo_compile(
             module,
@@ -239,7 +239,10 @@ def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Any:
     """
     from torch_tensorrt.dynamo.backend import torch_tensorrt_backend
 
-    boxed_fn = torch.compile(module, backend=torch_tensorrt_backend, options={**kwargs})
+    # TODO: Remove dynamic=False when SymInt Dynamic shape support is ready
+    boxed_fn = torch.compile(
+        module, backend=torch_tensorrt_backend, dynamic=False, options={**kwargs}
+    )
 
     return boxed_fn
 

py/torch_tensorrt/dynamo/_defaults.py

@@ -1,7 +1,9 @@
 import torch
+from torch_tensorrt._Device import Device
 
 PRECISION = torch.float32
 DEBUG = False
+DEVICE = None
 WORKSPACE_SIZE = 0
 MIN_BLOCK_SIZE = 5
 PASS_THROUGH_BUILD_FAILURES = False
@@ -12,3 +14,7 @@
 USE_PYTHON_RUNTIME = False
 USE_FAST_PARTITIONER = True
 ENABLE_EXPERIMENTAL_DECOMPOSITIONS = False
+
+
+def default_device() -> Device:
+    return Device(gpu_id=torch.cuda.current_device())

py/torch_tensorrt/dynamo/_settings.py

@@ -2,6 +2,7 @@
 from typing import Optional, Set
 
 import torch
+from torch_tensorrt._Device import Device
 from torch_tensorrt.dynamo._defaults import (
     DEBUG,
     ENABLE_EXPERIMENTAL_DECOMPOSITIONS,
@@ -15,6 +16,7 @@
     USE_PYTHON_RUNTIME,
     VERSION_COMPATIBLE,
     WORKSPACE_SIZE,
+    default_device,
 )
 
 
@@ -54,3 +56,4 @@ class CompilationSettings:
     truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE
     use_fast_partitioner: bool = USE_FAST_PARTITIONER
     enable_experimental_decompositions: bool = ENABLE_EXPERIMENTAL_DECOMPOSITIONS
+    device: Device = field(default_factory=default_device)

py/torch_tensorrt/dynamo/compile.py

@@ -2,7 +2,7 @@
 
 import collections.abc
 import logging
-from typing import Any, List, Optional, Sequence, Set, Tuple
+from typing import Any, List, Optional, Sequence, Set, Tuple, Union
 
 import torch
 import torch_tensorrt
@@ -13,6 +13,7 @@
 from torch_tensorrt.dynamo import CompilationSettings, partitioning
 from torch_tensorrt.dynamo._defaults import (
     DEBUG,
+    DEVICE,
     ENABLE_EXPERIMENTAL_DECOMPOSITIONS,
     MAX_AUX_STREAMS,
     MIN_BLOCK_SIZE,
@@ -29,7 +30,11 @@
     convert_module,
     repair_long_or_double_inputs,
 )
-from torch_tensorrt.dynamo.utils import prepare_device, prepare_inputs
+from torch_tensorrt.dynamo.utils import (
+    prepare_inputs,
+    to_torch_device,
+    to_torch_tensorrt_device,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +43,7 @@ def compile(
     gm: Any,
     inputs: Any,
     *,
-    device: Device = Device._current_device(),
+    device: Optional[Union[Device, torch.device, str]] = DEVICE,
     disable_tf32: bool = False,
     sparse_weights: bool = False,
     enabled_precisions: Set[torch.dtype] | Tuple[torch.dtype] = (torch.float32,),
@@ -82,7 +87,9 @@ def compile(
     if not isinstance(inputs, collections.abc.Sequence):
         inputs = [inputs]
 
-    _, torch_inputs = prepare_inputs(inputs, prepare_device(device))
+    device = to_torch_tensorrt_device(device)
+
+    _, torch_inputs = prepare_inputs(inputs, to_torch_device(device))
 
     if (
         torch.float16 in enabled_precisions
@@ -105,6 +112,7 @@ def compile(
     compilation_options = {
         "precision": precision,
         "debug": debug,
+        "device": device,
         "workspace_size": workspace_size,
         "min_block_size": min_block_size,
         "torch_executed_ops": torch_executed_ops