diff --git a/core/runtime/execute_engine.cpp b/core/runtime/execute_engine.cpp index 945f51b855..4ab1e3a973 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,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 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 +262,24 @@ 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 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 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()