adding test case for multiple outputs and moving op to register_meta_ops

Author: apbose

py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py

@@ -150,3 +150,24 @@ def __setstate__(self, serialized_state: List[str]) -> Any:
 
     def __getstate__(self) -> Any:
         pass
+
+
+@torch.library.custom_op(
+    "tensorrt::no_op_placeholder_for_execute_engine", mutates_args=()
+)
+def no_op_placeholder_for_execute_engine(
+    inputs: List[torch.Tensor],
+    abi_version: str,
+    name: str,
+    serialized_device_info: str,
+    serialized_engine: str,
+    serialized_in_binding_names: str,
+    serialized_out_binding_names: str,
+    serialized_hardware_compatible: str,
+    serialized_metadata: str,
+    serialized_target_platform: str,
+    serialized_require_output_allocator: str,
+) -> List[torch.Tensor]:
+    raise RuntimeError(
+        "The saved model is cross compiled for windows in Linux, should only be loadded in Windows via torch_tensorrt.load_cross_compiled_exported_program() api."
+    )

py/torch_tensorrt/runtime/_utils.py

@@ -128,24 +128,3 @@ def _get_most_compatible_device(
                 best_match = candidate
 
     return best_match
-
-
-@torch.library.custom_op(
-    "tensorrt::no_op_placeholder_for_execute_engine", mutates_args=()
-)
-def no_op_placeholder_for_execute_engine(
-    inputs: List[torch.Tensor],
-    abi_version: str,
-    name: str,
-    serialized_device_info: str,
-    serialized_engine: str,
-    serialized_in_binding_names: str,
-    serialized_out_binding_names: str,
-    serialized_hardware_compatible: str,
-    serialized_metadata: str,
-    serialized_target_platform: str,
-    serialized_require_output_allocator: str,
-) -> List[torch.Tensor]:
-    raise RuntimeError(
-        "The saved model is cross compiled for windows in Linux, should only be loadded in Windows via torch_tensorrt.load_cross_compiled_exported_program() api."
-    )

tests/py/dynamo/runtime/test_003_cross_compile_for_windows.py

@@ -22,6 +22,8 @@ class Add(torch.nn.Module):
             def forward(self, a, b):
                 return torch.add(a, b)
 
+        print("here")
+
         model = Add().eval().cuda()
         inputs = [torch.randn(2, 3).cuda(), torch.randn(2, 3).cuda()]
         trt_ep_path = os.path.join(tempfile.gettempdir(), "trt.ep")
@@ -63,3 +65,31 @@ def forward(self, a, b):
             )
         except Exception as e:
             pytest.fail(f"unexpected exception raised: {e}")
+
+    @unittest.skipIf(
+        platform.system() != "Linux" or platform.architecture()[0] != "64bit",
+        "Cross compile for windows can only be enabled on linux x86-64 platform",
+    )
+    @pytest.mark.unit
+    def test_dynamo_cross_compile_for_windows_multiple_output(self):
+        class Add(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.add(a, b), torch.add(a, b)
+
+        model = Add().eval().cuda()
+        inputs = (torch.randn(2, 3).cuda(), torch.randn(2, 3).cuda())
+        trt_ep_path = os.path.join(tempfile.gettempdir(), "trt.ep")
+        exp_program = torch.export.export(model, inputs)
+        compile_spec = {
+            "inputs": inputs,
+            "min_block_size": 1,
+        }
+        try:
+            trt_gm = torch_tensorrt.dynamo.cross_compile_for_windows(
+                exp_program, **compile_spec
+            )
+            torch_tensorrt.dynamo.save_cross_compiled_exported_program(
+                trt_gm, file_path=trt_ep_path
+            )
+        except Exception as e:
+            pytest.fail(f"unexpected exception raised: {e}")