feat: Add support for device compilation setting

- Add updated Device utilities and automatic context-aware device
detection for torch compile

Author: gs-olive

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/dynamo/_defaults.py

@@ -1,4 +1,5 @@
 import torch
+from torch_tensorrt._Device import Device
 
 PRECISION = torch.float32
 DEBUG = False
@@ -10,3 +11,4 @@
 OPTIMIZATION_LEVEL = None
 TRUNCATE_LONG_AND_DOUBLE = False
 USE_PYTHON_RUNTIME = False
+DEVICE = Device._current_device()

py/torch_tensorrt/dynamo/_settings.py

@@ -2,8 +2,10 @@
 from typing import Optional, Set
 
 import torch
+from torch_tensorrt._Device import Device
 from torch_tensorrt.dynamo._defaults import (
     DEBUG,
+    DEVICE,
     MAX_AUX_STREAMS,
     MIN_BLOCK_SIZE,
     OPTIMIZATION_LEVEL,
@@ -29,3 +31,4 @@ class CompilationSettings:
     optimization_level: Optional[int] = OPTIMIZATION_LEVEL
     use_python_runtime: Optional[bool] = USE_PYTHON_RUNTIME
     truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE
+    device: Device = DEVICE

py/torch_tensorrt/dynamo/compile.py

@@ -2,7 +2,7 @@
 
 import collections.abc
 import logging
-from typing import Any, List, Optional, Set, Tuple
+from typing import Any, List, Optional, Set, Tuple, Union
 
 import torch
 import torch_tensorrt
@@ -15,6 +15,7 @@
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo._defaults import (
     DEBUG,
+    DEVICE,
     MAX_AUX_STREAMS,
     MIN_BLOCK_SIZE,
     OPTIMIZATION_LEVEL,
@@ -31,7 +32,11 @@
     fuse_permute_linear,
     fuse_permute_matmul,
 )
-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,
+)
 from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
 
 logger = logging.getLogger(__name__)
@@ -41,7 +46,7 @@ def compile(
     gm: Any,
     inputs: Any,
     *,
-    device: Device = Device._current_device(),
+    device: Union[Device, torch.device, str] = DEVICE,
     disable_tf32: bool = False,
     sparse_weights: bool = False,
     enabled_precisions: Set[torch.dtype] | Tuple[torch.dtype] = (torch.float32,),
@@ -81,7 +86,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
@@ -104,6 +111,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

py/torch_tensorrt/dynamo/conversion/conversion.py

@@ -3,14 +3,13 @@
 import io
 from typing import Sequence
 
+import tensorrt as trt
 import torch
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.conversion import TRTInterpreter
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
 
-import tensorrt as trt
-
 
 def convert_module(
     module: torch.fx.GraphModule,
@@ -72,4 +71,5 @@ def convert_module(
             name=name,
             input_binding_names=list(interpreter_result.input_names),
             output_binding_names=list(interpreter_result.output_names),
+            target_device=settings.device,
         )

py/torch_tensorrt/dynamo/utils.py

@@ -2,7 +2,7 @@
 
 import logging
 from dataclasses import fields, replace
-from typing import Any, Callable, Dict, Optional, Sequence
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 import torch
 from torch_tensorrt._Device import Device
@@ -114,23 +114,37 @@ def prepare_inputs(
         )
 
 
-def prepare_device(device: Device | torch.device) -> torch.device:
-    _device: torch.device
+def to_torch_device(device: Union[Device, torch.device, str]) -> torch.device:
+    """Cast a device-type to torch.device
+
+    Returns the corresponding torch.device
+    """
     if isinstance(device, Device):
         if device.gpu_id != -1:
-            _device = torch.device(device.gpu_id)
+            return torch.device(device.gpu_id)
         else:
             raise ValueError("Invalid GPU ID provided for the CUDA device provided")
 
     elif isinstance(device, torch.device):
-        _device = device
+        return device
 
     else:
-        raise ValueError(
-            "Invalid device provided. Supported options: torch.device | torch_tensorrt.Device"
-        )
+        return torch.device(device)
 
-    return _device
+
+def to_torch_tensorrt_device(device: Union[Device, torch.device, str]) -> Device:
+    """Cast a device-type to torch_tensorrt.Device
+
+    Returns the corresponding torch_tensorrt.Device
+    """
+    if isinstance(device, Device):
+        return device
+
+    elif isinstance(device, torch.device):
+        return Device(gpu_id=device.index)
+
+    else:
+        return Device(device)
 
 
 def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings:
@@ -164,6 +178,19 @@ def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings:
     # Parse input runtime specification
     settings.use_python_runtime = use_python_runtime_parser(settings.use_python_runtime)
 
+    # Ensure device is a torch_tensorrt Device
+    settings.device = to_torch_tensorrt_device(settings.device)
+
+    # Check and update device settings
+    default_torch_gpu_idx = torch.cuda.default_stream().device.index
+    if "device" not in kwargs and default_torch_gpu_idx != settings.device.gpu_id:
+        logger.warning(
+            f"No device specified, detected differing gpu IDs for CUDA default: {settings.device.gpu_id} "
+            f"and Torch default: {default_torch_gpu_idx}. Using Torch default gpu ID: {default_torch_gpu_idx}. "
+            "If this is incorrect, please specify an input device, via the device keyword."
+        )
+        settings.device = Device(gpu_id=default_torch_gpu_idx)
+
     logger.debug(f"Compiling with Settings:\n{settings}")
 
     return settings