From 05aa22cd761951c8562c677daa7c61ef05fa1d4b Mon Sep 17 00:00:00 2001 From: jyz Date: Fri, 26 Jun 2026 01:13:43 +0000 Subject: [PATCH 1/3] Fix #1870: fuse_as_graphmodule input ordering --- torch/fx/passes/utils/fuser_utils.py | 34 ++++++++++++++++++++++------ 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/torch/fx/passes/utils/fuser_utils.py b/torch/fx/passes/utils/fuser_utils.py index 4930278afffbd..e6884562a49b4 100644 --- a/torch/fx/passes/utils/fuser_utils.py +++ b/torch/fx/passes/utils/fuser_utils.py @@ -155,6 +155,33 @@ def fuse_as_graphmodule( ] = {} # mapping of nodes from old graph to placeholder in new graph node_map: dict[Node, Node] = {} # mapping of nodes from old graph to new graph + def create_placeholder(input_node: Node) -> None: + placeholder_node = subgraph.placeholder( + input_node.name, type_expr=input_node.type + ) + placeholder_node.meta = copy.copy(input_node.meta) + node_to_placeholder[input_node] = placeholder_node + + external_inputs: list[Node] = [] + external_inputs_set: set[Node] = set() + for node in nodes: + for input_node in node.all_input_nodes: + if ( + input_node not in partition_lookup_table + and input_node not in external_inputs_set + ): + external_inputs.append(input_node) + external_inputs_set.add(input_node) + + if external_inputs and all( + input_node.op == "placeholder" for input_node in external_inputs + ): + graph_order = {node: i for i, node in enumerate(gm.graph.nodes)} + external_inputs.sort(key=lambda node: graph_order[node]) + + for input_node in external_inputs: + create_placeholder(input_node) + # handles inputs through graph.node_copy's arg_transform functions def remap_inputs(x: Node) -> Node: if x.op == "get_attr": @@ -167,13 +194,6 @@ def remap_inputs(x: Node) -> Node: # the node should have been copied already, as we are copying graph in the topological order return node_map[x] - if x not in node_to_placeholder: - # x is not in subgraph, create a new placeholder for subgraph - placeholder_node = subgraph.placeholder(x.name, type_expr=x.type) - # copy all meta fields, even if some fields might be irrelevant for the placeholder node - placeholder_node.meta = copy.copy(x.meta) - node_to_placeholder[x] = placeholder_node - return node_to_placeholder[x] # copy nodes in topological order From ab012acc8cd4ef01d1577ce079c3a324a40dbd09 Mon Sep 17 00:00:00 2001 From: jyz Date: Sun, 5 Jul 2026 20:56:48 +0000 Subject: [PATCH 2/3] [inductor] Fix nextafter low-precision lowering --- torch/_inductor/codegen/triton.py | 3 +- torch/_inductor/lowering.py | 54 ++++++++++++++++++++++- torch/_inductor/runtime/triton_helpers.py | 25 +++++++++++ 3 files changed, 79 insertions(+), 3 deletions(-) diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py index 98374bd0a74b1..fb207f54b4963 100644 --- a/torch/_inductor/codegen/triton.py +++ b/torch/_inductor/codegen/triton.py @@ -2068,10 +2068,9 @@ def ldexp(x, n): return f"libdevice.ldexp({x}, {n}.to(tl.int32))" @staticmethod - @maybe_upcast_float32() # pyrefly: ignore [bad-override] def nextafter(x, y): - return f"libdevice.nextafter({x}, {y})" + return f"triton_helpers.nextafter({x}, {y})" @staticmethod # pyrefly: ignore [bad-override] diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py index 9ed9412609c95..2c301b78677ee 100644 --- a/torch/_inductor/lowering.py +++ b/torch/_inductor/lowering.py @@ -8419,7 +8419,59 @@ def signbit(x): register_pointwise_numeric(aten.hypot) register_pointwise_numeric(aten.log10) register_pointwise_numeric(aten.log2) -register_pointwise_numeric(aten.nextafter) + +register_op_dtype_propagation_rules( + "nextafter", + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH, + override_return_dtype=None, +) + + +@register_lowering( + aten.nextafter, + broadcast=True, + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH, +) +@register_lowering(prims.nextafter, type_promotion_kind=None) +def nextafter(self, other): + dtype = self.get_dtype() + if not dtype.is_floating_point: + return fallback_handler(aten.nextafter.default)(self, other) + + device = None + for inp in (self, other): + inp_device = inp.get_device() + if inp_device is not None and is_gpu(inp_device.type): + device = inp_device + break + if device is None: + device = self.get_device() + + if dtype in (torch.float16, torch.bfloat16) and not is_triton(device): + return fallback_handler(aten.nextafter.default)(self, other) + + self_loader = self.make_loader() + other_loader = other.make_loader() + ranges = self.get_size() + + def inner_fn(index): + x = self_loader(index) + y = other_loader(index) + if dtype in (torch.float16, torch.bfloat16): + # Triton loads low precision floats as fp32 by default. nextafter + # must step in the result dtype, so restore the fp16/bf16 value + # before the bit-level Triton helper runs. + x = ops.to_dtype(x, dtype, use_compute_types=False) + y = ops.to_dtype(y, dtype, use_compute_types=False) + return ops.to_dtype(ops.nextafter(x, y), dtype) + return ops.nextafter(x, y) + + return Pointwise.create( + device=device, + dtype=dtype, + inner_fn=inner_fn, + ranges=ranges, + ) from .codegen.common import BackendFeature, pointwise_overrides_data diff --git a/torch/_inductor/runtime/triton_helpers.py b/torch/_inductor/runtime/triton_helpers.py index 124af40967d35..a259545001af3 100644 --- a/torch/_inductor/runtime/triton_helpers.py +++ b/torch/_inductor/runtime/triton_helpers.py @@ -191,6 +191,31 @@ def maximum(a, b): return tl.where(mask, a, b) +@triton.jit +def nextafter(x, y): + if not is_floating(x) or x.dtype.primitive_bitwidth != 16: + return libdevice.nextafter(x, y) + + idtype: tl.constexpr = tl.core.get_int_dtype( + x.dtype.primitive_bitwidth, signed=False + ) + ix = x.to(idtype, bitcast=True) + iy = y.to(idtype, bitcast=True) + sign_mask: tl.constexpr = 1 << (x.dtype.primitive_bitwidth - 1) + + x_is_zero = (ix & (sign_mask - 1)) == 0 + x_is_pos = (ix & sign_mask) == 0 + step_up = (y > x) == x_is_pos + stepped = ix + tl.where(step_up, 1, -1).to(idtype) + zero_step = (iy & sign_mask) | 1 + + result = ( + tl.where(x_is_zero, zero_step, stepped).to(idtype).to(x.dtype, bitcast=True) + ) + result = tl.where(x == y, y, result) + return tl.where((x != x) | (y != y), x + y, result) + + @triton.jit def min2(a, dim): return tl.reduce(a, dim, minimum) From 1798708ffc4817a7c5da7f342613d6f243d4445c Mon Sep 17 00:00:00 2001 From: jyz Date: Sun, 5 Jul 2026 21:00:24 +0000 Subject: [PATCH 3/3] Fix unbacked memo reuse during fake replay --- test/inductor/test_cuda_repro.py | 16 ++++++++++ test/test_fx.py | 40 ++++++++++++++++++++++++ torch/_inductor/fx_utils.py | 4 +++ torch/_inductor/ir.py | 2 ++ torch/_subclasses/fake_tensor.py | 25 +++++++++++++++ torch/fx/experimental/symbolic_shapes.py | 27 ++++++++++++++-- torch/fx/passes/fake_tensor_prop.py | 5 +++ torch/fx/passes/shape_prop.py | 4 +++ 8 files changed, 121 insertions(+), 2 deletions(-) diff --git a/test/inductor/test_cuda_repro.py b/test/inductor/test_cuda_repro.py index addcce90c8e21..158fff5cc326f 100644 --- a/test/inductor/test_cuda_repro.py +++ b/test/inductor/test_cuda_repro.py @@ -933,6 +933,22 @@ def fn(x, y): if not same(fn(a, b), fn_optimized(a, b)): raise AssertionError + @unittest.skipIf(not TEST_CUDA, "requires CUDA") + @dynamo_config.patch(capture_dynamic_output_shape_ops=True) + def test_bool_mask_index_retrace_unbacked_memo(self): + def fn(x, w): + mask = torch.eye(4, dtype=torch.bool) + y = w[:, ~mask] + sym_w = (w + w.transpose(1, 2)) / 2 + sym_y = sym_w[:, ~mask] + return x + y.sum() + sym_y.sum() + + x = torch.randn((), device="cuda") + w = torch.randn(3, 4, 4, device="cuda") + + actual = torch.compile(fn, fullgraph=True)(x, w) + self.assertEqual(actual, fn(x, w)) + def test_simplify_dims(self): def fn(a): return (a + 1,) diff --git a/test/test_fx.py b/test/test_fx.py index 580c58bc86319..d14124678f36a 100644 --- a/test/test_fx.py +++ b/test/test_fx.py @@ -2148,6 +2148,46 @@ def forward(self, x: torch.Tensor): shape_prop.ShapeProp(gm=gm, fake_mode=fake_mode).propagate(*inp) self.assertEqual(len(fake_mode.shape_env.pending_fresh_unbacked_symbols), 0) + def test_shape_prop_replay_repeated_bool_mask_unbacked_memo(self): + from torch._dynamo.utils import detect_fake_mode + from torch.fx.experimental.proxy_tensor import make_fx + + def fn(x, w): + mask = torch.eye(4, dtype=torch.bool) + y = w[:, ~mask] + sym_w = (w + w.transpose(1, 2)) / 2 + sym_y = sym_w[:, ~mask] + return x + y.sum() + sym_y.sum() + + gm = make_fx(fn, tracing_mode="symbolic")( + torch.randn(()), torch.randn(3, 4, 4) + ) + bitwise_not_nodes = [ + n + for n in gm.graph.nodes + if n.target == torch.ops.aten.bitwise_not.default + ] + index_nodes = [ + n for n in gm.graph.nodes if n.target == torch.ops.aten.index.Tensor + ] + self.assertEqual(len(bitwise_not_nodes), 2) + self.assertEqual(len(index_nodes), 2) + self.assertTrue(all(n.meta.get("unbacked_bindings") for n in index_nodes)) + + # ShapeProp is replaying a pre-existing graph. Make both old binding + # sites share the same bool-mask input while preserving their distinct + # old unbacked bindings, matching the failing replay invariant. + index_nodes[1].replace_input_with(bitwise_not_nodes[1], bitwise_not_nodes[0]) + gm.graph.erase_node(bitwise_not_nodes[1]) + gm.recompile() + + fake_inputs = [ + node.meta["val"] for node in gm.graph.nodes if node.op == "placeholder" + ] + fake_mode = detect_fake_mode(fake_inputs) + shape_prop.ShapeProp(gm=gm, fake_mode=fake_mode).propagate(*fake_inputs) + self.assertEqual(len(fake_mode.shape_env.pending_fresh_unbacked_symbols), 0) + def test_nn_module_stack(self): class SubModule(torch.nn.Module): def __init__(self) -> None: diff --git a/torch/_inductor/fx_utils.py b/torch/_inductor/fx_utils.py index 2c7034f9cb46e..5607947d88892 100644 --- a/torch/_inductor/fx_utils.py +++ b/torch/_inductor/fx_utils.py @@ -25,6 +25,7 @@ from torch._subclasses.fake_tensor import FakeTensorMode from torch.fx.experimental.symbolic_shapes import ( compute_unbacked_bindings, + invalidate_unbacked_memos_for_replay, rebind_unbacked, statically_known_true, sym_eq, @@ -629,6 +630,9 @@ def node_invokes_subgraph( if not any_output_updated and "val" in node.meta: continue + if node.meta.get("unbacked_bindings"): + invalidate_unbacked_memos_for_replay(node, args, kwargs) + with V.fake_mode, enable_python_dispatcher(): new_fake_tensor = node.target(*args, **kwargs) diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py index f6b69fea448d4..b505465daac88 100644 --- a/torch/_inductor/ir.py +++ b/torch/_inductor/ir.py @@ -62,6 +62,7 @@ free_unbacked_symbols, GuardOnDataDependentSymNode, has_free_unbacked_symbols, + invalidate_unbacked_memos_for_replay, IterateExprs, rebind_unbacked, resolve_unbacked_bindings, @@ -7236,6 +7237,7 @@ def unflatten_args( enable_python_dispatcher(), _fallback_kernel_symbol_tracking_context(shape_env), ): + invalidate_unbacked_memos_for_replay(V.current_node, new_args, new_kwargs) example_output = kernel(*new_args, **new_kwargs) unbacked_bindings: dict[sympy.Symbol, pytree.KeyPath] | None = None diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py index 7656fdbdbb219..eb1dceeae7fcd 100644 --- a/torch/_subclasses/fake_tensor.py +++ b/torch/_subclasses/fake_tensor.py @@ -313,6 +313,25 @@ def maybe_get_fake_mode(t: object) -> FakeTensorMode | None: return None +def invalidate_unbacked_memos(t: object) -> None: + from torch._subclasses.functional_tensor import FunctionalTensor + + if isinstance(t, FakeTensor): + t.invalidate_unbacked_memos() + elif is_traceable_wrapper_subclass(t): + attrs, _ = t.__tensor_flatten__() + for attr in attrs: + invalidate_unbacked_memos(getattr(t, attr)) + elif isinstance(t, FunctionalTensor): + invalidate_unbacked_memos(t.elem) + elif isinstance(t, Tensor) and torch._is_functional_tensor(t): + reapply_views = torch._C._functionalization_reapply_views_tls() + unwrapped = torch._C._functorch._unwrap_functional_tensor(t, reapply_views) + invalidate_unbacked_memos(unwrapped) + elif isinstance(t, Tensor) and is_functorch_wrapped_tensor(t): + invalidate_unbacked_memos(torch._C._functorch.get_unwrapped(t)) + + @functools.cache def get_schema_info(func: OpOverload) -> torch._C._SchemaInfo: return torch._C._SchemaInfo(func._schema) @@ -827,6 +846,12 @@ class FakeTensor(Tensor): # this is an "infra" mode with lower dispatching precedence. _mode_key = torch._C._TorchDispatchModeKey.FAKE + def invalidate_unbacked_memos(self) -> None: + self.nonzero_memo = None + self.item_memo = None + self.unique_memo = None + self.unique_consecutive_memo = None + @property # pyrefly: ignore [bad-override] def device(self) -> torch.device: diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py index dca4d9aef2f41..e295bd529aa7d 100644 --- a/torch/fx/experimental/symbolic_shapes.py +++ b/torch/fx/experimental/symbolic_shapes.py @@ -729,13 +729,31 @@ def rebind_unbacked( continue # The old and new could be the same if you improperly hit the memo - # while retracing. Make sure you updated FakeTensorMode.epoch + # while retracing. Make sure you invalidated unbacked memos for + # this retrace, typically by updating FakeTensorMode.epoch. if raw_u0 == raw_u1: raise AssertionError(f"{raw_u0} possible memo disaster") # Reuse the OLD symbol name shape_env._rename_unbacked_to(raw_u1, raw_u0) +def invalidate_unbacked_memos_for_replay( + n: torch.fx.Node, args: object, kwargs: object +) -> None: + """ + Before replaying an existing unbacked binding site, clear fake-tensor + unbacked memos on its inputs so the node allocates a fresh symbol for + rebind_unbacked instead of reusing one from an earlier replayed node. + """ + if not n.meta.get("unbacked_bindings"): + return + + from torch._subclasses.fake_tensor import invalidate_unbacked_memos + + for value in pytree.tree_leaves((args, kwargs)): + invalidate_unbacked_memos(value) + + # NB: You could try to expand this to cover more cases by simply # detecting whenever you have an int output, but this is a bit # dangerous in case someone adds a function that returns an int but is @@ -8977,10 +8995,15 @@ def run_node(self, n: torch.fx.Node) -> Result: """ from torch._guards import detect_fake_mode - result = super().run_node(n) fake_mode = detect_fake_mode() if fake_mode is None: raise AssertionError("fake_mode must not be None") + + if n.meta.get("unbacked_bindings"): + args, kwargs = self.fetch_args_kwargs_from_env(n) + invalidate_unbacked_memos_for_replay(n, args, kwargs) + + result = super().run_node(n) rebind_unbacked(fake_mode.shape_env, n, result) return result diff --git a/torch/fx/passes/fake_tensor_prop.py b/torch/fx/passes/fake_tensor_prop.py index 688d850f2a004..36bb1d63b7a1c 100644 --- a/torch/fx/passes/fake_tensor_prop.py +++ b/torch/fx/passes/fake_tensor_prop.py @@ -41,6 +41,7 @@ def __init__( def run_node(self, n: Node) -> Any: from torch.fx.experimental.symbolic_shapes import ( compute_unbacked_bindings, + invalidate_unbacked_memos_for_replay, rebind_unbacked, ) @@ -76,6 +77,10 @@ def run_node(self, n: Node) -> Any: getattr(self.module, n.args[0].target), mode=self._mode ).propagate(*example_inputs) + if n.meta.get("unbacked_bindings"): + args, kwargs = self.fetch_args_kwargs_from_env(n) + invalidate_unbacked_memos_for_replay(n, args, kwargs) + result = super().run_node(n) rebind_unbacked(self._mode.shape_env, n, result) diff --git a/torch/fx/passes/shape_prop.py b/torch/fx/passes/shape_prop.py index b9f89820767f4..7c3fe78cc4aa7 100644 --- a/torch/fx/passes/shape_prop.py +++ b/torch/fx/passes/shape_prop.py @@ -160,6 +160,7 @@ def __init__(self, gm: torch.fx.GraphModule, fake_mode: Any = None) -> None: def run_node(self, n: Node) -> Any: from torch.fx.experimental.symbolic_shapes import ( compute_unbacked_bindings, + invalidate_unbacked_memos_for_replay, rebind_unbacked, ) @@ -171,6 +172,9 @@ def run_node(self, n: Node) -> Any: try: if self.fake_mode is not None: with self.fake_mode, enable_python_dispatcher(): + if n.meta.get("unbacked_bindings"): + args, kwargs = self.fetch_args_kwargs_from_env(n) + invalidate_unbacked_memos_for_replay(n, args, kwargs) result = super().run_node(n) rebind_unbacked(self.fake_mode.shape_env, n, result) else: