From 05aa22cd761951c8562c677daa7c61ef05fa1d4b Mon Sep 17 00:00:00 2001 From: jyz Date: Fri, 26 Jun 2026 01:13:43 +0000 Subject: [PATCH 1/2] 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 4d4356a62f892e9fe00ad368ab89262716d0579b Mon Sep 17 00:00:00 2001 From: jyz Date: Sun, 5 Jul 2026 20:54:08 +0000 Subject: [PATCH 2/2] Fix low precision nextafter in Inductor --- test/inductor/test_cuda_repro.py | 56 +++++++++++++++++++++++ test/inductor/test_op_dtype_prop.py | 1 - torch/_inductor/codegen/triton.py | 3 +- torch/_inductor/lowering.py | 50 +++++++++++++++++++- torch/_inductor/runtime/triton_helpers.py | 25 ++++++++++ 5 files changed, 131 insertions(+), 4 deletions(-) diff --git a/test/inductor/test_cuda_repro.py b/test/inductor/test_cuda_repro.py index addcce90c8e21..660133d85adc0 100644 --- a/test/inductor/test_cuda_repro.py +++ b/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 diff --git a/test/inductor/test_op_dtype_prop.py b/test/inductor/test_op_dtype_prop.py index b8fc7010d4e0f..8a3fe21d8ff0e 100644 --- a/test/inductor/test_op_dtype_prop.py +++ b/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] 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..eda6d1f315fe4 100644 --- a/torch/_inductor/lowering.py +++ b/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 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)