diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 9d610bdc..92117d99 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -851,6 +851,11 @@ 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): + """Load one int64 seed from Inductor's seed buffer.""" + return CSEProxy.load(name, sympy.Integer(offset)) @staticmethod def store(name, index, value, mode=None): diff --git a/PyTorchSimFrontend/mlir/mlir_decomposition.py b/PyTorchSimFrontend/mlir/mlir_decomposition.py index f9ddbc31..6ebea41f 100644 --- a/PyTorchSimFrontend/mlir/mlir_decomposition.py +++ b/PyTorchSimFrontend/mlir/mlir_decomposition.py @@ -373,6 +373,48 @@ def decompose_native_multi_head_attention( else: return (output, None) +@register_decomposition(aten.frexp.Tensor) +def decompose_frexp(x: torch.Tensor): + """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. + """ + if x.dtype == torch.float16: + mantissa, exponent = decompose_frexp(x.float()) + return mantissa.half(), exponent + + if x.dtype != torch.float32: + raise NotImplementedError( + f"PyTorchSim frexp supports float32 and float16, got {x.dtype}" + ) + + 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 + # 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) + + # 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( + 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/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index aeb3ce26..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 @@ -117,19 +117,154 @@ 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), 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 + _PHILOX_10B = 0xBB67AE85 + _PHILOX_ROUNDS = 10 + + @staticmethod + def _u32_const(value): + """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): + """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) + 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): + """Run Philox4_32-10 for ``(seed, offset)``.""" + 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): + """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): + """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): + """Lower ``randn`` with Box-Muller on the first two Philox words. - def rand(self, *args, **kwargs): - raise NotImplementedError + 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")) + 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): + """Lower ``inductor_prims.randint`` over any valid int64 interval. - def randn(self, *args, **kwargs): - raise NotImplementedError + 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 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")) + ) + + # The unsigned range width is stored as an i64 bit pattern. + modulus = ops.sub(high, low) + + 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 randint64(self, *args, **kwargs): + def load_seed(self, *args, **kwargs): + """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() @@ -455,7 +590,66 @@ def tan(operand, *args, **kwargs): @staticmethod def lgamma(operand, *args, **kwargs): - raise NotImplementedError + """Approximate ``log(abs(gamma(x)))`` with Lanczos and reflection. + + 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] + + 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] + + 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) + + is_reflect = ops.lt(operand, half) + xr = ops.where(is_reflect, ops.sub(one, operand), operand) + + t = ops.add(xr, ops.constant(5.5, dtype)) + tmp = ops.sub(t, ops.mul(ops.add(xr, half), ops.log(t))) + + 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)) + + 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))) + + # |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) + + 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_inf, inf, + ops.where(is_pole, inf, + ops.where(is_reflect, refl, lg))) + return res, V.kernel.var_info[res] @staticmethod def erf(operand, *args, **kwargs): @@ -734,10 +928,62 @@ def erfc(operand, *args, **kwargs): @staticmethod def erfinv(operand, *args, **kwargs): - raise NotImplementedError + """Approximate ``erfinv`` with the single-precision Giles polynomials.""" + tile_size, dtype = V.kernel.var_info[operand] + + 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] + + if dtype != "f32": + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + # 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] + 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 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) + + product = ops.mul(ops.sub(one, x), ops.add(one, x)) + w = ops.neg(ops.log(product)) + + 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))) + + p = ops.where(is_central, central, tail) + res = ops.mul(p, x) + + 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 def frexp(operand, *args, **kwargs): + """Implemented in mlir_decomposition.py.""" raise NotImplementedError @staticmethod @@ -807,7 +1053,42 @@ def log1p(operand, *args, **kwargs): @staticmethod def nextafter(operand1, operand2, *args, **kwargs): - raise NotImplementedError + """Step ``operand1`` one representable value toward ``operand2``""" + 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}" + sign_mask = -(1 << (width - 1)) # the sign bit, as a signed int + + 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(x, ops.constant(0.0, ret_type)) + 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/_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 74e6656d..cac58671 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,186 @@ 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) + + 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, 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, + 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, + 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): + 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) + + 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 + + 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) + 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) + 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 + + 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") @@ -211,4 +391,13 @@ 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) + test_nextafter(device) + test_load_seed(device) + test_rand(device) + test_randn(device) + test_randint64(device) + test_rand_e2e(device) + + diff --git a/tests/ops/elementwise/test_transcendental.py b/tests/ops/elementwise/test_transcendental.py index c3a2ee0f..d6bbf0d7 100644 --- a/tests/ops/elementwise/test_transcendental.py +++ b/tests/ops/elementwise/test_transcendental.py @@ -51,6 +51,117 @@ 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 (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 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) + res = opt_fn(x) + 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)) + + # +/-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) + 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) + + # 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) + + # |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)) + + # 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 @@ -64,4 +175,6 @@ 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) + test_erfinv(device) \ No newline at end of file