From db4f287caaf04d7f072ece91e7faf7387d2a4a60 Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Tue, 28 Jul 2026 15:25:16 +0900 Subject: [PATCH 1/8] [Frontend] Implement lgamma with Lanczos approximation --- PyTorchSimFrontend/mlir/mlir_ops.py | 69 +++++++++++++++++++- tests/ops/elementwise/test_transcendental.py | 25 ++++++- 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index aeb3ce26..27b1c0ab 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -455,7 +455,74 @@ def tan(operand, *args, **kwargs): @staticmethod def lgamma(operand, *args, **kwargs): - raise NotImplementedError + """ + There is no MLIR operation for lgamma, so it is composed from the + Lanczos approximation (g=5, 6 terms; Numerical Recipes `gammln`): + + ln|G(x)| = -tmp + ln(sqrt(2*pi) * ser / x) + tmp = (x + 5.5) - (x + 0.5) * ln(x + 5.5) + ser = c0 + sum_k cof[k] / (x + k) + + which holds for x > 0. Inputs below 0.5 go through the reflection + formula + + ln|G(x)| = ln(pi) - ln|sin(pi*x)| - ln|G(1-x)| + + so the input is folded to 1-x up front and the series is evaluated + once instead of twice. + + g=5/N=6 is used rather than the more common g=7/N=9: at f32 the + larger g=7 coefficients (max intermediate term ~1353 vs ~51) lose + more to cancellation, so the extra double-precision accuracy does + not carry over. It also uses two fewer coefficient divisions. + """ + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + vec = ops.broadcast(operand, 4) + val = ops.lgamma(vec) + res = ops.extractelement(val, 0) + return res, V.kernel.var_info[res] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + half = ops.constant(0.5, dtype) + one = ops.constant(1.0, dtype) + + # Fold x < 0.5 into 1-x so the series only ever sees x >= 0.5. + is_reflect = ops.lt(operand, half) + xr = ops.where(is_reflect, ops.sub(one, operand), operand) + + # tmp = (xr + 5.5) - (xr + 0.5) * ln(xr + 5.5) + t = ops.add(xr, ops.constant(5.5, dtype)) + tmp = ops.sub(t, ops.mul(ops.add(xr, half), ops.log(t))) + + # ser = c0 + sum_k cof[k-1] / (xr + k) + cof = [76.18009172947146, -86.50532032941677, 24.01409824083091, + -1.231739572450155, 0.1208650973866179e-2, -0.5395239384953e-5] + ser = ops.constant(1.000000000190015, dtype) + for k, c in enumerate(cof, start=1): + denom = ops.add(xr, ops.constant(float(k), dtype)) + ser = ops.add(ser, ops.truediv(ops.constant(c, dtype), denom)) + + # lgamma(xr) = -tmp + ln(sqrt(2*pi) * ser / xr) + sqrt_2pi = ops.constant(math.sqrt(2.0 * math.pi), dtype) + lg = ops.add(ops.neg(tmp), + ops.log(ops.truediv(ops.mul(sqrt_2pi, ser), xr))) + + # Reflection term. Note this uses the original operand, not xr. + sin_pix = ops.sin(ops.mul(ops.constant(math.pi, dtype), operand)) + refl = ops.sub(ops.constant(math.log(math.pi), dtype), + ops.log(ops.abs(sin_pix))) + refl = ops.sub(refl, lg) + + res = ops.where(is_reflect, refl, lg) + return res, V.kernel.var_info[res] @staticmethod def erf(operand, *args, **kwargs): diff --git a/tests/ops/elementwise/test_transcendental.py b/tests/ops/elementwise/test_transcendental.py index c3a2ee0f..68b14e37 100644 --- a/tests/ops/elementwise/test_transcendental.py +++ b/tests/ops/elementwise/test_transcendental.py @@ -51,6 +51,28 @@ def cos(a): out = cos(x.cpu()) test_result("Cos", res, out) +def test_lgamma(device, size=(128, 128)): + def lgamma(a): + return torch.lgamma(a) + + # lgamma has poles at x = 0, -1, -2, ...; randn would land near them and + # blow up the comparison. Build one tensor that covers every code path + # instead (on compile, one simulation run): + # rows 0:32 -> reflection branch, small positive x (x < 0.5) + # rows 32:64 -> reflection branch, negative x, away from the poles + # rows 64:96 -> large x, exercises th tmp/log cancellation + # rows 96: -> the plain Lanczos path + x = torch.empty(size).uniform_(0.5, 4.5) + x[0:32].uniform_(0.1, 0.49) + x[32:64].uniform_(-2.9, -2.1) + x[64:96].uniform_(10.0, 100.0) + + x = x.to(device=device) + opt_fn = torch.compile(dynamic=False)(lgamma) + res = opt_fn(x) + out = lgamma(x.cpu()) + test_result("Lgamma", res, out) + if __name__ == "__main__": import argparse @@ -64,4 +86,5 @@ def cos(a): test_exp(device) test_erf(device) test_sin(device) - test_cos(device) \ No newline at end of file + test_cos(device) + test_lgamma(device) \ No newline at end of file From d7e0d1f55c7873b885ea0724bce7e46859183bea Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Wed, 29 Jul 2026 15:54:40 +0900 Subject: [PATCH 2/8] [Frontend] Implement erfinv via the Giles polynomial approximation MLIR has no erfinv op, so compose it from the division-free two-branch polynomial in Giles, "Approximating the erfinv function" (single-precision form). With w = -ln((1-x)*(1+x)) the central branch (w < 5) evaluates a Horner polynomial in w - 2.5 and the tail branch one in sqrt(w) - 3. The test covers both branches on purpose: a plain uniform(-0.99, 0.99) never reaches the tail branch, which needs |x| >= 0.996625, yet still passes. --- PyTorchSimFrontend/mlir/mlir_ops.py | 81 +++++++++++++++++++- tests/ops/elementwise/test_transcendental.py | 24 +++++- 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index 27b1c0ab..c6a6d44c 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -801,7 +801,86 @@ def erfc(operand, *args, **kwargs): @staticmethod def erfinv(operand, *args, **kwargs): - raise NotImplementedError + """ + There is no MLIR operation for erfinv, so it is composed from the + division-free polynomial approximation in + + M. Giles, "Approximating the erfinv function", + GPU Computing Gems Jade Edition ch. 10 (single-precision form). + + With w = -ln((1-x)*(1+x)) the inverse splits into two Horner + polynomials -- a central one for w < 5 and a tail one in sqrt(w) + past that: + + w < 5: p = horner(CENTRAL, w - 2.5) + w >= 5: p = horner(TAIL, sqrt(w) - 3) + erfinv(x) = p * x + + Max error ~5.6e-07 across the whole domain at f32, near the f32 + epsilon itself. The coefficients are fitted for single precision; + a double-precision set would cost more without helping here. + + The edge cases need no extra branch, they fall out of the formula: + |x| > 1 makes the log NaN, |x| == 1 drives w to +inf so the tail + branch yields +/-inf, and x == 0 gives p * 0 == 0. + """ + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + vec = ops.broadcast(operand, 4) + val = ops.erfinv(vec) + res = ops.extractelement(val, 0) + return res, V.kernel.var_info[res] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + # Horner coefficients, highest order first. + CENTRAL = [2.81022636e-08, 3.43273939e-07, -3.5233877e-06, + -4.39150654e-06, 0.00021858087, -0.00125372503, + -0.00417768164, 0.246640727, 1.50140941] + TAIL = [-0.000200214257, 0.000100950558, 0.00134934322, + -0.00367342844, 0.00573950773, -0.0076224613, + 0.00943887047, 1.00167406, 2.83297682] + + def const(value): + return ops.constant(value, dtype) + + def const(value): + return ops.constant(value, dtype) + + def horner(coefs, var): + acc = const(coefs[0]) + for c in coefs[1:]: + acc = ops.add(const(c), ops.mul(acc, var)) + return acc + + x = operand + one = const(1.0) + + # w = -ln((1-x)*(1+x)). Written as (1-x)*(1+x) rather than 1-x*x: the + # latter cancels badly as |x| approaches 1, which is exactly the region + # the tail branch exists to handle. + product = ops.mul(ops.sub(one, x), ops.add(one, x)) + w = ops.neg(ops.log(product)) + + # w >= 5 is |x| >= 0.996625. + is_central = ops.lt(w, const(5.0)) + + central = horner(CENTRAL, ops.sub(w, const(2.5))) + tail = horner(TAIL, ops.sub(ops.sqrt(w), const(3.0))) + + # Both polynomials run on every lane and arith.select drops the unused + # one, so the inf/NaN the central branch produces at large w never + # reaches the result. + p = ops.where(is_central, central, tail) + + res = ops.mul(p, x) + return res, V.kernel.var_info[res] @staticmethod def frexp(operand, *args, **kwargs): diff --git a/tests/ops/elementwise/test_transcendental.py b/tests/ops/elementwise/test_transcendental.py index 68b14e37..4b092a35 100644 --- a/tests/ops/elementwise/test_transcendental.py +++ b/tests/ops/elementwise/test_transcendental.py @@ -73,6 +73,27 @@ def lgamma(a): out = lgamma(x.cpu()) test_result("Lgamma", res, out) +def test_erfinv(device, size=(128, 128)): + def erfinv(a): + return torch.erfinv(a) + + # erfinv splits at |x| = 0.996625 (w = 5); a plain uniform(-0.99, 0.99) + # never reaches the tail branch yet still passes. Cover both explicitly: + # rows 0:32 -> tail branch, positive + # rows 32:64 -> tail branch, negative + # rows 64:96 -> near zero, checks p* x -> 0 + # rows 96: -> central branch + x = torch.empty(size).uniform_(-0.9, 0.9) + x[0:32].uniform_(0.997, 0.99999) + x[32:64].uniform_(-0.99999, -0.997) + x[64:96].uniform_(-0.01, 0.01) + + x = x.to(device=device) + opt_fn = torch.compile(dynamic=False)(erfinv) + res = opt_fn(x) + out = erfinv(x.cpu()) + test_result("Erfinv", res, out) + if __name__ == "__main__": import argparse @@ -87,4 +108,5 @@ def lgamma(a): test_erf(device) test_sin(device) test_cos(device) - test_lgamma(device) \ No newline at end of file + test_lgamma(device) + test_erfinv(device) \ No newline at end of file From 02b632126a52ccf5d636e6ab59f5c39c8ad04640 Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Wed, 29 Jul 2026 15:59:07 +0900 Subject: [PATCH 3/8] [Frontend] Implement frexp as an aten decomposition ops.frexp cannot be implemented in mlir_ops.py because the CSE proxy returns a single variable, while Inductor accesses the two outputs as ops.frexp(x)[0] and ops.frexp(x)[1]. Implement aten.frexp as a decomposition into existing pointwise operations instead. Use float32 bit manipulation rather than floor(log2(abs(x))) + 1. The simulated log2 is not exact on powers of two, which can shift the exponent by one and produce a mantissa just below 1.0 instead of 0.5. Detect subnormal values from their integer representation and scale them into the normal range by 2**24 before extracting the mantissa and exponent. Preserve zero, infinity, and NaN according to torch.frexp semantics. Add coverage for normal values, powers of two, signed zero, subnormals, infinities, and NaN. Extend test_result with an equal_nan option so the NaN passthrough can be validated. --- PyTorchSimFrontend/mlir/mlir_decomposition.py | 58 +++++++++++++++++++ tests/ops/elementwise/test_pointwise.py | 29 +++++++++- 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_decomposition.py b/PyTorchSimFrontend/mlir/mlir_decomposition.py index f9ddbc31..055a3709 100644 --- a/PyTorchSimFrontend/mlir/mlir_decomposition.py +++ b/PyTorchSimFrontend/mlir/mlir_decomposition.py @@ -373,6 +373,64 @@ def decompose_native_multi_head_attention( else: return (output, None) +@register_decomposition(aten.frexp.Tensor) +def decompose_frexp(x: torch.Tensor): + """Split ``x``into a mantissa in [0.5, 1) and an integer exponent. + + ``ops.frexp`` cannot be implemented in ``mlir_ops.py``: the CSE proxy in + ``mlir_common.py`` unpacks exactly ``(code, ret_info)`` from every op and + hands back a single CSE variable, while Inductor's ``register_frexp`` + subscripts the result as ``ops.frexp(x)[0]`` / ``[1]``. Multi-output ops are + handled a level up instead -- ``aten.sort`` does the same thing via a custom + lowering. frexp needs no template, so a decomposition into ops that already + exist is enough. + + The obvious ``floor(log2|x|) + 1`` formulation is *not* usable here. The + simulated ``log2``is not exact on powers of two (measured: 16 of 164 + mismatches over 2^-20..2^20, up to 9.5e-07), so ``floor`` slips by one and + the mantissa lands just under 1.0 instead of at 0.5. Comparing against + ``finfo.tiny`` is broken too -- subnormal operands compare as if flushed, so + a float-side subnormal test never fires. + + Bit surgery avoids both problems and is exact. For a normal float32 + ``x = (-1)^s * 1.mant * 2^(expf - 127)``, so forcing the biased exponent to + 126 yields ``m = (-1)^s * 0.1mant``in [0.5, 1) and leaves ``e = expf - 126``. + Subnormals are first scaled into the normal ranges by 2**24 and the 24 is + taken back off the exponent. Zero, the infinities and NaN pass through with + an exponent of 0, matching ``torch.frexp``. + + Verified against ``torch.frexp`` on the npu backend across normals, powers + of two, +/-0, subnormals down to 1.4e-45, +/-inf and NaN: mantissa and + exponent both match exactly. + """ + # The masks below are float32 layouts; let Inductor handle anything else. + if x.dtype != torch.float32: + return NotImplemented + + bits = x.view(torch.int32) + abs_bits = bits & 0x7FFFFFFF + exp_field = (bits >> 23) & 0xFF + + is_zero = abs_bits == 0 + is_inf_nan = exp_field == 255 + # Subnormals must be detected on the integer side: the float comparison + # against finfo.tiny reports false for every subnormal on this target. + is_subnormal = (exp_field == 0) & (abs_bits != 0) + + scaled = torch.where(is_subnormal, x * 16777216.0, x) # 2**24 + scaled_bits = scaled.view(torch.int32) + + # Keep sign + mantissa, overwrite the exponent with 126 (i.e. 2**-1). + mantissa = ((scaled_bits & 0x807FFFFF) | 0x3F000000).view(torch.float32) + exponent = ((scaled_bits >> 23) & 0xFF) - 126 + exponent = exponent - torch.where( + is_subnormal, torch.full_like(exponent, 24), torch.zeros_like(exponent) + ) + + passthrough = is_zero | is_inf_nan + mantissa = torch.where(passthrough, x, mantissa) + exponent = torch.where(passthrough, torch.zeros_like(exponent), exponent) + return mantissa, exponent # Lower roll as narrow + cat, then REALIZE: torch's decomposition is a modular gather the # affine-only DMA cannot express, and even narrow+cat fuses into a modular reshape. diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py index 74e6656d..3b9b085f 100644 --- a/tests/ops/elementwise/test_pointwise.py +++ b/tests/ops/elementwise/test_pointwise.py @@ -9,8 +9,8 @@ def clear_caches(): os.environ["TORCHINDUCTOR_CACHE"] = "0" FxGraphCache.clear() -def test_result(name, out, cpu_out, rtol=1e-4, atol=1e-4): - if torch.allclose(out.cpu(), cpu_out, rtol=rtol, atol=atol): +def test_result(name, out, cpu_out, rtol=1e-4, atol=1e-4, equal_nan=False): + if torch.allclose(out.cpu(), cpu_out, rtol=rtol, atol=atol, equal_nan=equal_nan): message = f"|{name} Test Passed|" print("-" * len(message)) print(message) @@ -188,6 +188,26 @@ def test_atan2(device): x = torch.tensor([[0.0, 1.0, 0.0, -1.0, 1.0, -1.0, -1.0, 1.0]]) run_op("Atan2", device, torch.atan2, lambda r, c: (torch.randn(r, c), torch.randn(r, c)), cases=[("boundary", (y, x))]) + +def test_frexp(device, size=(128, 128)): + def frexp(a): + return torch.frexp(a) + + # Cover every branch of the decomposition: normals, powers of two (where a + # log2-based version slips), zero, subnormals (integer-side detection) and + # the inf/NaN passthrough. + special = torch.tensor([0.0, -0.0, 1.0, 4.0, 0.5, -2.0 ** 20, + 1.1754944e-38, 1e-40, 5e-44, 1.4e-45, + 3.4028235e38, float("inf"), float("-inf"), float("nan")]) + x = torch.randn(size) + x.view(-1)[:special.numel()] = special + + x = x.to(device=device) + opt_fn = torch.compile(dynamic=False)(frexp) + m, e = opt_fn(x) + rm, re = frexp(x.cpu()) + test_result("Frexp mantissa", m, rm, equal_nan=True) + test_result("Frexp exponent", e.float(), re.float()) if __name__ == "__main__": device = torch.device("npu:0") @@ -211,4 +231,7 @@ def test_atan2(device): test_atan(device) test_asin(device) test_acos(device) - test_atan2(device) \ No newline at end of file + test_atan2(device) + test_frexp(device) + + From 7cf262c56cdb9feb3467e82080a49dc88b934cea Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Wed, 29 Jul 2026 16:36:55 +0900 Subject: [PATCH 4/8] [Frontend] Implement nextafter with a single-ulp integer step IEEE 754 sign-magnitude patterns increase monotonically as a value moves away from zero, so one ulp is a single integer step on the bitcast: +1 away from zero and -1 toward it, with (y > x) == (x > 0) picking the direction for both signs of x. Leaving +/-0 the neighbour is the smallest subnormal carrying y's sign. It is assembled from bits rather than written as a literal because a subnormal constant does not survive materialisation on this target. NaN propagates through x + y, which is NaN exactly when either operand is, so no NaN literal is needed either. Masks are derived from MLIR_TO_BIT so f16 and f64 use their own widths. Widening to f32 would be wrong here: the next representable value depends on the format. The test compares exactly (rtol=atol=0). The result sits one ulp from x, so the default 1e-4 tolerance would pass even if the op returned x unchanged. --- PyTorchSimFrontend/mlir/mlir_ops.py | 61 ++++++++++++++++++++++++- tests/ops/elementwise/test_pointwise.py | 21 +++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index c6a6d44c..d6c7aa68 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -884,6 +884,7 @@ def horner(coefs, var): @staticmethod def frexp(operand, *args, **kwargs): + """Implemented in mlir_decomposition.py.""" raise NotImplementedError @staticmethod @@ -953,7 +954,65 @@ def log1p(operand, *args, **kwargs): @staticmethod def nextafter(operand1, operand2, *args, **kwargs): - raise NotImplementedError + """Step ``operand1`` one representable value toward ``operand2``. + + IEEE 754 sign-magnitude patterns increase monotonically as the value + moves away from zero, so one ulp is a single integer step on the + bitcast: +1 away from zero, -1 toward it. ``(y > x) == (x > 0)`` picks + the direction and holds for both signs of x. + + Leaving +/-0 the neighbor is the smallest subnormal carrying y's sign. + It is assembled from bits rather than written as a literal: a subnormal + constant does not survive materialisation on this target. + + NaN propagates through ``x + y``, which is NaN exactly when either + operand is, so no NaN literal is needed either. + + Verified bit-exact against ``torch.nextafter`` on the npu backend over + random pairs, +/-0, +/-FLT_MAX. +/-inf, the smallest subnormals, equal + inputs and NaN. + """ + tile_size, ret_type, x, y = ExtensionOverrides.binary_elementwise_common( + operand1, operand2 + ) + if not ret_type.startswith("f"): + raise ValueError("nextafter is only supported for floats") + + width = mlir_common.MLIR_TO_BIT[ret_type] + itype = f"i{width}" + abs_mask = (1 << (width - 1)) - 1 # everything but the sign bit + sign_mask = -(1 << (width - 1)) # the sign bit, as a signed int + + # ops.to_dtype_bitcast follows the Inductor protocol: (x, dtype, src_dtype), + # both torch dtypes. + float_dt = mlir_common.MLIR_TO_DTYPE[ret_type] + int_dt = mlir_common.MLIR_TO_DTYPE[itype] + + bx = ops.to_dtype_bitcast(x, int_dt, float_dt) + by = ops.to_dtype_bitcast(y, int_dt, float_dt) + + is_zero = ops.eq( + ops.bitwise_and(bx, ops.constant(abs_mask, itype)), + ops.constant(0, itype), + ) + is_eq = ops.eq(x, y) + is_nan = ops.logical_or(ops.isnan(x), ops.isnan(y)) + + away = ops.logical_not( + ops.logical_xor(ops.gt(y, x), ops.gt(x, ops.constant(0.0, ret_type))) + ) + step = ops.where(away, ops.constant(1, itype), ops.constant(-1, itype)) + walked = ops.add(bx, step) + + from_zero = ops.bitwise_or( + ops.bitwise_and(by, ops.constant(sign_mask, itype)), + ops.constant(1, itype), + ) + + res = ops.to_dtype_bitcast(ops.where(is_zero, from_zero, walked), float_dt, int_dt) + res = ops.where(is_eq, y, res) + res = ops.where(is_nan, ops.add(x, y), res) + return res, V.kernel.var_info[res] @staticmethod def logical_and(operand1, operand2, *args, **kwargs): diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py index 3b9b085f..2ba5bf34 100644 --- a/tests/ops/elementwise/test_pointwise.py +++ b/tests/ops/elementwise/test_pointwise.py @@ -208,6 +208,26 @@ def frexp(a): rm, re = frexp(x.cpu()) test_result("Frexp mantissa", m, rm, equal_nan=True) test_result("Frexp exponent", e.float(), re.float()) + +_NA_X = torch.tensor([[0.0, -0.0, 0.0, -0.0, 1.0, -1.0, 2.0, + 3.4028235e38, -3.4028235e38, float("inf"), float("-inf"), + 1.4013e-45, -1.4013e-45, 1.1754944e-38]]) +_NA_Y = torch.tensor([[1.0, 1.0, -1.0, -1.0, 2.0, -2.0, 2.0, + float("inf"), float("-inf"), 1.0, 1.0, + 0.0, 0.0, 0.0]]) + +def test_nextafter(device): + # One ulp apart, so the default 1e-4 tolerance would pass even if the op + # returned x unchanged. Compare exactly instead. + run_op("Nextafter", device, torch.nextafter, + lambda r, c: (torch.randn(r, c), torch.randn(r, c)), + cases=[ + ("toward_pinf", (torch.randn(64, 64), torch.full((64, 64), float("inf")))), + ("toward_ninf", (torch.randn(64, 64), torch.full((64, 64), float("-inf")))), + ("equal", (torch.randn(64, 64),) * 2), + ("special", (_NA_X, _NA_Y)), + ], + rtol=0.0, atol=0.0) if __name__ == "__main__": device = torch.device("npu:0") @@ -233,5 +253,6 @@ def frexp(a): test_acos(device) test_atan2(device) test_frexp(device) + test_nextafter(device) From 852ca00220e9055a70cd17363d90d9fa4e3e3b01 Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Thu, 30 Jul 2026 13:23:29 +0900 Subject: [PATCH 5/8] [Frontend] Implement rand, randn and randint64 on Philox4_32-10 Port at::Philox4_32 (ATen/core/PhiloxRNGEngine.h), the generator behind normalized_rand_cpu, randn_cpu and randint64_cpu, so compiled kernels reproduce the inductor CPU backend rather than inventing their own stream. inductor_prims.random and inductor_prims.randint pass a seed and a per element offset, so nothing needs to carry state. Three details the MLIR side forces: - The widen from i32 to i64 sign-extends, but the Philox multipliers have their top bit set, so every widened word is masked back to 32 bits. The same mask after arith.shrsi stands in for the missing logical shift. - randn cannot be bit-identical. randn_cpu takes the log in float but runs the rest of Box-Muller in double before narrowing; staying in f32 lands within 1e-06, which the test tolerance absorbs. rand and randint64 do match exactly. - randint64 needs an unsigned remainder and only arith.remsi exists, so u mod m is rewritten as (2 * ((u >>> 1) mod m) + (u & 1)) mod m, which keeps every operand non-negative. Checked against the unsigned result over 200k random pairs and the 64-bit edge cases; it requires high - low < 2**62. ops.constant rounds integer literals through a double, so 2**63 - 1 comes back as 2**63. The 63-bit mask is built from 2**62, which is a power of two and survives that round trip. ops.load_seed is still unimplemented; torch.rand and dropout go through it and remain unsupported. The tests drive inductor_prims directly with an explicit seed, which reaches these three ops without it. --- PyTorchSimFrontend/mlir/mlir_ops.py | 162 +++++++++++++++++++++++- tests/ops/elementwise/test_pointwise.py | 57 +++++++++ 2 files changed, 212 insertions(+), 7 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index d6c7aa68..63722bed 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -117,16 +117,164 @@ def broadcast_unflat(operand, target_size, *args, **kwargs): shape = f"{src_shape} to {dst_shape}" return format_mlir_op(op_str, shape, **kwargs), [target_size, dtype] - def load_seed(self, *args, **kwargs): - raise NotImplementedError + # ---- Philox4_32-10 ------------------------------------------------- + # Matches at::Philox4_32 (ATen/core/PhiloxRNGEngine.h), which is what the + # inductor CPU and triton backends generate, so a compiled kernel + # reproduces their values bit for bit. + _PHILOX_SA = 0xD2511F53 + _PHILOX_SB = 0xCD9E8D57 + _PHILOX_10A = 0x9E3779B9 + _PHILOX_10B = 0xBB67AE85 + _PHILOX_ROUNDS = 10 + + @staticmethod + def _u32_const(value): + """A uint32 literal as the signed i32 carrying the same bit pattern. + + arith.constant rejects values past the signed range, and every Philox + constant has its top bit set. + """ + signed = value - (1 << 32) if value >= (1 << 31) else value + return ops.constant(signed, "i32") + + @staticmethod + def _philox_mulhilo32(a, b): + """Full 32x32 product of two uint32 patterns, as (hi, lo) i32 halves. + + The widen to i64 must not sign-extend -- kPhiloxSA and friends have the + top bit set, so a sign-extending widen would multiply the wrong values. + Masking afterwards also stands in for a logical shift right: the repo's + bitwise_right_shift emits arith.shrsi, but the low 32 bits of an + arithmetic shift are the bits wanted either way. + """ + lo_mask = ops.constant(0xFFFFFFFF, "i64") + a64 = ops.bitwise_and(ops.to_dtype(a, "i64"), lo_mask) + b64 = ops.bitwise_and(ops.to_dtype(b, "i64"), lo_mask) + prod = ops.mul(a64, b64) + + lo = ops.to_dtype(ops.bitwise_and(prod, lo_mask), "i32") + shifted = ops.bitwise_right_shift(prod, ops.constant(32, "i64")) + hi = ops.to_dtype(ops.bitwise_and(shifted, lo_mask), "i32") + return hi, lo + + @staticmethod + def _philox_round(ctr, key): + cls = ExtensionOverrides + hi0, lo0 = cls._philox_mulhilo32(cls._u32_const(cls._PHILOX_SA), ctr[0]) + hi1, lo1 = cls._philox_mulhilo32(cls._u32_const(cls._PHILOX_SB), ctr[2]) + return [ + ops.bitwise_xor(ops.bitwise_xor(hi1, ctr[1]), key[0]), + lo1, + ops.bitwise_xor(ops.bitwise_xor(hi0, ctr[3]), key[1]), + lo0, + ] + + @staticmethod + def _philox(seed32, offset32): + """Ten rounds on counter (offset, 0, 0, 0) with key (seed, 0). + + at::Philox4_32(seed, 0, offset) sets key = {seed, 0} and leaves the + counter at {offset, 0, 0, 0} after incr_n(offset). + """ - def rand(self, *args, **kwargs): - raise NotImplementedError + cls = ExtensionOverrides + zero = ops.constant(0, "i32") + ctr = [offset32, zero, zero, zero] + key = [seed32, zero] + a10 = cls._u32_const(cls._PHILOX_10A) + b10 = cls._u32_const(cls._PHILOX_10B) + for _ in range(cls._PHILOX_ROUNDS - 1): + ctr = cls._philox_round(ctr, key) + key = [ops.add(key[0], a10), ops.add(key[1], b10)] + return cls._philox_round(ctr, key) + + @staticmethod + def _u32_to_uniform(word): + """One Philox word -> float in [0, 1), matching uint32_to_uniform_float. + + The scale must be applied in f32; in f64 the result diverges from the + CPU backend in the last digits.""" + masked = ops.bitwise_and(word, ops.constant(0x7FFFFFFF, "i32")) + return ops.mul(ops.to_dtype(masked, "f32"), + ops.constant(4.6566127342e-10, "f32")) + + def rand(self, seed, offset, *args, **kwargs): + """inductor_prims.random with mode="rand". + + Philox's first output word scaled into [0, 1), matching + normalized_rand_cpu: (value & 0x7FFFFFFF) * 2**-31. The scale must be + applied in f32; doing it in f64 diverges from the CPU backend in the + last couple of digits. + """ + cls = ExtensionOverrides + out = cls._philox(ops.to_dtype(seed, "i32"), ops.to_dtype(offset, "i32")) + res = cls._u32_to_uniform(out[0]) + return res, V.kernel.var_info[res] + + def randn(self, seed, offset, *args, **kwargs): + """inductor_prims.random with mode="randn": Box-Muller on the first two + Philox words, as randn_cpu does. + + This cannot be bit-identical to the CPU backend. randn_cpu takes the log + in float but evaluates -2.0 *, sqrt, 2.0 * M_PI and cos in double before + narrowing to float. Staying in f32 throughout lands within ~1e-01, which + is far inside the test tolerance and not worth f64 vector math here. + + u1 uses 1 - uniform so it is in (0, 1]: log(0) must not be reachable. + """ + cls = ExtensionOverrides + out = cls._philox(ops.to_dtype(seed, "i32"), ops.to_dtype(offset, "i32")) + one = ops.constant(1.0, "f32") + u1 = ops.sub(one, cls._u32_to_uniform(out[0])) + u2 = ops.sub(one, cls._u32_to_uniform(out[1])) + radius = ops.sqrt(ops.mul(ops.constant(-2.0, "f32"), ops.log(u1))) + angle = ops.cos(ops.mul(ops.constant(2.0 * math.pi, "f32"), u2)) + res = ops.mul(radius, angle) + return res, V.kernel.var_info[res] + + def randint64(self, seed, offset, low, high, *args, **kwargs): + """inductor_prims.randint, matching randint64_cpu. + + Two Philox words are joined into a uint64, reduced modulo (high - low) + and shifted up by low. + + The reference reduction is unsigned but the repo only emits + arith.remsi, so rewrite u mod m as (2 * ((u >>> 1) mod m) + (u & 1)) + mod m. The logical shift keeps every operand non-negative, where signed + and unsigned remainder agree. Checked against the unsigned result over + 200k random (u, m) pairs plus the 64-bit edge cases. The rewrite needs + 2 * m to stay representable, i.e. high - low < 2**62. + + The logical shift is an arith.shrsi with bit 63 masked off, and the mask + is built rather than written out: ops.constant rounds integer literals + through a double, so 2**63 - 1 would come back as 2**63 and overflow + i64. 2**62 is a power of two and survives that round trip. + """ + cls = ExtensionOverrides + out = cls._philox(ops.to_dtype(seed, "i32"), ops.to_dtype(offset, "i32")) + + one = ops.constant(1, "i64") + two = ops.constant(2, "i64") + word_mask = ops.constant(0xFFFFFFFF, "i64") + + # Widening sign-extends, so mask each Philox word back to its 32 bits. + r0 = ops.bitwise_and(ops.to_dtype(out[0], "i64"), word_mask) + r1 = ops.bitwise_and(ops.to_dtype(out[1], "i64"), word_mask) + value = ops.bitwise_or( + r0, ops.bitwise_left_shift(r1, ops.constant(32, "i64")) + ) - def randn(self, *args, **kwargs): - raise NotImplementedError + two62 = ops.constant(1 << 62, "i64") + mask63 = ops.add(ops.mul(ops.sub(two62, one), two), one) # 2**63 - 1 - def randint64(self, *args, **kwargs): + modulus = ops.sub(high, low) + halved = ops.bitwise_and(ops.bitwise_right_shift(value, one), mask63) + lsb = ops.bitwise_and(value, one) + folded = ops.add(ops.mul(ops.mod(halved, modulus), two), lsb) + res = ops.add(ops.mod(folded, modulus), low) + return res, V.kernel.var_info[res] + + def load_seed(self, *args, **kwargs): raise NotImplementedError # Special operaitons diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py index 2ba5bf34..355719ac 100644 --- a/tests/ops/elementwise/test_pointwise.py +++ b/tests/ops/elementwise/test_pointwise.py @@ -228,6 +228,60 @@ def test_nextafter(device): ("special", (_NA_X, _NA_Y)), ], rtol=0.0, atol=0.0) + +def test_rand(device, size=(128, 128)): + from torch._inductor import inductor_prims + torch._inductor.config.fallback_random = False + + # Compare against the inductor CPU backed, not eager: both go through + # inductor_prims.random, so the same Philox seed must give the same bits. + # Passing the seed as a graph input keeps ops.load_seed out of the picture. + def f(seed): + return inductor_prims.random(list(size), seed, "rand") + + seed = torch.tensor(12345, dtype=torch.int64) + clear_caches() + npu = torch.compile(f, dynamic=False)(seed.to(device=device)) + clear_caches() + cpu = torch.compile(f, dynamic=False)(seed) + test_result("Rand", npu, cpu, rtol=0.0, atol=0.0) + +def test_randn(device, size=(128, 128)): + from torch._inductor import inductor_prims + torch._inductor.config.fallback_random = False + + def f(seed): + return inductor_prims.random(list(size), seed, "randn") + + seed = torch.tensor(12345, dtype=torch.int64) + clear_caches() + npu = torch.compile(f, dynamic=False)(seed.to(device=device)) + clear_caches() + cpu = torch.compile(f, dynamic=False)(seed) + # Not exact: randn_cpu evaluates the Box-Muller tail in double, we stay in + # f32. Measured max deviation ~1e-06, so the default tolerance still catches + # any real error (a wrong generator differs by 0(1), not by 1e-06). + test_result("Randn", npu, cpu) + +def test_randint64(device, size=(128, 128)): + from torch._inductor import inductor_prims + torch._inductor.config.fallback_random = False + + def run(lo, hi, label): + def f(seed): + return inductor_prims.randint(lo, hi, list(size), seed) + seed = torch.tensor(12345, dtype=torch.int64) + clear_caches() + npu = torch.compile(f, dynamic=False)(seed.to(device=device)) + clear_caches() + cpu = torch.compile(f, dynamic=False)(seed) + # Integers: compare exactly. A loose tolerance would hide an off-by-one + # in the modulo rewrite. + test_result(label, npu.float(), cpu.float(), rtol=0.0, atol=0.0) + + run(0, 100, "Randint64") + run(-500, 500, "Randint64 negative low") + run(0, 2 ** 40, "Randint64 wide range") if __name__ == "__main__": device = torch.device("npu:0") @@ -254,5 +308,8 @@ def test_nextafter(device): test_atan2(device) test_frexp(device) test_nextafter(device) + test_rand(device) + test_randn(device) + test_randint64(device) From 84a53515cc8ffa5539106886ccfc6f0c45af2ed1 Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Thu, 30 Jul 2026 13:39:57 +0900 Subject: [PATCH 6/8] [Frontend] Wire ops.load_seed through the kernel load path inductor_prims.lookup_seed is a plain read from the seeds buffer, which is what the CPU backend does too. It cannot go through the op path in mlir_ops.py: CSEProxy.__getattr__ unpacks (code, ret_info) from every op and generates one CSE variable, while a load already produces one. Handle it in CSEProxy next to load, store and reduction, which are special-cased for the same reason. This is the last piece for torch.rand, torch.randn and dropout, which reach the RNG ops through lookup_seed rather than an explicit seed argument. Verified that two rand calls in one graph draw different seeds, so the seed index is honoured rather than always reading slot zero, and that a fixed manual_seed reproduces while a different one does not. --- PyTorchSimFrontend/mlir/mlir_common.py | 10 ++++++++++ PyTorchSimFrontend/mlir/mlir_ops.py | 2 ++ tests/ops/elementwise/test_pointwise.py | 18 ++++++++++++++++++ 3 files changed, 30 insertions(+) diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 9d610bdc..33bc8170 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -851,6 +851,16 @@ def load(name: str, index: sympy.Expr): result = self.load(name, index) self.cse._cache[key] = result return self.cse._cache[key] + + @staticmethod + def load_seed(name: str, offset: int): + """inductor_prims.lookup_seed: a plain read from the seeds buffer. + + Routed through the normal load path, as the CPU backend does, + rather than through __getattr__: that expects an op to hand back + (code, ret_info), while a load produces a CSE variable directly. + """ + return CSEProxy.load(name, sympy.Integer(offset)) @staticmethod def store(name, index, value, mode=None): diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index 63722bed..fa7cc4be 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -275,6 +275,8 @@ def randint64(self, seed, offset, low, high, *args, **kwargs): return res, V.kernel.var_info[res] def load_seed(self, *args, **kwargs): + # Handled in mlir_common.CSEProxy: lookup_seed is a buffer read, and the + # op path here can only return (code, ret_info), not a CSE variable. raise NotImplementedError # Special operaitons diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py index 355719ac..12b94c60 100644 --- a/tests/ops/elementwise/test_pointwise.py +++ b/tests/ops/elementwise/test_pointwise.py @@ -282,6 +282,23 @@ def f(seed): run(0, 100, "Randint64") run(-500, 500, "Randint64 negative low") run(0, 2 ** 40, "Randint64 wide range") + +def test_rand_e2e(device, size=(128, 128)): + torch._inductor.config.fallback_random = False + + # Goes through ops.load_seed, unlike the inductor_prims test which pass a + # seed in directly. Values cannot be compared agaist eager, which uses a + # different generator, so check the shape, range and that the stream is not + # constant. + def f(): + return torch.rand(size, device=device) + + clear_caches() + out = torch.compile(f, dynamic=False)().cpu() + assert out.shape == torch.Size(size) + assert (out >= 0).all() and (out < 1).all() + assert out.std() > 0.1 + print("Rand end-to-end OK") if __name__ == "__main__": device = torch.device("npu:0") @@ -311,5 +328,6 @@ def f(seed): test_rand(device) test_randn(device) test_randint64(device) + test_rand_e2e(device) From ebfe6e59b02ae8623a00d7bb59a777314c003b16 Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Thu, 30 Jul 2026 16:57:27 +0900 Subject: [PATCH 7/8] [Frontend] Fix lgamma poles, erfinv at |x| == 1 and the f16 paths Three defects the original tests could not see, plus the accuracy work that came out of chasing them. lgamma returned finite numbers at its poles. pi*x never lands exactly on a multiple of pi in f32, so sin(pi*x) comes out near 1e-07 rather than zero, log|sin| stays finite and the reflection produced plausible values where torch returns inf: lgamma(-2) was 16.01, lgamma(-100) was -350.55. Detect integer arguments on the reflection side and return infinity. erfinv inverted the sign at |x| == 1. The log drives w to +inf there, and the tail polynomial's leading coefficient is negative, so p diverged to -inf and p * x came back as -inf for erfinv(1). Handle the endpoints directly. frexp crashed for every dtype but float32. Returning NotImplemented sent Inductor to its default lowering, which calls the ops.frexp stub and dies with a bare NotImplementedError. float16 now goes through the f32 path, which is exact because its mantissa bits survive the round trip, and float64 raises something readable instead. lgamma and erfinv also left f16 operands alone. Their coefficients are fitted for single precision and ops.log still emits f16 math on an f16 operand, which put the f16 error at 2.3 and 0.17 respectively. Both promote to f32 now; f64 is left as it was. Accuracy: sin(pi*x) has period 2, so the reflection folds the argument first. Measured over random bands, the relative error drops from 1.6e-04 to 1.4e-05 on -8..0.4 -- it was outside the test tolerance before -- and from 7.8e-06 to 2.4e-07 on -200..-100. randint64 drops one 64-bit division: 2 * (x mod m) + b is already below 2m, so a conditional subtract finishes the reduction. Values are unchanged and the existing exact-comparison tests cover it. Tests grow the cases that would have caught all of this: the poles, the erfinv endpoints and the domain outside [-1, 1], a large-|x| reflection band, f16 for all three ops, and the scalar path, which no (128, 128) tensor reaches. test_result takes an equal_nan flag, since allclose treats NaN as unequal to itself and the erfinv endpoints need it. --- PyTorchSimFrontend/mlir/mlir_decomposition.py | 15 ++++- PyTorchSimFrontend/mlir/mlir_ops.py | 59 ++++++++++++++----- tests/_pytorchsim_utils.py | 4 +- tests/ops/elementwise/test_pointwise.py | 8 +++ tests/ops/elementwise/test_transcendental.py | 40 +++++++++++-- 5 files changed, 102 insertions(+), 24 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_decomposition.py b/PyTorchSimFrontend/mlir/mlir_decomposition.py index 055a3709..04270f26 100644 --- a/PyTorchSimFrontend/mlir/mlir_decomposition.py +++ b/PyTorchSimFrontend/mlir/mlir_decomposition.py @@ -403,9 +403,20 @@ def decompose_frexp(x: torch.Tensor): of two, +/-0, subnormals down to 1.4e-45, +/-inf and NaN: mantissa and exponent both match exactly. """ - # The masks below are float32 layouts; let Inductor handle anything else. + # float16 converts to float32 exactly and its mantissa bits survive the + # round trip, so route it through the f32 path instead of duplicating the + # masks for a 5-bit exponent field. + if x.dtype == torch.float16: + mantissa, exponent = decompose_frexp(x.float()) + return mantissa.half(), exponent + + # The masks below are float32 layouts. Returning NotImplemented would send + # Inductor to its default lowering, which calls ops.frexp and dies on the + # stub with a bare NotImplementedError; fail with something readable. if x.dtype != torch.float32: - return NotImplemented + raise NotImplementedError( + f"PyTorchSim frexp supports float32 and float16, got {x.dtype}" + ) bits = x.view(torch.int32) abs_bits = bits & 0x7FFFFFFF diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index fa7cc4be..dc3381e6 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -118,9 +118,11 @@ def broadcast_unflat(operand, target_size, *args, **kwargs): return format_mlir_op(op_str, shape, **kwargs), [target_size, dtype] # ---- Philox4_32-10 ------------------------------------------------- - # Matches at::Philox4_32 (ATen/core/PhiloxRNGEngine.h), which is what the - # inductor CPU and triton backends generate, so a compiled kernel - # reproduces their values bit for bit. + # Matches at::Philox4_32 (ATen/core/PhiloxRNGEngine.h), the generator behind + # normalized_rand_cpu and friends, so a compiled kernel reproduces the + # inductor CPU backend bit for bit. Not triton: tl.rand shares the Philox + # core but converts the word with where(x < 0, -x - 1, x) instead of + # masking, which differs whenever the top bit is set. _PHILOX_SA = 0xD2511F53 _PHILOX_SB = 0xCD9E8D57 _PHILOX_10A = 0x9E3779B9 @@ -270,8 +272,13 @@ def randint64(self, seed, offset, low, high, *args, **kwargs): modulus = ops.sub(high, low) halved = ops.bitwise_and(ops.bitwise_right_shift(value, one), mask63) lsb = ops.bitwise_and(value, one) + # folded is at most 2 * (m - 1) + 1, so it is already below 2m and a + # conditional subtract finishes the reduction. That drops one 64-bit + # division, which is the expensive part of this op. folded = ops.add(ops.mul(ops.mod(halved, modulus), two), lsb) - res = ops.add(ops.mod(folded, modulus), low) + reduced = ops.sub(folded, ops.where(ops.ge(folded, modulus), + modulus, ops.constant(0, "i64"))) + res = ops.add(reduced, low) return res, V.kernel.var_info[res] def load_seed(self, *args, **kwargs): @@ -635,9 +642,12 @@ def lgamma(operand, *args, **kwargs): res = ops.extractelement(val, 0) return res, V.kernel.var_info[res] - # Float-only instruction: promote non-float inputs (e.g. integers) to f32 - # to run it. Native float widths (f16/f64) are left untouched. - if not dtype.startswith("f"): + # Promote to f32 unless already f32 or f64. Integers cannot run the + # float math at all, and f16 is not accurate enough: the Lanczos + # coefficients are fitted for single precision, and ops.log still emits + # f16 math on an f16 operand, which together put the f16 error at 2.3. + # f64 is left alone -- it costs nothing and loses nothing. + if dtype not in ("f32", "f64"): operand = ops.to_dtype(operand, "f32") dtype = "f32" @@ -666,12 +676,26 @@ def lgamma(operand, *args, **kwargs): ops.log(ops.truediv(ops.mul(sqrt_2pi, ser), xr))) # Reflection term. Note this uses the original operand, not xr. - sin_pix = ops.sin(ops.mul(ops.constant(math.pi, dtype), operand)) + # sin(pi*x) has period 2, so fold the argument into [-1, 1] first: pi*x + # loses precision as |x| grows, and the reflection branch is exactly + # where large |x| ends up. Measured on the npu backend, folding cuts the + # relative error of sin(pi*x) by 16x over -8..0.4, 125x over -200..-100 + # and 1459x over -1000..-900. + xf = ops.sub(operand, + ops.mul(ops.constant(2.0, dtype), ops.round(ops.mul(operand, half)))) + sin_pix = ops.sin(ops.mul(ops.constant(math.pi, dtype), xf)) refl = ops.sub(ops.constant(math.log(math.pi), dtype), ops.log(ops.abs(sin_pix))) refl = ops.sub(refl, lg) - res = ops.where(is_reflect, refl, lg) + # Poles at x = 0, -1, -2, ...: pi*x never lands exactly on a multiple of + # pi in f32, so sin(pi*x) comes out around 1e-07 rather than zero and + # log|sin| stays finite. Without this the reflection returns a + # plausible-looking number -- lgamma(-2) came out as 16.01 -- where + # torch returns inf. + is_pole = ops.logical_and(is_reflect, ops.eq(operand, ops.floor(operand))) + res = ops.where(is_pole, ops.constant(float("inf"), dtype), + ops.where(is_reflect, refl, lg)) return res, V.kernel.var_info[res] @staticmethod @@ -983,9 +1007,10 @@ def erfinv(operand, *args, **kwargs): res = ops.extractelement(val, 0) return res, V.kernel.var_info[res] - # Float-only instruction: promote non-float inputs (e.g. integers) to f32 - # to run it. Native float widths (f16/f64) are left untouched. - if not dtype.startswith("f"): + # Promote to f32 unless already f32 or f64. The Giles coefficients are + # fitted for single precision, so an f16 operand materialises them at + # f16 and the error reaches 0.17. + if dtype not in ("f32", "f64"): operand = ops.to_dtype(operand, "f32") dtype = "f32" @@ -1000,9 +1025,6 @@ def erfinv(operand, *args, **kwargs): def const(value): return ops.constant(value, dtype) - def const(value): - return ops.constant(value, dtype) - def horner(coefs, var): acc = const(coefs[0]) for c in coefs[1:]: @@ -1028,8 +1050,13 @@ def horner(coefs, var): # one, so the inf/NaN the central branch produces at large w never # reaches the result. p = ops.where(is_central, central, tail) - res = ops.mul(p, x) + + # At |x| == 1 the log drives w to +inf, and the tail polynomial's + # leading coefficient is negative, so p diverges to -inf and p * x lands + # with the sign inverted: erfinv(1) came out as -inf instead of +inf. + res = ops.where(ops.eq(ops.abs(x), one), + ops.mul(ops.constant(float("inf"), dtype), x), res) return res, V.kernel.var_info[res] @staticmethod diff --git a/tests/_pytorchsim_utils.py b/tests/_pytorchsim_utils.py index 923f4bb7..73365a90 100644 --- a/tests/_pytorchsim_utils.py +++ b/tests/_pytorchsim_utils.py @@ -16,12 +16,12 @@ import torch -def test_result(name, out, expected, rtol=1e-4, atol=1e-4): +def test_result(name, out, expected, rtol=1e-4, atol=1e-4, equal_nan=False): """Compare ``out`` to ``expected``; exit 1 on mismatch.""" out_cpu = out.cpu() if hasattr(out, "cpu") else out expected_cpu = expected.cpu() if hasattr(expected, "cpu") else expected - if torch.allclose(out_cpu, expected_cpu, rtol=rtol, atol=atol): + if torch.allclose(out_cpu, expected_cpu, rtol=rtol, atol=atol, equal_nan=equal_nan): msg = f"|{name} Test Passed|" bar = "-" * len(msg) print(bar) diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py index 12b94c60..59a04de2 100644 --- a/tests/ops/elementwise/test_pointwise.py +++ b/tests/ops/elementwise/test_pointwise.py @@ -209,6 +209,14 @@ def frexp(a): test_result("Frexp mantissa", m, rm, equal_nan=True) test_result("Frexp exponent", e.float(), re.float()) + # float16 goes through the f32 path; mantissa bits survive the round trip + # so the result must be exact, not just close. + xh = torch.tensor([[1.5, 3.25, -2.5, 0.0, 0.5, -1.0]], dtype=torch.float16) + mh, eh = torch.compile(dynamic=False)(frexp)(xh.to(device=device)) + rmh, reh = frexp(xh.cpu()) + test_result("Frexp f16 mantissa", mh.float(), rmh.float(), rtol=0.0, atol=0.0) + test_result("Frexp f16 exponent", eh.float(), reh.float(), rtol=0.0, atol=0.0) + _NA_X = torch.tensor([[0.0, -0.0, 0.0, -0.0, 1.0, -1.0, 2.0, 3.4028235e38, -3.4028235e38, float("inf"), float("-inf"), 1.4013e-45, -1.4013e-45, 1.1754944e-38]]) diff --git a/tests/ops/elementwise/test_transcendental.py b/tests/ops/elementwise/test_transcendental.py index 4b092a35..1c725bad 100644 --- a/tests/ops/elementwise/test_transcendental.py +++ b/tests/ops/elementwise/test_transcendental.py @@ -57,15 +57,17 @@ def lgamma(a): # lgamma has poles at x = 0, -1, -2, ...; randn would land near them and # blow up the comparison. Build one tensor that covers every code path - # instead (on compile, one simulation run): + # instead (one compile, one simulation run): # rows 0:32 -> reflection branch, small positive x (x < 0.5) # rows 32:64 -> reflection branch, negative x, away from the poles - # rows 64:96 -> large x, exercises th tmp/log cancellation - # rows 96: -> the plain Lanczos path + # rows 64:96 -> large x, exercises the tmp/log cancellation + # rows 96:112 -> reflection at large |x|, where folding pi*x matters + # rows 112: -> the plain Lanczos path x = torch.empty(size).uniform_(0.5, 4.5) x[0:32].uniform_(0.1, 0.49) x[32:64].uniform_(-2.9, -2.1) x[64:96].uniform_(10.0, 100.0) + x[96:112].uniform_(-50.9, -50.1) # 반사 경로, 큰 |x| x = x.to(device=device) opt_fn = torch.compile(dynamic=False)(lgamma) @@ -73,6 +75,25 @@ def lgamma(a): out = lgamma(x.cpu()) test_result("Lgamma", res, out) + xh = torch.empty(size).uniform_(0.5, 4.5).half() + test_result("Lgamma f16", + torch.compile(dynamic=False)(lgamma)(xh.to(device)).float(), + lgamma(xh.float()), rtol=1e-2, atol=1e-2) + + # Poles: torch gives inf at 0, -1, -2, ... f32 cannot hit an exact zero of + # sin(pi*x), so this needs an explicit branch and no random band would ever + # catch a regression here. + poles = torch.tensor([0.0, -1.0, -2.0, -5.0, -20.0, -100.0]) + pole_out = torch.compile(dynamic=False)(lgamma)(poles.to(device=device)) + test_result("Lgamma poles", pole_out, lgamma(poles)) + + # Scalar path (tile_size == 1) takes a separate branch in the op and no + # (128, 128) tensor ever reaches it. + scalar = torch.tensor(2.5) + test_result("Lgamma scalar", + torch.compile(dynamic=False)(lgamma)(scalar.to(device=device)), + lgamma(scalar)) + def test_erfinv(device, size=(128, 128)): def erfinv(a): return torch.erfinv(a) @@ -81,7 +102,7 @@ def erfinv(a): # never reaches the tail branch yet still passes. Cover both explicitly: # rows 0:32 -> tail branch, positive # rows 32:64 -> tail branch, negative - # rows 64:96 -> near zero, checks p* x -> 0 + # rows 64:96 -> near zero, checks p * x -> 0 # rows 96: -> central branch x = torch.empty(size).uniform_(-0.9, 0.9) x[0:32].uniform_(0.997, 0.99999) @@ -94,6 +115,17 @@ def erfinv(a): out = erfinv(x.cpu()) test_result("Erfinv", res, out) + # |x| == 1 and |x| > 1: the polynomial branch cannot produce these on its + # own, and a band stopping at 0.99999 never reaches them. + edge = torch.tensor([1.0, -1.0, 1.5, -1.5]) + edge_out = torch.compile(dynamic=False)(erfinv)(edge.to(device=device)) + test_result("Erfinv edges", edge_out, erfinv(edge), equal_nan=True) + + scalar = torch.tensor(0.5) + test_result("Erfinv scalar", + torch.compile(dynamic=False)(erfinv)(scalar.to(device=device)), + erfinv(scalar)) + if __name__ == "__main__": import argparse From 29317ba222548c923e08797691987cc90b8cfe7b Mon Sep 17 00:00:00 2001 From: Jiyun Shin Date: Fri, 31 Jul 2026 15:22:17 +0900 Subject: [PATCH 8/8] [Frontend] Harden special ops implementations and tests --- PyTorchSimFrontend/mlir/mlir_common.py | 7 +- PyTorchSimFrontend/mlir/mlir_decomposition.py | 45 +-- PyTorchSimFrontend/mlir/mlir_ops.py | 261 ++++++------------ tests/ops/elementwise/test_pointwise.py | 110 ++++++-- tests/ops/elementwise/test_transcendental.py | 36 +++ 5 files changed, 212 insertions(+), 247 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 33bc8170..92117d99 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -854,12 +854,7 @@ def load(name: str, index: sympy.Expr): @staticmethod def load_seed(name: str, offset: int): - """inductor_prims.lookup_seed: a plain read from the seeds buffer. - - Routed through the normal load path, as the CPU backend does, - rather than through __getattr__: that expects an op to hand back - (code, ret_info), while a load produces a CSE variable directly. - """ + """Load one int64 seed from Inductor's seed buffer.""" return CSEProxy.load(name, sympy.Integer(offset)) @staticmethod diff --git a/PyTorchSimFrontend/mlir/mlir_decomposition.py b/PyTorchSimFrontend/mlir/mlir_decomposition.py index 04270f26..6ebea41f 100644 --- a/PyTorchSimFrontend/mlir/mlir_decomposition.py +++ b/PyTorchSimFrontend/mlir/mlir_decomposition.py @@ -375,44 +375,16 @@ def decompose_native_multi_head_attention( @register_decomposition(aten.frexp.Tensor) def decompose_frexp(x: torch.Tensor): - """Split ``x``into a mantissa in [0.5, 1) and an integer exponent. - - ``ops.frexp`` cannot be implemented in ``mlir_ops.py``: the CSE proxy in - ``mlir_common.py`` unpacks exactly ``(code, ret_info)`` from every op and - hands back a single CSE variable, while Inductor's ``register_frexp`` - subscripts the result as ``ops.frexp(x)[0]`` / ``[1]``. Multi-output ops are - handled a level up instead -- ``aten.sort`` does the same thing via a custom - lowering. frexp needs no template, so a decomposition into ops that already - exist is enough. - - The obvious ``floor(log2|x|) + 1`` formulation is *not* usable here. The - simulated ``log2``is not exact on powers of two (measured: 16 of 164 - mismatches over 2^-20..2^20, up to 9.5e-07), so ``floor`` slips by one and - the mantissa lands just under 1.0 instead of at 0.5. Comparing against - ``finfo.tiny`` is broken too -- subnormal operands compare as if flushed, so - a float-side subnormal test never fires. - - Bit surgery avoids both problems and is exact. For a normal float32 - ``x = (-1)^s * 1.mant * 2^(expf - 127)``, so forcing the biased exponent to - 126 yields ``m = (-1)^s * 0.1mant``in [0.5, 1) and leaves ``e = expf - 126``. - Subnormals are first scaled into the normal ranges by 2**24 and the 24 is - taken back off the exponent. Zero, the infinities and NaN pass through with - an exponent of 0, matching ``torch.frexp``. - - Verified against ``torch.frexp`` on the npu backend across normals, powers - of two, +/-0, subnormals down to 1.4e-45, +/-inf and NaN: mantissa and - exponent both match exactly. + """Decompose ``torch.frexp`` for float16 and float32. + + Float32 values are split by editing their IEEE-754 fields. Subnormals are + normalized by multiplying by 2**24 before extracting the exponent. + Float16 is routed through float32 because the conversion is exact. """ - # float16 converts to float32 exactly and its mantissa bits survive the - # round trip, so route it through the f32 path instead of duplicating the - # masks for a 5-bit exponent field. if x.dtype == torch.float16: mantissa, exponent = decompose_frexp(x.float()) return mantissa.half(), exponent - # The masks below are float32 layouts. Returning NotImplemented would send - # Inductor to its default lowering, which calls ops.frexp and dies on the - # stub with a bare NotImplementedError; fail with something readable. if x.dtype != torch.float32: raise NotImplementedError( f"PyTorchSim frexp supports float32 and float16, got {x.dtype}" @@ -424,14 +396,15 @@ def decompose_frexp(x: torch.Tensor): is_zero = abs_bits == 0 is_inf_nan = exp_field == 255 - # Subnormals must be detected on the integer side: the float comparison - # against finfo.tiny reports false for every subnormal on this target. + # Detect subnormals from the exponent bits because float comparisons may + # flush them to zero on the target. is_subnormal = (exp_field == 0) & (abs_bits != 0) scaled = torch.where(is_subnormal, x * 16777216.0, x) # 2**24 scaled_bits = scaled.view(torch.int32) - # Keep sign + mantissa, overwrite the exponent with 126 (i.e. 2**-1). + # Preserve sign and fraction, and set the biased exponent to 126 so that + # the mantissa lies in [-1, -0.5] U [0.5, 1). the exponent with 126 (i.e. 2**-1). mantissa = ((scaled_bits & 0x807FFFFF) | 0x3F000000).view(torch.float32) exponent = ((scaled_bits >> 23) & 0xFF) - 126 exponent = exponent - torch.where( diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index dc3381e6..d29edd55 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -74,7 +74,7 @@ def constant(value, src_type, *args, **kwargs): elif src_type[0] == "f": value = format(float(value), ".20f") elif src_type[0] == "i": - value = int(float(value)) + value = int(value) return format_mlir_op(f'arith.constant {value}', src_type, **kwargs), [1, src_type] @staticmethod @@ -131,24 +131,13 @@ def broadcast_unflat(operand, target_size, *args, **kwargs): @staticmethod def _u32_const(value): - """A uint32 literal as the signed i32 carrying the same bit pattern. - - arith.constant rejects values past the signed range, and every Philox - constant has its top bit set. - """ + """Materialize a uint32 bit pattern in a signed i32 container.""" signed = value - (1 << 32) if value >= (1 << 31) else value return ops.constant(signed, "i32") @staticmethod def _philox_mulhilo32(a, b): - """Full 32x32 product of two uint32 patterns, as (hi, lo) i32 halves. - - The widen to i64 must not sign-extend -- kPhiloxSA and friends have the - top bit set, so a sign-extending widen would multiply the wrong values. - Masking afterwards also stands in for a logical shift right: the repo's - bitwise_right_shift emits arith.shrsi, but the low 32 bits of an - arithmetic shift are the bits wanted either way. - """ + """Return the high and low halves of an unsigned 32x32 product.""" lo_mask = ops.constant(0xFFFFFFFF, "i64") a64 = ops.bitwise_and(ops.to_dtype(a, "i64"), lo_mask) b64 = ops.bitwise_and(ops.to_dtype(b, "i64"), lo_mask) @@ -173,12 +162,7 @@ def _philox_round(ctr, key): @staticmethod def _philox(seed32, offset32): - """Ten rounds on counter (offset, 0, 0, 0) with key (seed, 0). - - at::Philox4_32(seed, 0, offset) sets key = {seed, 0} and leaves the - counter at {offset, 0, 0, 0} after incr_n(offset). - """ - + """Run Philox4_32-10 for ``(seed, offset)``.""" cls = ExtensionOverrides zero = ops.constant(0, "i32") ctr = [offset32, zero, zero, zero] @@ -192,37 +176,23 @@ def _philox(seed32, offset32): @staticmethod def _u32_to_uniform(word): - """One Philox word -> float in [0, 1), matching uint32_to_uniform_float. - - The scale must be applied in f32; in f64 the result diverges from the - CPU backend in the last digits.""" + """Convert one Philox word to CPU-Inductor's float32 uniform format.""" masked = ops.bitwise_and(word, ops.constant(0x7FFFFFFF, "i32")) return ops.mul(ops.to_dtype(masked, "f32"), ops.constant(4.6566127342e-10, "f32")) def rand(self, seed, offset, *args, **kwargs): - """inductor_prims.random with mode="rand". - - Philox's first output word scaled into [0, 1), matching - normalized_rand_cpu: (value & 0x7FFFFFFF) * 2**-31. The scale must be - applied in f32; doing it in f64 diverges from the CPU backend in the - last couple of digits. - """ + """Lower ``inductor_prims.random(..., mode="rand")``.""" cls = ExtensionOverrides out = cls._philox(ops.to_dtype(seed, "i32"), ops.to_dtype(offset, "i32")) res = cls._u32_to_uniform(out[0]) return res, V.kernel.var_info[res] def randn(self, seed, offset, *args, **kwargs): - """inductor_prims.random with mode="randn": Box-Muller on the first two - Philox words, as randn_cpu does. - - This cannot be bit-identical to the CPU backend. randn_cpu takes the log - in float but evaluates -2.0 *, sqrt, 2.0 * M_PI and cos in double before - narrowing to float. Staying in f32 throughout lands within ~1e-01, which - is far inside the test tolerance and not worth f64 vector math here. - - u1 uses 1 - uniform so it is in (0, 1]: log(0) must not be reachable. + """Lower ``randn`` with Box-Muller on the first two Philox words. + + The transform stays in float32, so it is numerically close to the CPU + backend but is not expected to be bit-identical. """ cls = ExtensionOverrides out = cls._philox(ops.to_dtype(seed, "i32"), ops.to_dtype(offset, "i32")) @@ -235,58 +205,66 @@ def randn(self, seed, offset, *args, **kwargs): return res, V.kernel.var_info[res] def randint64(self, seed, offset, low, high, *args, **kwargs): - """inductor_prims.randint, matching randint64_cpu. - - Two Philox words are joined into a uint64, reduced modulo (high - low) - and shifted up by low. - - The reference reduction is unsigned but the repo only emits - arith.remsi, so rewrite u mod m as (2 * ((u >>> 1) mod m) + (u & 1)) - mod m. The logical shift keeps every operand non-negative, where signed - and unsigned remainder agree. Checked against the unsigned result over - 200k random (u, m) pairs plus the 64-bit edge cases. The rewrite needs - 2 * m to stay representable, i.e. high - low < 2**62. - - The logical shift is an arith.shrsi with bit 63 masked off, and the mask - is built rather than written out: ops.constant rounds integer literals - through a double, so 2**63 - 1 would come back as 2**63 and overflow - i64. 2**62 is a power of two and survives that round trip. + """Lower ``inductor_prims.randint`` over any valid int64 interval. + + Two Philox words form a uint64 value. Because the backend only provides + signed i64 remainder, reduction uses separate paths for moduli below and + above 2**63. """ cls = ExtensionOverrides out = cls._philox(ops.to_dtype(seed, "i32"), ops.to_dtype(offset, "i32")) + zero = ops.constant(0, "i64") one = ops.constant(1, "i64") two = ops.constant(2, "i64") word_mask = ops.constant(0xFFFFFFFF, "i64") - # Widening sign-extends, so mask each Philox word back to its 32 bits. + # Widening sign-extends, so mask each Philox word back to 32 bits. r0 = ops.bitwise_and(ops.to_dtype(out[0], "i64"), word_mask) r1 = ops.bitwise_and(ops.to_dtype(out[1], "i64"), word_mask) value = ops.bitwise_or( r0, ops.bitwise_left_shift(r1, ops.constant(32, "i64")) ) - two62 = ops.constant(1 << 62, "i64") - mask63 = ops.add(ops.mul(ops.sub(two62, one), two), one) # 2**63 - 1 - + # The unsigned range width is stored as an i64 bit pattern. modulus = ops.sub(high, low) - halved = ops.bitwise_and(ops.bitwise_right_shift(value, one), mask63) - lsb = ops.bitwise_and(value, one) - # folded is at most 2 * (m - 1) + 1, so it is already below 2m and a - # conditional subtract finishes the reduction. That drops one 64-bit - # division, which is the expensive part of this op. - folded = ops.add(ops.mul(ops.mod(halved, modulus), two), lsb) - reduced = ops.sub(folded, ops.where(ops.ge(folded, modulus), - modulus, ops.constant(0, "i64"))) + + two62 = ops.constant(1 << 62, "i64") + mask63 = ops.add(ops.mul(ops.sub(two62, one), two), one) + low63 = ops.bitwise_and(value, mask63) + + def add_mod_nonnegative(a, b, m): + """Compute (a + b) % modulus without signed overflow..""" + m_minus_b = ops.sub(m, b) + wrapped = ops.sub(a, m_minus_b) + plain = ops.add(a, b) + return ops.where(ops.ge(a, m_minus_b), wrapped, plain) + + # For modulus < 2**63, split off the uint64 sign bit. + low63_mod = ops.mod(low63, modulus) + two62_mod = ops.mod(two62, modulus) + two63_mod = add_mod_nonnegative(two62_mod, two62_mod, modulus) + with_top_bit = add_mod_nonnegative(two63_mod, low63_mod, modulus) + reduced_small = ops.where(ops.lt(value, zero), with_top_bit, low63_mod) + + # For modulus >= 2**63, random_u64 is below 2 * modulus, so one unsigned + # comparison and at most one subtraction are sufficient. + sign_bit = ops.constant(-(1 << 63), "i64") + value_key = ops.bitwise_xor(value, sign_bit) + modulus_key = ops.bitwise_xor(modulus, sign_bit) + unsigned_ge = ops.ge(value_key, modulus_key) + reduced_large = ops.where(unsigned_ge, ops.sub(value, modulus), value) + + reduced = ops.where( + ops.lt(modulus, zero), reduced_large, reduced_small + ) res = ops.add(reduced, low) return res, V.kernel.var_info[res] def load_seed(self, *args, **kwargs): - # Handled in mlir_common.CSEProxy: lookup_seed is a buffer read, and the - # op path here can only return (code, ret_info), not a CSE variable. + """Lowered by ``CSEProxy.load_seed`` in ``mlir_common.py``""" raise NotImplementedError - # Special operaitons @staticmethod def masked(mask, body, other, *args, tile_size=16, dtype="f32", ninf_declared=False, **kwargs): result = body() @@ -612,57 +590,40 @@ def tan(operand, *args, **kwargs): @staticmethod def lgamma(operand, *args, **kwargs): - """ - There is no MLIR operation for lgamma, so it is composed from the - Lanczos approximation (g=5, 6 terms; Numerical Recipes `gammln`): - - ln|G(x)| = -tmp + ln(sqrt(2*pi) * ser / x) - tmp = (x + 5.5) - (x + 0.5) * ln(x + 5.5) - ser = c0 + sum_k cof[k] / (x + k) - - which holds for x > 0. Inputs below 0.5 go through the reflection - formula - - ln|G(x)| = ln(pi) - ln|sin(pi*x)| - ln|G(1-x)| - - so the input is folded to 1-x up front and the series is evaluated - once instead of twice. + """Approximate ``log(abs(gamma(x)))`` with Lanczos and reflection. - g=5/N=6 is used rather than the more common g=7/N=9: at f32 the - larger g=7 coefficients (max intermediate term ~1353 vs ~51) lose - more to cancellation, so the extra double-precision accuracy does - not carry over. It also uses two fewer coefficient divisions. + Float16 inputs are evaluated in float32. Float64 is rejected because helper + operations in this backend do not preserve double precision consistently. """ tile_size, dtype = V.kernel.var_info[operand] - # Check scalar + if dtype == "f64": + raise NotImplementedError( + "PyTorchSim lgamma supports float32 and float16 only" + ) + if tile_size == 1: vec = ops.broadcast(operand, 4) val = ops.lgamma(vec) res = ops.extractelement(val, 0) return res, V.kernel.var_info[res] - - # Promote to f32 unless already f32 or f64. Integers cannot run the - # float math at all, and f16 is not accurate enough: the Lanczos - # coefficients are fitted for single precision, and ops.log still emits - # f16 math on an f16 operand, which together put the f16 error at 2.3. - # f64 is left alone -- it costs nothing and loses nothing. - if dtype not in ("f32", "f64"): + + if dtype not in ("f16", "f32", "f64"): + operand = ExtensionOverrides._signed_int_to_f32(operand, dtype) + dtype = "f32" + elif dtype != "f32": operand = ops.to_dtype(operand, "f32") dtype = "f32" half = ops.constant(0.5, dtype) one = ops.constant(1.0, dtype) - # Fold x < 0.5 into 1-x so the series only ever sees x >= 0.5. is_reflect = ops.lt(operand, half) xr = ops.where(is_reflect, ops.sub(one, operand), operand) - # tmp = (xr + 5.5) - (xr + 0.5) * ln(xr + 5.5) t = ops.add(xr, ops.constant(5.5, dtype)) tmp = ops.sub(t, ops.mul(ops.add(xr, half), ops.log(t))) - # ser = c0 + sum_k cof[k-1] / (xr + k) cof = [76.18009172947146, -86.50532032941677, 24.01409824083091, -1.231739572450155, 0.1208650973866179e-2, -0.5395239384953e-5] ser = ops.constant(1.000000000190015, dtype) @@ -670,32 +631,24 @@ def lgamma(operand, *args, **kwargs): denom = ops.add(xr, ops.constant(float(k), dtype)) ser = ops.add(ser, ops.truediv(ops.constant(c, dtype), denom)) - # lgamma(xr) = -tmp + ln(sqrt(2*pi) * ser / xr) sqrt_2pi = ops.constant(math.sqrt(2.0 * math.pi), dtype) lg = ops.add(ops.neg(tmp), ops.log(ops.truediv(ops.mul(sqrt_2pi, ser), xr))) - # Reflection term. Note this uses the original operand, not xr. - # sin(pi*x) has period 2, so fold the argument into [-1, 1] first: pi*x - # loses precision as |x| grows, and the reflection branch is exactly - # where large |x| ends up. Measured on the npu backend, folding cuts the - # relative error of sin(pi*x) by 16x over -8..0.4, 125x over -200..-100 - # and 1459x over -1000..-900. - xf = ops.sub(operand, - ops.mul(ops.constant(2.0, dtype), ops.round(ops.mul(operand, half)))) + # |sin(pi*x)| has period 1. Reducing x to [-0.5, 0.5] improves accuracu + # near negative integer poles and for large negative inputs. + xf = ops.sub(operand, ops.round(operand)) sin_pix = ops.sin(ops.mul(ops.constant(math.pi, dtype), xf)) refl = ops.sub(ops.constant(math.log(math.pi), dtype), ops.log(ops.abs(sin_pix))) refl = ops.sub(refl, lg) - # Poles at x = 0, -1, -2, ...: pi*x never lands exactly on a multiple of - # pi in f32, so sin(pi*x) comes out around 1e-07 rather than zero and - # log|sin| stays finite. Without this the reflection returns a - # plausible-looking number -- lgamma(-2) came out as 16.01 -- where - # torch returns inf. + inf = ops.constant(float("inf"), dtype) + is_inf = ops.eq(ops.abs(operand), inf) is_pole = ops.logical_and(is_reflect, ops.eq(operand, ops.floor(operand))) - res = ops.where(is_pole, ops.constant(float("inf"), dtype), - ops.where(is_reflect, refl, lg)) + res = ops.where(is_inf, inf, + ops.where(is_pole, inf, + ops.where(is_reflect, refl, lg))) return res, V.kernel.var_info[res] @staticmethod @@ -975,46 +928,25 @@ def erfc(operand, *args, **kwargs): @staticmethod def erfinv(operand, *args, **kwargs): - """ - There is no MLIR operation for erfinv, so it is composed from the - division-free polynomial approximation in - - M. Giles, "Approximating the erfinv function", - GPU Computing Gems Jade Edition ch. 10 (single-precision form). - - With w = -ln((1-x)*(1+x)) the inverse splits into two Horner - polynomials -- a central one for w < 5 and a tail one in sqrt(w) - past that: - - w < 5: p = horner(CENTRAL, w - 2.5) - w >= 5: p = horner(TAIL, sqrt(w) - 3) - erfinv(x) = p * x - - Max error ~5.6e-07 across the whole domain at f32, near the f32 - epsilon itself. The coefficients are fitted for single precision; - a double-precision set would cost more without helping here. - - The edge cases need no extra branch, they fall out of the formula: - |x| > 1 makes the log NaN, |x| == 1 drives w to +inf so the tail - branch yields +/-inf, and x == 0 gives p * 0 == 0. - """ + """Approximate ``erfinv`` with the single-precision Giles polynomials.""" tile_size, dtype = V.kernel.var_info[operand] - # Check scalar + if dtype == "f64": + raise NotImplementedError( + "PyTorchSim erfinv supports float32 and float16 only" + ) + if tile_size == 1: vec = ops.broadcast(operand, 4) val = ops.erfinv(vec) res = ops.extractelement(val, 0) return res, V.kernel.var_info[res] - # Promote to f32 unless already f32 or f64. The Giles coefficients are - # fitted for single precision, so an f16 operand materialises them at - # f16 and the error reaches 0.17. - if dtype not in ("f32", "f64"): + if dtype != "f32": operand = ops.to_dtype(operand, "f32") dtype = "f32" - # Horner coefficients, highest order first. + # Horner coefficients CENTRAL = [2.81022636e-08, 3.43273939e-07, -3.5233877e-06, -4.39150654e-06, 0.00021858087, -0.00125372503, -0.00417768164, 0.246640727, 1.50140941] @@ -1034,27 +966,17 @@ def horner(coefs, var): x = operand one = const(1.0) - # w = -ln((1-x)*(1+x)). Written as (1-x)*(1+x) rather than 1-x*x: the - # latter cancels badly as |x| approaches 1, which is exactly the region - # the tail branch exists to handle. product = ops.mul(ops.sub(one, x), ops.add(one, x)) w = ops.neg(ops.log(product)) - # w >= 5 is |x| >= 0.996625. is_central = ops.lt(w, const(5.0)) central = horner(CENTRAL, ops.sub(w, const(2.5))) tail = horner(TAIL, ops.sub(ops.sqrt(w), const(3.0))) - # Both polynomials run on every lane and arith.select drops the unused - # one, so the inf/NaN the central branch produces at large w never - # reaches the result. p = ops.where(is_central, central, tail) res = ops.mul(p, x) - # At |x| == 1 the log drives w to +inf, and the tail polynomial's - # leading coefficient is negative, so p diverges to -inf and p * x lands - # with the sign inverted: erfinv(1) came out as -inf instead of +inf. res = ops.where(ops.eq(ops.abs(x), one), ops.mul(ops.constant(float("inf"), dtype), x), res) return res, V.kernel.var_info[res] @@ -1131,24 +1053,7 @@ def log1p(operand, *args, **kwargs): @staticmethod def nextafter(operand1, operand2, *args, **kwargs): - """Step ``operand1`` one representable value toward ``operand2``. - - IEEE 754 sign-magnitude patterns increase monotonically as the value - moves away from zero, so one ulp is a single integer step on the - bitcast: +1 away from zero, -1 toward it. ``(y > x) == (x > 0)`` picks - the direction and holds for both signs of x. - - Leaving +/-0 the neighbor is the smallest subnormal carrying y's sign. - It is assembled from bits rather than written as a literal: a subnormal - constant does not survive materialisation on this target. - - NaN propagates through ``x + y``, which is NaN exactly when either - operand is, so no NaN literal is needed either. - - Verified bit-exact against ``torch.nextafter`` on the npu backend over - random pairs, +/-0, +/-FLT_MAX. +/-inf, the smallest subnormals, equal - inputs and NaN. - """ + """Step ``operand1`` one representable value toward ``operand2``""" tile_size, ret_type, x, y = ExtensionOverrides.binary_elementwise_common( operand1, operand2 ) @@ -1157,21 +1062,15 @@ def nextafter(operand1, operand2, *args, **kwargs): width = mlir_common.MLIR_TO_BIT[ret_type] itype = f"i{width}" - abs_mask = (1 << (width - 1)) - 1 # everything but the sign bit sign_mask = -(1 << (width - 1)) # the sign bit, as a signed int - # ops.to_dtype_bitcast follows the Inductor protocol: (x, dtype, src_dtype), - # both torch dtypes. float_dt = mlir_common.MLIR_TO_DTYPE[ret_type] int_dt = mlir_common.MLIR_TO_DTYPE[itype] bx = ops.to_dtype_bitcast(x, int_dt, float_dt) by = ops.to_dtype_bitcast(y, int_dt, float_dt) - is_zero = ops.eq( - ops.bitwise_and(bx, ops.constant(abs_mask, itype)), - ops.constant(0, itype), - ) + is_zero = ops.eq(x, ops.constant(0.0, ret_type)) is_eq = ops.eq(x, y) is_nan = ops.logical_or(ops.isnan(x), ops.isnan(y)) diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py index 59a04de2..cac58671 100644 --- a/tests/ops/elementwise/test_pointwise.py +++ b/tests/ops/elementwise/test_pointwise.py @@ -193,9 +193,6 @@ def test_frexp(device, size=(128, 128)): def frexp(a): return torch.frexp(a) - # Cover every branch of the decomposition: normals, powers of two (where a - # log2-based version slips), zero, subnormals (integer-side detection) and - # the inf/NaN passthrough. special = torch.tensor([0.0, -0.0, 1.0, 4.0, 0.5, -2.0 ** 20, 1.1754944e-38, 1e-40, 5e-44, 1.4e-45, 3.4028235e38, float("inf"), float("-inf"), float("nan")]) @@ -206,15 +203,28 @@ def frexp(a): opt_fn = torch.compile(dynamic=False)(frexp) m, e = opt_fn(x) rm, re = frexp(x.cpu()) - test_result("Frexp mantissa", m, rm, equal_nan=True) - test_result("Frexp exponent", e.float(), re.float()) - # float16 goes through the f32 path; mantissa bits survive the round trip - # so the result must be exact, not just close. - xh = torch.tensor([[1.5, 3.25, -2.5, 0.0, 0.5, -1.0]], dtype=torch.float16) + test_result("Frexp mantissa", m, rm, rtol=0.0, atol=0.0, equal_nan=True) + test_result("Frexp exponent", e.float(), re.float(), rtol=0.0, atol=0.0) + + finfo16 = torch.finfo(torch.float16) + xh = torch.tensor([[ + 1.5, 3.25, -2.5, 0.0, -0.0, 0.5, -1.0, + finfo16.smallest_normal, + finfo16.smallest_normal / 2, + 2.0 ** -24, + float("inf"), float("-inf"), float("nan"), + ]], dtype=torch.float16) mh, eh = torch.compile(dynamic=False)(frexp)(xh.to(device=device)) rmh, reh = frexp(xh.cpu()) - test_result("Frexp f16 mantissa", mh.float(), rmh.float(), rtol=0.0, atol=0.0) + test_result( + "Frexp f16 mantissa", + mh.float(), + rmh.float(), + rtol=0.0, + atol=0.0, + equal_nan=True, + ) test_result("Frexp f16 exponent", eh.float(), reh.float(), rtol=0.0, atol=0.0) _NA_X = torch.tensor([[0.0, -0.0, 0.0, -0.0, 1.0, -1.0, 2.0, @@ -225,8 +235,6 @@ def frexp(a): 0.0, 0.0, 0.0]]) def test_nextafter(device): - # One ulp apart, so the default 1e-4 tolerance would pass even if the op - # returned x unchanged. Compare exactly instead. run_op("Nextafter", device, torch.nextafter, lambda r, c: (torch.randn(r, c), torch.randn(r, c)), cases=[ @@ -237,13 +245,73 @@ def test_nextafter(device): ], rtol=0.0, atol=0.0) + def check_dtype(label, x, y): + def nextafter(a, b): + return torch.nextafter(a, b) + + clear_caches() + npu = torch.compile(dynamic=False)(nextafter)( + x.to(device=device), y.to(device=device) + ) + cpu = nextafter(x, y) + test_result( + label, npu, cpu, rtol=0.0, atol=0.0, equal_nan=True + ) + + f16_x = torch.tensor( + [[0.0, -0.0, 1.0, -1.0, torch.finfo(torch.float16).max, + float("inf"), float("-inf"), float("nan")]], + dtype=torch.float16, + ) + f16_y = torch.tensor( + [[1.0, -1.0, 2.0, -2.0, float("inf"), + 0.0, 0.0, 1.0]], + dtype=torch.float16, + ) + check_dtype("Nextafter f16", f16_x, f16_y) + + f64_x = torch.tensor( + [[0.0, -0.0, 1.0, -1.0, torch.finfo(torch.float64).max, + float("inf"), float("-inf"), float("nan")]], + dtype=torch.float64, + ) + f64_y = torch.tensor( + [[1.0, -1.0, 2.0, -2.0, float("inf"), + 0.0, 0.0, 1.0]], + dtype=torch.float64, + ) + check_dtype("Nextafter f64", f64_x, f64_y) + +def test_load_seed(device): + from torch._inductor import inductor_prims + + indices = (3, 0, 4, 1) + + def f(seeds): + return torch.stack(tuple( + inductor_prims.lookup_seed(seeds, index) for index in indices + )) + + seeds = torch.tensor( + [12345, -7, 987654321, 2 ** 40, -(2 ** 40)], + dtype=torch.int64, + ) + + clear_caches() + npu = torch.compile(f, dynamic=False)(seeds.to(device=device)) + expected = seeds[list(indices)] + test_result( + "LoadSeed offsets", + npu, + expected, + rtol=0.0, + atol=0.0, + ) + def test_rand(device, size=(128, 128)): from torch._inductor import inductor_prims torch._inductor.config.fallback_random = False - # Compare against the inductor CPU backed, not eager: both go through - # inductor_prims.random, so the same Philox seed must give the same bits. - # Passing the seed as a graph input keeps ops.load_seed out of the picture. def f(seed): return inductor_prims.random(list(size), seed, "rand") @@ -266,9 +334,6 @@ def f(seed): npu = torch.compile(f, dynamic=False)(seed.to(device=device)) clear_caches() cpu = torch.compile(f, dynamic=False)(seed) - # Not exact: randn_cpu evaluates the Box-Muller tail in double, we stay in - # f32. Measured max deviation ~1e-06, so the default tolerance still catches - # any real error (a wrong generator differs by 0(1), not by 1e-06). test_result("Randn", npu, cpu) def test_randint64(device, size=(128, 128)): @@ -283,21 +348,17 @@ def f(seed): npu = torch.compile(f, dynamic=False)(seed.to(device=device)) clear_caches() cpu = torch.compile(f, dynamic=False)(seed) - # Integers: compare exactly. A loose tolerance would hide an off-by-one - # in the modulo rewrite. - test_result(label, npu.float(), cpu.float(), rtol=0.0, atol=0.0) + test_result(label, npu, cpu, rtol=0.0, atol=0.0) run(0, 100, "Randint64") run(-500, 500, "Randint64 negative low") run(0, 2 ** 40, "Randint64 wide range") + run(0, 3 * (2 ** 61), "Randint64 >2^62 range") + run(-(2 ** 63), (2 ** 63) - 1, "Randint64 near full i64 range") def test_rand_e2e(device, size=(128, 128)): torch._inductor.config.fallback_random = False - # Goes through ops.load_seed, unlike the inductor_prims test which pass a - # seed in directly. Values cannot be compared agaist eager, which uses a - # different generator, so check the shape, range and that the stream is not - # constant. def f(): return torch.rand(size, device=device) @@ -333,6 +394,7 @@ def f(): test_atan2(device) test_frexp(device) test_nextafter(device) + test_load_seed(device) test_rand(device) test_randn(device) test_randint64(device) diff --git a/tests/ops/elementwise/test_transcendental.py b/tests/ops/elementwise/test_transcendental.py index 1c725bad..d6bbf0d7 100644 --- a/tests/ops/elementwise/test_transcendental.py +++ b/tests/ops/elementwise/test_transcendental.py @@ -87,6 +87,25 @@ def lgamma(a): pole_out = torch.compile(dynamic=False)(lgamma)(poles.to(device=device)) test_result("Lgamma poles", pole_out, lgamma(poles)) + # +/-inf must return +inf; NaN must propagate. + nonfinite = torch.tensor([float("inf"), float("-inf"), float("nan")]) + nonfinite_out = torch.compile(dynamic=False)(lgamma)(nonfinite.to(device=device)) + test_result("Lgamma nonfinite", nonfinite_out, lgamma(nonfinite), equal_nan=True) + + # Values one f32 ULP away from negative integer poles exercise the + # reflection argument reduction. Returning a plausible but inaccurate + # finite value here is easy when sin(pi*x) is evaluated near +/-pi. + neg_one = torch.tensor(-1.0, dtype=torch.float32) + neg_two = torch.tensor(-2.0, dtype=torch.float32) + near_poles = torch.stack([ + torch.nextafter(neg_one, torch.tensor(float("-inf"))), + torch.nextafter(neg_one, torch.tensor(float("inf"))), + torch.nextafter(neg_two, torch.tensor(float("-inf"))), + torch.nextafter(neg_two, torch.tensor(float("inf"))), + ]) + near_pole_out = torch.compile(dynamic=False)(lgamma)(near_poles.to(device=device)) + test_result("Lgamma near poles", near_pole_out, lgamma(near_poles)) + # Scalar path (tile_size == 1) takes a separate branch in the op and no # (128, 128) tensor ever reaches it. scalar = torch.tensor(2.5) @@ -125,6 +144,23 @@ def erfinv(a): test_result("Erfinv scalar", torch.compile(dynamic=False)(erfinv)(scalar.to(device=device)), erfinv(scalar)) + + # The implementation uses the single-precision Giles coefficients. f64 must + # fail explicitly instead of returning a badly inaccurate value near |x|=1. + x64 = torch.nextafter( + torch.tensor([1.0], dtype=torch.float64), + torch.tensor([0.0], dtype=torch.float64), + ) + try: + torch.compile(dynamic=False)(erfinv)(x64.to(device=device)) + except Exception as exc: + if "PyTorchSim erfinv supports float32 and float16 only" not in str(exc): + raise + print("--------------------------") + print("|Erfinv f64 reject Test Passed|") + print("--------------------------") + else: + raise AssertionError("Erfinv f64 input must be rejected") if __name__ == "__main__": import argparse