[Frontend] Implement special and random operations - #307
Open
jyshin1201 wants to merge 8 commits into
Open
Conversation
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.
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.
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.
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.
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.
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.
jyshin1201
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Target branch