🐛 [Bug] Module using slice-based ops fails at inference time when compiled with dynamo and dynamic batch size

Bug Description

Module triggers a runtime error at inference time when compiled with dynamo and dynamic batch sizes. Error appears to be related to using slice-based assignment (minimum example included). This appears to cause TRT (maybe originating from Myelin?) to generate a handful of CPU tensors internally. Dynamo runtime module does not handle this edge-case.

Runtime Error

  File "/opt/venv/lib/python3.10/site-packages/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py", line 204, in forward
    outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine(
  File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1024, in __call__
    return self_._op(*args, **(kwargs or {}))

RuntimeError: [Error thrown at core/runtime/execute_engine.cpp:190] Expected inputs[i].is_cuda() to be true but got false
Expected input tensors to have device cuda, found device cpu

Workaround

Trivial workaround is to cast inputs to cuda and incur the copy cost at inference time.
https://github.com/pytorch/TensorRT/blob/abed8f06f057d5ec6652049d2d5770bb5ac6a4ce/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py#L199-L202
Modified to:

input_tensors: List[torch.Tensor] = [
           (i if isinstance(i, torch.Tensor) else torch.tensor(i)).cuda()
           for i in inputs
       ]

Or more concisely:

input_tensors: List[torch.Tensor] = [
    torch.as_tensor(i).cuda() for i in inputs
]

To Reproduce

Minimal reproducible example provided below.

import torch
import torch.nn as nn
import torch_tensorrt

class Network(nn.Module):
    def __init__(self):
        super().__init__()

        self.h = nn.Sequential(
            nn.Conv2d(1, 4, 3, padding='same'),
            nn.ReLU()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return decode(
            self.h(x).permute(0, 2, 3, 1)
        )

def decode(x: torch.Tensor) -> torch.Tensor:
    return x[..., :2].exp_()

def main():
    x = torch.randn(4, 1, 16, 16).cuda()

    net = Network().eval().cuda()
    spec = torch_tensorrt.Input(min_shape=[1, 1, 16, 16],
                                opt_shape=[4, 1, 16, 16],
                                max_shape=[8, 1, 16, 16],
                                dtype=torch.float)
    with torch_tensorrt.logging.debug():
        net_ts = torch_tensorrt.compile(net, ir="dynamo", inputs=spec, enabled_precisions={torch.float}, min_block_size=3)
    net_ts.print_readable()
    net_ts(x)

if __name__ == "__main__":
    with torch.no_grad():
        main()

Readable printout of torch.fx.GraphModule:

Note that full and full_1 are created on cpu.

class GraphModule(torch.nn.Module):
    def forward(self, x):
        x: "f32[s0, 1, 16, 16]"; 
    
        x, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
        # No stacktrace found for following nodes
        _run_on_acc_0 = self._run_on_acc_0(x);  x = None
        getitem = _run_on_acc_0[0]
        getitem_1 = _run_on_acc_0[1]
        getitem_2 = _run_on_acc_0[2];  _run_on_acc_0 = None
        _run_on_gpu_1 = self._run_on_gpu_1(getitem);  getitem = None
        getitem_3 = _run_on_gpu_1[0]
        getitem_4 = _run_on_gpu_1[1];  _run_on_gpu_1 = None
        _run_on_acc_2 = self._run_on_acc_2(getitem_3, getitem_4);  getitem_3 = getitem_4 = None
        _run_on_gpu_3 = self._run_on_gpu_3(getitem_1, _run_on_acc_2, getitem_2);  getitem_1 = _run_on_acc_2 = getitem_2 = None
        _run_on_acc_4 = self._run_on_acc_4(_run_on_gpu_3);  _run_on_gpu_3 = None
        return pytree.tree_unflatten((_run_on_acc_4,), self._out_spec)
        
    class GraphModule(torch.nn.Module):
        def forward(self, sym_size_int: "Sym(s0)"):
            # File: /workspace/slice_error.py:20 in decode, code: return x[..., :2].exp_()
            full: "i64[s0, 16, 16]" = torch.ops.aten.full.default([sym_size_int, 16, 16], 1, dtype = torch.int64, layout = torch.strided, device = device(type='cpu'), pin_memory = False)
            full_1: "i64[s0, 16, 16]" = torch.ops.aten.full.default([sym_size_int, 16, 16], 1, dtype = torch.int64, layout = torch.strided, device = device(type='cpu'), pin_memory = False);  sym_size_int = None
            return (full, full_1)
            
    class GraphModule(torch.nn.Module):
        def forward(self, permute_1: "f32[s0, 16, 16, 4]", _to_copy: "i64[s0, 16, 16, 2]", exp: "f32[s0, 16, 16, 2]"):
            # File: /workspace/slice_error.py:20 in decode, code: return x[..., :2].exp_()
            scatter: "f32[s0, 16, 16, 4]" = torch.ops.aten.scatter.src(permute_1, 3, _to_copy, exp);  permute_1 = _to_copy = exp = None
            return scatter
            

Both full and full_1 seem related to the Myelin backend. (snippet from debug output)

Layer(Myelin): {ForeignNode[__/mul_1_rhs + __/mul_1_broadcast_rhs_broadcast...(Unnamed Layer* 10) [Cast]]}, Tactic: 0x0000000000000000, full (Int64[-1,16,16]), full_1 (Int64[-1,16,16]) -> output0 (Int64[-1,16,16,2])

Expected behavior

No runtime error.

Honestly I'm surprised the resulting GraphModule is so different between static vs dynamic batching.

Environment

Use Pytorch, CUDA, and TensorRT distributed in NGC container.

Image: nvcr.io/nvidia/pytorch:24.06-py3
Torch-TensorRT: Built from source from commit abed8f0
CPU: AMD Ryzen Threadripper 3rd gen
GPU: Nvidia RTX A6000
Python: 3.10
Build Command: I build with a minimal makefile / modified setup.py found here: TensorRT-Make.

Additional context

Common relevant use-case is in the decoding step of any number of object detection networks.