Fix unbacked memo reuse during fake replay

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,)

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:

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]

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)
 

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

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
 

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)

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:

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
 

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)
 

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:

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