feat: Add support for device compilation setting

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

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/_compile.py

@@ -192,12 +192,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,

py/torch_tensorrt/dynamo/_defaults.py

@@ -1,4 +1,5 @@
 import torch
+from torch_tensorrt._Device import Device
 
 PRECISION = torch.float32
 DEBUG = False
@@ -11,3 +12,4 @@
 TRUNCATE_LONG_AND_DOUBLE = False
 USE_PYTHON_RUNTIME = False
 USE_FAST_PARTITIONER = True
+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,
@@ -31,3 +33,4 @@ class CompilationSettings:
     use_python_runtime: Optional[bool] = USE_PYTHON_RUNTIME
     truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE
     use_fast_partitioner: bool = USE_FAST_PARTITIONER
+    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
@@ -14,6 +14,7 @@
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo._defaults import (
     DEBUG,
+    DEVICE,
     MAX_AUX_STREAMS,
     MIN_BLOCK_SIZE,
     OPTIMIZATION_LEVEL,
@@ -30,7 +31,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,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -39,7 +44,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,),
@@ -80,7 +85,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
@@ -103,6 +110,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

tests/py/dynamo/backend/test_compiler_utils.py

@@ -1,26 +1,61 @@
-from torch_tensorrt.dynamo.utils import prepare_device, prepare_inputs
-from utils import same_output_format
-import torch_tensorrt
 import unittest
+
 import torch
+import torch_tensorrt
+from torch_tensorrt.dynamo.utils import (
+    prepare_inputs,
+    to_torch_device,
+    to_torch_tensorrt_device,
+)
+from utils import same_output_format
 
 
-class TestPrepareDevice(unittest.TestCase):
-    def test_prepare_cuda_device(self):
+class TestToTorchDevice(unittest.TestCase):
+    def test_cast_cuda_device(self):
         gpu_id = 0
         device = torch.device(f"cuda:{gpu_id}")
-        prepared_device = prepare_device(device)
+        prepared_device = to_torch_device(device)
         self.assertTrue(isinstance(prepared_device, torch.device))
         self.assertTrue(prepared_device.index == gpu_id)
 
-    def test_prepare_trt_device(self):
+    def test_cast_trt_device(self):
         gpu_id = 4
         device = torch_tensorrt.Device(gpu_id=gpu_id)
-        prepared_device = prepare_device(device)
+        prepared_device = to_torch_device(device)
+        self.assertTrue(isinstance(prepared_device, torch.device))
+        self.assertTrue(prepared_device.index == gpu_id)
+
+    def test_cast_str_device(self):
+        gpu_id = 2
+        device = f"cuda:{2}"
+        prepared_device = to_torch_device(device)
         self.assertTrue(isinstance(prepared_device, torch.device))
         self.assertTrue(prepared_device.index == gpu_id)
 
 
+class TestToTorchTRTDevice(unittest.TestCase):
+    def test_cast_cuda_device(self):
+        gpu_id = 0
+        device = torch.device(f"cuda:{gpu_id}")
+        prepared_device = to_torch_tensorrt_device(device)
+        self.assertTrue(isinstance(prepared_device, torch_tensorrt.Device))
+        self.assertTrue(prepared_device.gpu_id == gpu_id)
+
+    def test_cast_trt_device(self):
+        gpu_id = 4
+        device = torch_tensorrt.Device(gpu_id=gpu_id)
+        prepared_device = to_torch_tensorrt_device(device)
+        self.assertTrue(isinstance(prepared_device, torch_tensorrt.Device))
+        self.assertTrue(prepared_device.gpu_id == gpu_id)
+
+    def test_cast_str_device(self):
+        gpu_id = 2
+        device = f"cuda:{2}"
+        prepared_device = to_torch_tensorrt_device(device)
+        self.assertTrue(isinstance(prepared_device, torch_tensorrt.Device))
+        self.assertTrue(prepared_device.gpu_id == gpu_id)
+
+
 class TestPrepareInputs(unittest.TestCase):
     def test_prepare_single_tensor_input(self):
         inputs = [torch.ones((4, 4))]