fix: accept host-resident shape-tensor inputs without a device round trip

core/runtime/execute_engine.cpp

@@ -110,8 +110,11 @@ void setup_input_tensors(
     const auto& binding = compiled_engine->input_binding_infos[i];
     const auto& name = binding.name;
 
+    // A shape-tensor input is read from host memory (see the shape-tensor branch below), so a
+    // host tensor is accepted for it. Every other input is expected to be on the device.
     TORCHTRT_CHECK(
-        inputs[i].is_cuda(), "Expected input tensors to have device cuda, found device " << inputs[i].device());
+        binding.is_shape_tensor || inputs[i].is_cuda(),
+        "Expected input tensors to have device cuda, found device " << inputs[i].device());
 
     TORCHTRT_CHECK(
         inputs[i].dtype() == binding.expected_type,
@@ -122,10 +125,15 @@ void setup_input_tensors(
     LOG_DEBUG("Input Name: " << name << " Shape: " << dims << " isShapeInferenceIO: " << binding.is_shape_tensor);
 
     if (binding.is_shape_tensor) {
-      // Shape tensor inputs are casted to int64 explicitly.
-      // Refer to
+      // TensorRT reads a shape tensor's values from host memory, so they are copied to the host
+      // and cast to int64. Refer to
       // https://github.com/NVIDIA/TensorRT/blob/d2f4ef789a9a6ffdf37b55c3f81b486225f6b380/samples/common/sampleInference.cpp#L435
-      auto input_cpu = inputs[i].clone().contiguous().cpu().to(torch::kInt64);
+      // A device tensor is copied back with .cpu(), which synchronizes the stream; a host tensor
+      // is used as is, which does not, so passing shape inputs on the host avoids a per-call sync.
+      // The values are deep-copied into active_shape_tensor_values below, so input_cpu is transient
+      // and no clone is needed.
+      auto input_cpu = inputs[i].is_cuda() ? inputs[i].contiguous().cpu().to(torch::kInt64)
+                                           : inputs[i].contiguous().to(torch::kInt64);
       std::vector<int64_t> inputs_cpu_vec(
           input_cpu.data_ptr<int64_t>(), input_cpu.data_ptr<int64_t>() + input_cpu.numel());
       compiled_engine->active_shape_tensor_values.emplace_back(std::move(inputs_cpu_vec));
@@ -254,15 +262,24 @@ void create_output_allocator(c10::intrusive_ptr<TRTEngine> compiled_engine) {
 }
 
 std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intrusive_ptr<TRTEngine> compiled_engine) {
-  // All inputs are expected to be on CUDA. Warn and move any that are not.
-  for (auto& inp : inputs) {
-    if (inp.defined() && !inp.is_cuda()) {
-      LOG_WARNING(
-          "Input tensor is not on a CUDA device. Moving it to CUDA automatically. "
-          "For best performance, ensure all inputs are on the correct CUDA device before "
-          "calling the TensorRT engine (e.g. tensor.cuda() or tensor.to(device)).");
-      inp = inp.cuda();
+  // All inputs are expected to be on CUDA. Warn and move any that are not. A host-resident
+  // shape-tensor input is an exception: TensorRT reads it from host memory, so it is left where it
+  // is rather than moved to the device only for setup_input_tensors to copy it back. A shape tensor
+  // on any other non-CUDA device still falls through to the move below.
+  for (size_t i = 0; i < inputs.size(); i++) {
+    auto& inp = inputs[i];
+    if (!inp.defined() || inp.is_cuda()) {
+      continue;
     }
+    if (inp.is_cpu() && i < compiled_engine->input_binding_infos.size() &&
+        compiled_engine->input_binding_infos[i].is_shape_tensor) {
+      continue;
+    }
+    LOG_WARNING(
+        "Input tensor is not on a CUDA device. Moving it to CUDA automatically. "
+        "For best performance, ensure all inputs are on the correct CUDA device before "
+        "calling the TensorRT engine (e.g. tensor.cuda() or tensor.to(device)).");
+    inp = inp.cuda();
   }
 
 #ifdef ENABLE_TRT_NCCL_COLLECTIVES

tests/py/dynamo/runtime/test_host_shape_inputs.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+import torch_tensorrt
+from torch.testing._internal.common_utils import TestCase, run_tests
+from torch_tensorrt.dynamo.runtime import TorchTensorRTModule
+
+
+class ShapeInputModel(nn.Module):
+    """``arange(0, n, 1)`` makes ``n`` a TensorRT shape-tensor input binding."""
+
+    def forward(self, x: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
+        a = torch.arange(0, n, 1, device=x.device).to(torch.float32)
+        return (x * 2.0 + a).sum(dim=0, keepdim=True)
+
+
+class TestHostShapeInput(TestCase):
+    """A shape-tensor input is read from host memory, so the C++ runtime accepts it on the
+    host without a device round trip. The arange converter tests only exercise a CUDA shape
+    input; these cover the host path."""
+
+    def _compile(self):
+        model = ShapeInputModel().eval().cuda()
+        x = torch.randn(64, device="cuda")
+        n = torch.tensor(64, dtype=torch.int64, device="cuda")
+        dim = torch.export.Dim("d", min=2, max=256)
+        ep = torch.export.export(
+            model, (x, n), dynamic_shapes={"x": {0: dim}, "n": None}
+        )
+        gm = torch_tensorrt.dynamo.compile(
+            ep,
+            inputs=[x, n],
+            min_block_size=1,
+            use_python_runtime=False,
+            pass_through_build_failures=True,
+            assume_dynamic_shape_support=True,
+        )
+        trt_mod = next(
+            m for _, m in gm.named_children() if isinstance(m, TorchTensorRTModule)
+        )
+        self.assertEqual(
+            set(trt_mod.input_binding_names),
+            {"x", "_local_scalar_dense"},
+            "unexpected engine input bindings; the shape-input mapping below needs updating",
+        )
+        return model, trt_mod
+
+    def _run_engine(self, trt_mod, x, shape_val, shape_device):
+        args = {
+            "x": x,
+            "_local_scalar_dense": torch.tensor(
+                shape_val, dtype=torch.int64, device=shape_device
+            ),
+        }
+        ordered = [args[name] for name in trt_mod.input_binding_names]
+        return torch.ops.tensorrt.execute_engine(ordered, trt_mod.engine)[0]
+
+    def test_host_shape_input_matches_eager(self):
+        model, trt_mod = self._compile()
+        for shape_val in (8, 64, 200):
+            x = torch.randn(shape_val, device="cuda")
+            eager = model(x, torch.tensor(shape_val, device="cuda"))
+            host_out = self._run_engine(trt_mod, x, shape_val, "cpu")
+            torch.testing.assert_close(host_out, eager, rtol=1e-3, atol=1e-3)
+
+    def test_host_and_device_shape_input_agree(self):
+        model, trt_mod = self._compile()
+        x = torch.randn(64, device="cuda")
+        host_out = self._run_engine(trt_mod, x, 64, "cpu")
+        device_out = self._run_engine(trt_mod, x, 64, "cuda")
+        torch.testing.assert_close(host_out, device_out, rtol=1e-5, atol=1e-5)
+
+
+if __name__ == "__main__":
+    run_tests()