Fix low precision nextafter in Inductor

test/inductor/test_cuda_repro.py

@@ -3074,6 +3074,62 @@ def compiled_divide(x, y):
 
                 self.assertEqual(eager_div, compiled_div)
 
+    @unittest.skipIf(not TEST_CUDA, "requires CUDA")
+    @parametrize("dtype", [torch.float16, torch.bfloat16])
+    def test_nextafter_low_precision_non_contiguous(self, dtype):
+        if dtype is torch.bfloat16 and not SM80OrLater:
+            self.skipTest("bfloat16 requires SM >= 80")
+
+        def fn(x, y):
+            return torch.nextafter(x, y)
+
+        x_base = torch.tensor(
+            [
+                [-0.0, 0.0, 1.0, -1.0],
+                [float("inf"), -float("inf"), 2.0, -2.0],
+                [0.0, -0.0, 1.0, -1.0],
+            ],
+            dtype=dtype,
+            device=device_type,
+        )
+        y_base = torch.tensor(
+            [
+                [0.0, -1.0, 0.0, 0.0],
+                [0.0, 0.0, float("inf"), -float("inf")],
+                [1.0, -1.0, float("nan"), float("nan")],
+            ],
+            dtype=dtype,
+            device=device_type,
+        )
+        x = x_base.t()
+        y = y_base.t()
+
+        self.assertFalse(x.is_contiguous())
+        self.assertFalse(y.is_contiguous())
+
+        expected = fn(x, y)
+        actual = torch.compile(fn, backend="inductor", fullgraph=True)(x, y)
+
+        self.assertEqual(torch.isnan(actual), torch.isnan(expected))
+        non_nan = ~torch.isnan(expected)
+        self.assertEqual(
+            actual[non_nan].view(torch.int16),
+            expected[non_nan].view(torch.int16),
+        )
+
+    @unittest.skipIf(not TEST_CUDA, "requires CUDA")
+    @parametrize("dtype", [torch.int16, torch.uint16])
+    def test_nextafter_integer_unsupported(self, dtype):
+        def fn(x, y):
+            return torch.nextafter(x, y)
+
+        x = torch.tensor([1, 2, 3], dtype=dtype, device=device_type)
+        y = torch.tensor([2, 1, 4], dtype=dtype, device=device_type)
+        compiled = torch.compile(fn, backend="inductor", fullgraph=True)
+
+        with self.assertRaisesRegex(NotImplementedError, "nextafter_cuda"):
+            compiled(x, y)
+
     @skipIfXpu(msg="triton dependency - torch-xpu-ops: 2554")
     @config.patch({"eager_numerics.division_rounding": False})
     @xfailIfROCm

test/inductor/test_op_dtype_prop.py

@@ -159,7 +159,6 @@ def test_op_dtype_support(self):
             "exp2",
             "abs",
             "hypot",
-            "nextafter",
         ]:
             # These ops do not support float16 and bfloat16.
             supported_dtypes = OpDtypeSupport.supported_dtypes[op_name]

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/lowering.py

@@ -8419,7 +8419,55 @@ 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)
+
+    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)
+
+    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()
+
+    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/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