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,46 @@ 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:
+        kwargs = dict(node.kwargs)
+        dtype = kwargs.get("dtype")
+        inp = node.args[0]
+        from_complex = self._is_complex_layout_node(inp)
+        to_complex = dtype is None or dtype in COMPLEX_DTYPES
+        if dtype is not None and to_complex:
+            kwargs["dtype"] = COMPLEX_TO_REAL_DTYPE[dtype]
+
+        with SubgraphBuilder(self.gm.graph, node) as b:
+            if to_complex and from_complex:
+                # remap dtype, [..., 2] layout unchanged
+                out = b(torch.ops.aten._to_copy.default, inp)
+                out.kwargs = kwargs
+                out.meta["is_complex_layout"] = True
+            elif to_complex:
+                # a real input needs a zero imaginary half, so 1 -> [1, 0]
+                re = b(torch.ops.aten._to_copy.default, inp)
+                re.kwargs = kwargs
+                im = b(torch.ops.aten.zeros_like.default, re)
+                out = self._inline_cat_re_im(b, re, im)
+            elif dtype == torch.bool:
+                # bool(a+bi) tests both halves for nonzero, not just a
+                re = b(torch.ops.aten.select.int, inp, -1, 0)
+                im = b(torch.ops.aten.select.int, inp, -1, 1)
+                re_bool = b(torch.ops.aten._to_copy.default, re)
+                re_bool.kwargs = kwargs
+                im_bool = b(torch.ops.aten._to_copy.default, im)
+                im_bool.kwargs = kwargs
+                out = b(torch.ops.aten.logical_or.default, re_bool, im_bool)
+            else:
+                # a real target discards the imaginary half
+                re = b(torch.ops.aten.select.int, inp, -1, 0)
+                out = b(torch.ops.aten._to_copy.default, re)
+                out.kwargs = kwargs
+            node.replace_all_uses_with(out)
+            self.gm.graph.erase_node(node)
+            return True
+
     # ------------------------------------------------------------------
     # Shape-manipulation handlers
     #
@@ -1297,6 +1337,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,74 @@ 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_to_copy_complex_to_bool():
+    """bool(0+1j) is True, so the imaginary half alone has to set the result."""
+
+    class M(nn.Module):
+        def forward(self, z):
+            return torch.ops.aten._to_copy.default(z, dtype=torch.bool)
+
+    # len 3 so a surviving [..., 2] layout is a shape mismatch, not a silent pass
+    z = torch.tensor([0 + 1j, 0 + 0j, 2 - 3j], dtype=torch.complex64)
+    _check_op(M(), (z,), "to_copy_complex_to_bool")
+
+
+@pytest.mark.unit
+def test_to_copy_real_to_complex():
+    """x.to(complex) pairs each element with a zero imaginary half."""
+
+    class M(nn.Module):
+        def forward(self, x):
+            return torch.ops.aten._to_copy.default(x, dtype=torch.complex64)
+
+    _check_op(M(), (torch.tensor([1.0, 2.0, 3.0]),), "to_copy_real_to_complex")
+
+
 @pytest.mark.unit
 def test_reshape_batch():
     class M(nn.Module):