From 632e8dcc8012d2dc3b83852288be3da556329b7a Mon Sep 17 00:00:00 2001 From: Kshitij Srivastava Date: Tue, 15 Sep 2026 16:59:33 +0000 Subject: [PATCH 1/2] fix: accept host-resident shape-tensor inputs without a device round trip A TensorRT shape-tensor input is read from host memory, but setup_input_tensors required every input on the device and copied a shape input back with .cpu(), and execute_engine moved any host input to the device first. A host shape value was therefore pushed device-side and pulled straight back, and the .cpu() copy synchronizes the stream. Accept a host-resident shape-tensor input as is; device inputs and all non-shape inputs are unchanged. --- core/runtime/execute_engine.cpp | 37 ++++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/core/runtime/execute_engine.cpp b/core/runtime/execute_engine.cpp index 945f51b855..ce41a63fd9 100644 --- a/core/runtime/execute_engine.cpp +++ b/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,13 @@ 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. + auto input_cpu = inputs[i].is_cuda() ? inputs[i].clone().contiguous().cpu().to(torch::kInt64) + : inputs[i].contiguous().to(torch::kInt64); std::vector inputs_cpu_vec( input_cpu.data_ptr(), input_cpu.data_ptr() + input_cpu.numel()); compiled_engine->active_shape_tensor_values.emplace_back(std::move(inputs_cpu_vec)); @@ -254,15 +260,22 @@ void create_output_allocator(c10::intrusive_ptr compiled_engine) { } std::vector execute_engine(std::vector inputs, c10::intrusive_ptr 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 shape-tensor input is + // an exception: TensorRT reads it from host memory, so a host tensor is left where it is rather + // than moved to the device only for setup_input_tensors to copy it back. + for (size_t i = 0; i < inputs.size(); i++) { + auto& inp = inputs[i]; + if (!inp.defined() || inp.is_cuda()) { + continue; } + if (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 From b7b611fc5db88d2a96df96b7dc7a37c0b4a264e2 Mon Sep 17 00:00:00 2001 From: Kshitij Srivastava Date: Fri, 18 Sep 2026 18:35:33 +0000 Subject: [PATCH 2/2] address review: drop redundant clone, guard preamble, add host shape-input test - setup_input_tensors: remove clone() on the device arm; .cpu() already returns a fresh host tensor and values are deep-copied into active_shape_tensor_values. - execute_engine preamble: only leave a shape-tensor input in place when it is on CPU; a shape tensor on any other non-CUDA device falls through to the CUDA move. - add tests/py/dynamo/runtime/test_host_shape_inputs.py covering the host shape input path (arange converter tests only exercise a CUDA shape input). --- core/runtime/execute_engine.cpp | 14 ++-- .../dynamo/runtime/test_host_shape_inputs.py | 74 +++++++++++++++++++ 2 files changed, 83 insertions(+), 5 deletions(-) create mode 100644 tests/py/dynamo/runtime/test_host_shape_inputs.py diff --git a/core/runtime/execute_engine.cpp b/core/runtime/execute_engine.cpp index ce41a63fd9..4ab1e3a973 100644 --- a/core/runtime/execute_engine.cpp +++ b/core/runtime/execute_engine.cpp @@ -130,7 +130,9 @@ void setup_input_tensors( // https://github.com/NVIDIA/TensorRT/blob/d2f4ef789a9a6ffdf37b55c3f81b486225f6b380/samples/common/sampleInference.cpp#L435 // 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. - auto input_cpu = inputs[i].is_cuda() ? inputs[i].clone().contiguous().cpu().to(torch::kInt64) + // 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 inputs_cpu_vec( input_cpu.data_ptr(), input_cpu.data_ptr() + input_cpu.numel()); @@ -260,15 +262,17 @@ void create_output_allocator(c10::intrusive_ptr compiled_engine) { } std::vector execute_engine(std::vector inputs, c10::intrusive_ptr compiled_engine) { - // All inputs are expected to be on CUDA. Warn and move any that are not. A shape-tensor input is - // an exception: TensorRT reads it from host memory, so a host tensor is left where it is rather - // than moved to the device only for setup_input_tensors to copy it back. + // 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 (i < compiled_engine->input_binding_infos.size() && compiled_engine->input_binding_infos[i].is_shape_tensor) { + if (inp.is_cpu() && i < compiled_engine->input_binding_infos.size() && + compiled_engine->input_binding_infos[i].is_shape_tensor) { continue; } LOG_WARNING( diff --git a/tests/py/dynamo/runtime/test_host_shape_inputs.py b/tests/py/dynamo/runtime/test_host_shape_inputs.py new file mode 100644 index 0000000000..f607ebba86 --- /dev/null +++ b/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()