Fix noncontiguous AOT inputs

py/torch_tensorrt/_compile.py

@@ -337,8 +337,12 @@ def _fx_input_interface(
         elif not isinstance(arg_inputs, collections.abc.Sequence):
             arg_inputs = [arg_inputs]
 
-        torchtrt_arg_inputs = prepare_inputs(arg_inputs)
-        torchtrt_kwarg_inputs = prepare_inputs(kwarg_inputs)
+        torchtrt_arg_inputs = prepare_inputs(
+            arg_inputs, disable_memory_format_check=True
+        )
+        torchtrt_kwarg_inputs = prepare_inputs(
+            kwarg_inputs, disable_memory_format_check=True
+        )
 
         if module_type == _ModuleType.ep:
             exp_program = module
@@ -435,8 +439,10 @@ def cross_compile_for_windows(
         arg_inputs = [arg_inputs]  # type: ignore
 
     # Export the module
-    torchtrt_arg_inputs = prepare_inputs(arg_inputs)
-    torchtrt_kwarg_inputs = prepare_inputs(kwarg_inputs)
+    torchtrt_arg_inputs = prepare_inputs(arg_inputs, disable_memory_format_check=True)
+    torchtrt_kwarg_inputs = prepare_inputs(
+        kwarg_inputs, disable_memory_format_check=True
+    )
 
     exp_program = dynamo_trace(
         module, torchtrt_arg_inputs, kwarg_inputs=torchtrt_kwarg_inputs, **kwargs
@@ -561,10 +567,14 @@ def convert_method_to_trt_engine(
             normalized_arg_inputs = [arg_inputs]
 
         # Export the module
-        torchtrt_arg_inputs = prepare_inputs(normalized_arg_inputs)
-        torchtrt_kwarg_inputs = prepare_inputs(kwarg_inputs)
+        torchtrt_arg_inputs = prepare_inputs(
+            normalized_arg_inputs, disable_memory_format_check=True
+        )
+        torchtrt_kwarg_inputs = prepare_inputs(
+            kwarg_inputs, disable_memory_format_check=True
+        )
 
-        exp_program = torch_tensorrt.dynamo.trace(
+        exp_program = dynamo_trace(
             module, torchtrt_arg_inputs, kwarg_inputs=torchtrt_kwarg_inputs, **kwargs
         )
 

py/torch_tensorrt/dynamo/_compiler.py

@@ -319,8 +319,12 @@ def cross_compile_for_windows(
         arg_inputs = [arg_inputs]  # type: ignore
 
     # Prepare torch_trt inputs
-    trt_arg_inputs: Sequence[Input] = prepare_inputs(arg_inputs)
-    trt_kwarg_inputs: Optional[dict[Any, Any]] = prepare_inputs(kwarg_inputs)
+    trt_arg_inputs: Sequence[Input] = prepare_inputs(
+        arg_inputs, disable_memory_format_check=True
+    )
+    trt_kwarg_inputs: Optional[dict[Any, Any]] = prepare_inputs(
+        kwarg_inputs, disable_memory_format_check=True
+    )
     device = to_torch_tensorrt_device(device)
 
     compilation_options = {
@@ -719,8 +723,12 @@ def compile(
         arg_inputs = [arg_inputs]  # type: ignore
 
     # Prepare torch_trt inputs
-    trt_arg_inputs: Sequence[Input] = prepare_inputs(arg_inputs)
-    trt_kwarg_inputs: Optional[dict[Any, Any]] = prepare_inputs(kwarg_inputs)
+    trt_arg_inputs: Sequence[Input] = prepare_inputs(
+        arg_inputs, disable_memory_format_check=True
+    )
+    trt_kwarg_inputs: Optional[dict[Any, Any]] = prepare_inputs(
+        kwarg_inputs, disable_memory_format_check=True
+    )
     device = to_torch_tensorrt_device(device)
 
     engine_cache = None
@@ -1951,8 +1959,12 @@ def convert_exported_program_to_serialized_trt_engine(
         arg_inputs = [arg_inputs]  # type: ignore
 
     # Prepare torch_trt inputs
-    trt_arg_inputs: Sequence[Input] = prepare_inputs(arg_inputs)
-    trt_kwarg_inputs: Optional[dict[str, Any]] = prepare_inputs(kwarg_inputs)
+    trt_arg_inputs: Sequence[Input] = prepare_inputs(
+        arg_inputs, disable_memory_format_check=True
+    )
+    trt_kwarg_inputs: Optional[dict[str, Any]] = prepare_inputs(
+        kwarg_inputs, disable_memory_format_check=True
+    )
     device = to_torch_tensorrt_device(device)
 
     engine_cache = None

tests/py/dynamo/runtime/test_000_compiler_utils.py

@@ -138,6 +138,61 @@ def test_prepare_scalar_inputs(self):
         bool_result = prepare_inputs(True)
         self.assertIsInstance(bool_result, torch_tensorrt.Input)
 
+    def test_prepare_noncontiguous_requires_flag(self):
+        x = torch.randn(1, 4, 8).transpose(1, 2)
+        self.assertFalse(x.is_contiguous())
+        with self.assertRaises(ValueError):
+            prepare_inputs([x])
+        prepared = prepare_inputs([x], disable_memory_format_check=True)
+        self.assertEqual(tuple(prepared[0].shape), tuple(x.shape))
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available")
+class TestNoncontiguousArgInputs(unittest.TestCase):
+    def test_dynamo_compile_accepts_transposed_view(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        x = torch.randn(1, 4, 8, device="cuda").transpose(1, 2)
+        self.assertFalse(x.is_contiguous())
+        ep = torch.export.export(M(), (x,))
+        trt_mod = torch_tensorrt.dynamo.compile(
+            ep, arg_inputs=[x], min_block_size=1, use_python_runtime=True
+        )
+        with torch.no_grad():
+            torch.testing.assert_close(trt_mod(x), M().cuda()(x), rtol=1e-2, atol=1e-2)
+
+    def test_top_level_compile_accepts_transposed_view(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        x = torch.randn(1, 4, 8, device="cuda").transpose(1, 2)
+        self.assertFalse(x.is_contiguous())
+        trt_mod = torch_tensorrt.compile(
+            M().eval().cuda(),
+            ir="dynamo",
+            arg_inputs=[x],
+            min_block_size=1,
+            use_python_runtime=True,
+        )
+        with torch.no_grad():
+            torch.testing.assert_close(trt_mod(x), M().cuda()(x), rtol=1e-2, atol=1e-2)
+
+    def test_convert_method_to_trt_engine_accepts_transposed_view(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        x = torch.randn(1, 4, 8, device="cuda").transpose(1, 2)
+        self.assertFalse(x.is_contiguous())
+        engine = torch_tensorrt.convert_method_to_trt_engine(
+            M().eval().cuda(), "forward", arg_inputs=[x], ir="dynamo"
+        )
+        self.assertIsInstance(engine, bytes)
+        self.assertGreater(len(engine), 0)
+
 
 if __name__ == "__main__":
     unittest.main()