| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # exact-solve |
|
|
| exact-solve computes the exact rational solution of an integer linear system |
| `A x = b` on INT8 tensor cores, and is loadable through `kernels`. `A` |
| (`[n, n]`) and `b` (`[n, c]`) are integer; the solution is returned as |
| `fractions.Fraction`. It uses high-order Dixon p-adic lifting on the residue |
| (CRT) engine shared with |
| [exact-gemm](https://huggingface.co/kernels/phanerozoic/exact-gemm), |
| [fp64-emu](https://huggingface.co/kernels/phanerozoic/fp64-emu), and |
| [fp128-emu](https://huggingface.co/kernels/phanerozoic/fp128-emu). The |
| reference baseline is FLINT `fmpq_mat.solve`, against which it measures up to |
| 3.6x at `n = 1024`. |
|
|
| Linear systems with condition numbers past ~10^16 are unsolvable in double |
| precision: every returned component can be pure noise. Exact rational solving |
| sidesteps conditioning entirely, returning the true numerators and |
| denominators, and has been a CPU library operation. This kernel runs the lift |
| on GPU tensor cores, so systems that floating point cannot touch, and that |
| occupy a CPU for seconds, solve exactly in fractions of a second. |
|
|
|  |
|
|
| *The order-30 Pascal system, condition ~10^32. fp64 returns every component |
| wrong, off by up to 10^12 and in sign; exact-solve returns the true 18-digit |
| integer solution (Pascal is unimodular, so the rationals are integers), |
| verified `A @ x == b` in integer arithmetic, in 4 ms.* |
|
|
| ## Usage |
|
|
| ```python |
| import torch |
| from kernels import get_kernel |
| |
| es = get_kernel("phanerozoic/exact-solve", version=1, trust_remote_code=True) |
| |
| A = torch.randint(-10**6, 10**6, (256, 256), dtype=torch.int64, device="cuda") |
| b = torch.randint(-10**6, 10**6, (256, 1), dtype=torch.int64, device="cuda") |
| X = es.solve(A, b) # [256][1] list of fractions.Fraction (exact) |
| ``` |
|
|
| `version` selects the release branch; `trust_remote_code` is required by |
| `kernels` for publishers without the trusted-publisher mark. `b` may have |
| multiple columns. `solve` returns `None` if `A` is singular over the |
| rationals. `gmpy2` accelerates the reconstruction and is recommended. |
|
|
| ## API |
|
|
| | Symbol | Purpose | |
| |---|---| |
| | `solve(A, b)` | exact rational solution of `A x = b`: `[n][c]` nested lists of `fractions.Fraction`; `None` if `A` is singular over the rationals | |
|
|
| ## Method |
|
|
| Dixon (1982) recovers the exact solution from a p-adic expansion: with |
| `B = A^{-1} mod P`, `r_0 = b`, and `x_i = B r_i mod P`, |
| `r_{i+1} = (r_i - A x_i) / P`, the digits `x_i` reconstruct `x` once the |
| precision `P^m` exceeds twice the Hadamard bound on numerator and denominator. |
|
|
| The modulus is a squarefree product `P = prod(q_j)` over primes `q_j < 256`. |
| Each step's `A^{-1} mod q_j` applied to the residual is then a native INT8 |
| tensor-core GEMM, one per prime, and balanced-Garner reconstruction assembles |
| the wide base-`P` digit directly. A product of primes in place of a single |
| word prime divides the sequential step count by the prime count (~460 |
| word-prime steps at `n = 256` become ~44), which is what puts the lift on the |
| tensor cores rather than in a long scalar recurrence. |
|
|
| Three parts run on the GPU: a batched modular inverse (`A^{-1} mod q_j` for |
| every prime by Gauss-Jordan over the augmented matrix, one grid per |
| elimination step spanning all primes, Barrett reduction in place of hardware |
| modulo); the lift (per-step per-prime products as one batched INT8 GEMM, CRT |
| reconstruction of the wide digit, and the residual update in one limb as |
| `(r_i - A x_i) P^{-1} mod 2^64`, `P` odd); and the final rational |
| reconstruction (common denominator from a small sample of components, then one |
| modular multiply per numerator; this CPU stage uses `gmpy2` and a process pool |
| when available). |
|
|
| Range: entries satisfy `n * max|A_ij| < 2^63`, so the residual is a single |
| signed 64-bit limb throughout; `K = n <= 2^17` (the INT8 accumulator bound). |
|
|
| ## Measured |
|
|
| Median wall-clock of the full solve against FLINT `fmpq_mat.solve` on the same |
| host, 24-bit integer entries. Ratio > 1 is faster than FLINT. |
|
|
| RTX PRO 6000 Blackwell / NVIDIA H200: |
|
|
| | n | c | exact-solve | FLINT | ratio | |
| |---|---|---|---|---| |
| | 128 | 1 | 33 / 33 ms | 20 / 18 ms | 0.62 / 0.61 | |
| | 256 | 1 | 96 / 97 ms | 144 / 142 ms | 1.52 / 1.46 | |
| | 512 | 1 | 413 / 437 ms | 1074 / 1089 ms | 2.60 / 2.49 | |
| | 1024 | 1 | 2411 / 2235 ms | 8122 / 8141 ms | 3.37 / 3.64 | |
| | 256 | 16 | 446 / 448 ms | 943 / 938 ms | 2.12 / 2.09 | |
| | 512 | 16 | 2615 / 2624 ms | 6637 / 6749 ms | 2.54 / 2.57 | |
|
|
| The relative advantage grows with `n`: FLINT's single-word Dixon and the GPU |
| lift both scale as `O(n^3)` in the exponent, but the tensor-core lift carries |
| a smaller constant and the reconstruction parallelizes across cores. For a |
| small single right-hand side (`n = 128`) FLINT's ~18 ms is below the fixed |
| cost of the GPU inverse, lift, and reconstruction; the single-RHS crossover is |
| near `n = 256`. Verified additionally on L4, where the same ordering holds |
| with smaller multipliers. |
|
|
| ## Correctness |
|
|
| The solution is checked bitwise against an exact `Fraction` |
| Gaussian-elimination oracle on small systems, against FLINT `fmpq_mat.solve` |
| on every timed case, and by the integer identity `A @ N == d * b` (numerators |
| `N`, common denominator `d`) at larger sizes; a planted integer solution and a |
| scaled-Hilbert system with dense denominators are included. Results are |
| deterministic. |
|
|
| ## Requirements and limits |
|
|
| - NVIDIA GPU with compute capability 8.0+ and torch with INT8 tensor-core |
| GEMM. |
| - Integer `A` `[n, n]`, `b` `[n, c]`, with `n * max|A_ij| < 2^63` and |
| `n <= 2^17`. |
| - Working memory scales with the prime count and `n`; the batched inverse |
| holds an augmented `[E, n, 2n]` int32 buffer. |
|
|
| ## References |
|
|
| Dixon, "Exact solution of linear equations using p-adic expansions" |
| (Numerische Mathematik 40, 1982); Ozaki, Ogita, Oishi, Rump 2012 (Numerical |
| Algorithms 59); Ozaki, Uchino, Imamura 2025 (Ozaki Scheme II, |
| arXiv:2504.08009); Hart et al., FLINT (Fast Library for Number Theory). |
|
|
| ## License |
|
|
| Apache-2.0. |
|
|
