From a39ce2952369fcfae79a1f104d7fdfb7a7b2e4aa Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Tue, 18 Aug 2026 21:04:40 +0000 Subject: [PATCH 1/3] Complex rewrite --- .../lowering/passes/complex_graph_rewrite.py | 22 +++++++++++++- .../dynamo/lowering/test_complex_rewrite.py | 30 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py index c8b7ad8571..3880562e6a 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py +++ b/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,25 @@ 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: + # In the [..., 2] real layout the copy is unchanged; only a complex + # dtype request needs remapping to its real counterpart. + kwargs = dict(node.kwargs) + dtype = kwargs.get("dtype") + if dtype is not None and dtype not in COMPLEX_DTYPES: + # complex -> real changes meaning, not layout; use the fallback. + return False + if dtype is not None: + kwargs["dtype"] = COMPLEX_TO_REAL_DTYPE[dtype] + with SubgraphBuilder(self.gm.graph, node) as b: + out = b(torch.ops.aten._to_copy.default, node.args[0]) + out.kwargs = kwargs + 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 +1316,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 diff --git a/tests/py/dynamo/lowering/test_complex_rewrite.py b/tests/py/dynamo/lowering/test_complex_rewrite.py index bae2a6bfcd..7cf543d93d 100644 --- a/tests/py/dynamo/lowering/test_complex_rewrite.py +++ b/tests/py/dynamo/lowering/test_complex_rewrite.py @@ -594,6 +594,36 @@ 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.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 + _check_op(M(), (_z(),), "to_copy_complex_dtype") + + @pytest.mark.unit def test_reshape_batch(): class M(nn.Module): From 8922660903c64740014dde233a391c2b32d0c78b Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 21 Aug 2026 18:37:50 +0000 Subject: [PATCH 2/3] review --- .../lowering/passes/complex_graph_rewrite.py | 18 ++++++----- .../dynamo/lowering/test_complex_rewrite.py | 32 +++++++++++++------ 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py index 3880562e6a..6c694cf690 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py @@ -1284,19 +1284,21 @@ def _rewrite_scalar_tensor(self, node: Node) -> bool: @_complex_unpacker(torch.ops.aten._to_copy.default) def _rewrite_to_copy(self, node: Node) -> bool: - # In the [..., 2] real layout the copy is unchanged; only a complex - # dtype request needs remapping to its real counterpart. + # 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") - if dtype is not None and dtype not in COMPLEX_DTYPES: - # complex -> real changes meaning, not layout; use the fallback. - return False - if dtype is not None: + 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: - out = b(torch.ops.aten._to_copy.default, node.args[0]) + 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 - out.meta["is_complex_layout"] = True + if not to_real: + out.meta["is_complex_layout"] = True node.replace_all_uses_with(out) self.gm.graph.erase_node(node) return True diff --git a/tests/py/dynamo/lowering/test_complex_rewrite.py b/tests/py/dynamo/lowering/test_complex_rewrite.py index 7cf543d93d..3cb9f61d51 100644 --- a/tests/py/dynamo/lowering/test_complex_rewrite.py +++ b/tests/py/dynamo/lowering/test_complex_rewrite.py @@ -601,9 +601,7 @@ 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" - } + 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") @@ -613,15 +611,31 @@ def forward(self, z): def test_to_copy_complex_dtype(): class M(nn.Module): def forward(self, z): - return torch.ops.aten._to_copy.default(z, dtype=torch.complex128) + return torch.ops.aten._to_copy.default(z, dtype=torch.complex64) - gm = _export_and_lower(M(), (_z(),)) - targets = { - n.target for n in gm.graph.nodes if n.op == "call_function" - } + # 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 - _check_op(M(), (_z(),), "to_copy_complex_dtype") + 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 From 592d404ea16490605218a70f358392d13ca29d97 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 21 Aug 2026 23:42:43 +0000 Subject: [PATCH 3/3] review --- .../lowering/passes/complex_graph_rewrite.py | 37 ++++++++++++++----- .../dynamo/lowering/test_complex_rewrite.py | 24 ++++++++++++ 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py index 6c694cf690..1309fb8883 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py @@ -1284,21 +1284,40 @@ def _rewrite_scalar_tensor(self, node: Node) -> bool: @_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: + 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_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: + 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 diff --git a/tests/py/dynamo/lowering/test_complex_rewrite.py b/tests/py/dynamo/lowering/test_complex_rewrite.py index 3cb9f61d51..447d716a20 100644 --- a/tests/py/dynamo/lowering/test_complex_rewrite.py +++ b/tests/py/dynamo/lowering/test_complex_rewrite.py @@ -638,6 +638,30 @@ def forward(self, z): _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):