complex _reshape_copy / _to_copy

py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py

@@ -12,7 +12,7 @@
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
     clean_up_graph_after_modifications,
 )
-from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES
+from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES, COMPLEX_TO_REAL_DTYPE
 
 logger = logging.getLogger(__name__)
 
@@ -1282,6 +1282,27 @@ def _rewrite_scalar_tensor(self, node: Node) -> bool:
             self.gm.graph.erase_node(node)
             return True
 
+    @_complex_unpacker(torch.ops.aten._to_copy.default)
+    def _rewrite_to_copy(self, node: Node) -> bool:
+        # complex target: remap dtype, [..., 2] layout unchanged
+        # real target: the cast discards the imaginary part, so select re
+        kwargs = dict(node.kwargs)
+        dtype = kwargs.get("dtype")
+        inp = node.args[0]
+        to_real = dtype is not None and dtype not in COMPLEX_DTYPES
+        if dtype is not None and not to_real:
+            kwargs["dtype"] = COMPLEX_TO_REAL_DTYPE[dtype]
+        with SubgraphBuilder(self.gm.graph, node) as b:
+            if to_real:
+                inp = b(torch.ops.aten.select.int, inp, -1, 0)
+            out = b(torch.ops.aten._to_copy.default, inp)
+            out.kwargs = kwargs
+            if not to_real:
+                out.meta["is_complex_layout"] = True
+            node.replace_all_uses_with(out)
+            self.gm.graph.erase_node(node)
+            return True
+
     # ------------------------------------------------------------------
     # Shape-manipulation handlers
     #
@@ -1297,6 +1318,7 @@ def _rewrite_scalar_tensor(self, node: Node) -> bool:
         torch.ops.aten.reshape.default,
         torch.ops.aten.view.default,
         torch.ops.aten._unsafe_view.default,
+        torch.ops.aten._reshape_copy.default,
     )
     def _rewrite_reshape_view(self, node: Node) -> bool:
         # Append 2 to the target shape so the trailing real/imag dim is

tests/py/dynamo/lowering/test_complex_rewrite.py

@@ -594,6 +594,50 @@ def forward(self, z):
     _check_op(M(), (_z(),), "reshape")
 
 
+@pytest.mark.unit
+def test_reshape_copy():
+    class M(nn.Module):
+        def forward(self, z):
+            return torch.ops.aten._reshape_copy.default(z, [12])
+
+    gm = _export_and_lower(M(), (_z(),))
+    targets = {n.target for n in gm.graph.nodes if n.op == "call_function"}
+    assert torch.ops.aten.view_as_complex.default not in targets
+    assert torch.ops.aten.view_as_real.default not in targets
+    _check_op(M(), (_z(),), "reshape_copy")
+
+
+@pytest.mark.unit
+def test_to_copy_complex_dtype():
+    class M(nn.Module):
+        def forward(self, z):
+            return torch.ops.aten._to_copy.default(z, dtype=torch.complex64)
+
+    # complex64's real counterpart (float32) is TRT-convertible; float64 is not
+    z = torch.randn(3, 4, dtype=torch.complex128)
+    gm = _export_and_lower(M(), (z,))
+    targets = {n.target for n in gm.graph.nodes if n.op == "call_function"}
+    assert torch.ops.aten.view_as_complex.default not in targets
+    assert torch.ops.aten.view_as_real.default not in targets
+    assert any(
+        node.target == torch.ops.aten._to_copy.default
+        and node.kwargs.get("dtype") == torch.float32
+        for node in gm.graph.nodes
+    ), "a complex target must be remapped to its real counterpart"
+    _check_op(M(), (z,), "to_copy_complex_dtype")
+
+
+@pytest.mark.unit
+def test_to_copy_complex_to_real():
+    """z.to(float) discards the imaginary part and the trailing real/imag dim."""
+
+    class M(nn.Module):
+        def forward(self, z):
+            return torch.ops.aten._to_copy.default(z, dtype=torch.float32)
+
+    _check_op(M(), (_z(3, 5),), "to_copy_complex_to_real")  # shape (3,5) so last dim≠2
+
+
 @pytest.mark.unit
 def test_reshape_batch():
     class M(nn.Module):