exact-solve / README.md
phanerozoic's picture
Card: standardized form with hero
12ead3f verified
|
Raw
History Blame
6.16 kB
metadata
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, fp64-emu, and 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.

Digit rain locks into the exact integer solution of a Pascal system while fp64's wrong values are struck out

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

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.