Uploaded using `kernel-builder`.

by phanerozoic - opened 16 days ago

base: refs/heads/v1

←

from: refs/pr/1

Discussion Files changed

+387

-0

Files changed (6) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +265 -0
build/torch213-cxx11-cu130-aarch64-linux/_exact_solve_cuda_31a8fa7_dirty.abi3.so +3 -0
build/torch213-cxx11-cu130-aarch64-linux/_exact_solve_recon.py +45 -0
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-aarch64-linux/exact_solve/__init__.py +26 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +39 -0

build/torch213-cxx11-cu130-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,265 @@

+"""exact-solve: exact rational linear-system solver on INT8 tensor cores.
+Solves ``A x = b`` exactly over the rationals for integer ``A`` (``[n, n]``) and
+integer ``b`` (``[n, c]``) by high-order Dixon p-adic lifting on the residue
+(CRT) engine, and returns the exact solution as ``fractions.Fraction``.
+The lift uses a squarefree modulus ``P = prod(q_j)`` over primes ``q_j < 256``:
+each step's ``A^{-1} mod q_j`` matmuls are native INT8 tensor-core GEMMs (one per
+prime), CRT-reconstructed into a wide base-``P`` digit; the single-limb residual
+update is ``(r - A x_i) * P^{-1} mod 2^64`` (``P`` odd), which needs only
+``A x_i mod 2^64``. The batched modular inverse and the lift run on the GPU; the
+final rational reconstruction runs on the CPU (accelerated by ``gmpy2`` and a
+process pool, ``fork`` or ``spawn``, when available).
+    solve(A, b)              -> [n][c] list of fractions.Fraction (exact)
+Reference: Dixon, "Exact solution of linear equations using p-adic expansions"
+(Numerische Mathematik 40, 1982).
+"""
+import math
+import os
+import sys
+from fractions import Fraction
+import torch
+from ._ops import ops
+# odd primes < 256 (descending); 2 is excluded so P is odd and invertible mod 2^64
+_ODD_PRIMES = [251, 241, 239, 233, 229, 227, 223, 211, 199, 197, 193, 191, 181,
+               179, 173, 167, 163, 157, 151, 149, 139, 137, 131, 127, 113, 109,
+               107, 103, 101, 97, 89, 83, 79, 73, 71, 67, 61, 59, 53, 47, 43, 41,
+               37, 31, 29, 23, 19, 17, 13, 11, 7, 5, 3]
+_LMAX = 4
+_STRIDE = 2 * _LMAX + 1
+_M64 = (1 << 64) - 1
+# ---------------------------------------------------------------- tables / ops
+def _crt_tables(primes, device):
+    P = len(primes)
+    inv = [0] * (P * P)
+    for a in range(P):
+        for b in range(a + 1, P):
+            inv[a * P + b] = pow(primes[a] % primes[b], -1, primes[b])
+    mu = [(1 << 62) // p for p in primes]
+    pow32 = [0] * (P * _STRIDE)
+    for pi, p in enumerate(primes):
+        for k in range(_STRIDE):
+            pow32[pi * _STRIDE + k] = pow(2, 32 * k, p)
+    return (torch.tensor(primes, dtype=torch.int32, device=device),
+            torch.tensor(inv, dtype=torch.int32, device=device),
+            torch.tensor(mu, dtype=torch.int64, device=device),
+            torch.tensor(pow32, dtype=torch.int32, device=device))
+def _pad8(x):
+    return max(24, ((x + 7) // 8) * 8)
+def _extract(Xi64, rows_pad, cols_pad, primes, mu, pow32):
+    rows, cols = Xi64.shape
+    if rows == rows_pad and cols == cols_pad:
+        Y = Xi64.contiguous()
+    else:
+        Y = torch.zeros(rows_pad, cols_pad, dtype=torch.int64, device=Xi64.device)
+        Y[:rows, :cols] = Xi64
+    planes = torch.empty(primes.numel(), rows_pad, cols_pad, dtype=torch.int8, device=Xi64.device)
+    ops.exact_solve_extract(planes, Y.unsqueeze(0).contiguous(), primes, mu, pow32)
+    return planes
+def _matmul(A_planes, B_planes, Lc, primes, inv, mu):
+    C = torch.empty(Lc, A_planes.size(1), B_planes.size(1), dtype=torch.int64, device=A_planes.device)
+    ops.exact_solve_matmul(C, A_planes, B_planes, primes, inv, mu)
+    return C
+def _batched_modinv(A, primes):
+    n = A.size(0)
+    E = primes.numel()
+    Bout = torch.empty(E, n, n, dtype=torch.int8, device=A.device)
+    work = torch.empty((int)(E) * n * 2 * n, dtype=torch.int32, device=A.device)
+    singular = torch.zeros(E, dtype=torch.int32, device=A.device)
+    ops.exact_solve_batched_modinv(Bout, A, primes, work, singular)
+    return Bout, singular
+def _i64_matmul(A, X):
+    out = torch.empty(A.size(0), X.size(1), dtype=torch.int64, device=A.device)
+    ops.exact_solve_i64_matmul(out, A, X)
+    return out
+def _padic(r, axlow, Pinv64):
+    out = torch.empty_like(r)
+    ops.exact_solve_padic_residual(out, r, axlow, Pinv64)
+    return out
+def _gpu_recon(digits, Pl, Pm, Rr, Pmhalf, dl, LO):
+    out = torch.empty(LO, digits.size(2), dtype=torch.int64, device=digits.device)
+    ops.exact_solve_gpu_recon(out, digits, Pl, Pm, Rr, Pmhalf, dl)
+    return out
+def _s64(u):
+    return u - (1 << 64) if u >> 63 else u
+def _to_limbs(x, L):
+    return torch.tensor([_s64((x >> (64 * i)) & _M64) for i in range(L)], dtype=torch.int64)
+def _from_limbs(col, LO):
+    u = 0
+    for i in range(LO):
+        u |= (int(col[i]) & _M64) << (64 * i)
+    return u - (1 << (64 * LO)) if u >> (64 * LO - 1) else u
+# ---------------------------------------------------------------- reconstruction
+# Reconstruction lives in the torch-free _exact_solve_recon module (imported
+# top-level, its dir on sys.path) so spawn-pool workers load it without torch.
+_RECON = None
+def _recon():
+    global _RECON
+    if _RECON is None:
+        import importlib
+        d = os.path.dirname(os.path.abspath(__file__))
+        if d not in sys.path:
+            sys.path.append(d)
+        _RECON = importlib.import_module("_exact_solve_recon")
+    return _RECON
+_POOL = None
+def _pool():
+    global _POOL
+    if _POOL is None:
+        import multiprocessing as mp
+        from concurrent.futures import ProcessPoolExecutor
+        _recon()  # place the backend dir on sys.path before any worker spawns
+        ctx = mp.get_context("fork" if hasattr(os, "fork") else "spawn")
+        _POOL = ProcessPoolExecutor(max_workers=min(int(os.cpu_count() or 8), 32), mp_context=ctx)
+    return _POOL
+def _hadamard_bits(A, B):
+    n = len(A)
+    col_sq = [sum(int(A[i][j]) ** 2 for i in range(n)) for j in range(n)]
+    logcol = [0.5 * math.log2(max(s, 1)) for s in col_sq]
+    logden = sum(logcol)
+    b_sq_max = max(sum(int(B[i][k]) ** 2 for i in range(n)) for k in range(len(B[0])))
+    lognum = sum(logcol) - min(logcol) + 0.5 * math.log2(max(b_sq_max, 1))
+    return lognum, logden
+def _to_int_lists(T, rows, cols):
+    if isinstance(T, torch.Tensor):
+        t = T.detach().cpu().tolist()
+    else:
+        t = [list(r) for r in T]
+    return [[int(t[i][j]) for j in range(cols)] for i in range(rows)]
+# ---------------------------------------------------------------- public solve
+def solve(A, b, parallel=True):
+    """Exact rational solution of ``A x = b``.
+    Args:
+        A: ``[n, n]`` integer matrix (torch int tensor, list, or numpy array).
+        b: ``[n, c]`` integer right-hand side(s). A 1-D ``[n]`` vector is
+           treated as a single column.
+        parallel: use a process pool (fork or spawn) for reconstruction when
+           available.
+    Returns:
+        ``[n][c]`` list of ``fractions.Fraction`` (the exact solution), or
+        ``None`` if ``A`` is singular over the rationals.
+    Requires a CUDA device; entries with ``n * max|A_ij| < 2^63``.
+    """
+    import numpy as np
+    dev = A.device if isinstance(A, torch.Tensor) else torch.device("cuda")
+    if isinstance(b, torch.Tensor) and b.dim() == 1:
+        b = b.unsqueeze(1)
+    n = A.shape[0] if isinstance(A, torch.Tensor) else len(A)
+    c = (b.shape[1] if isinstance(b, torch.Tensor) else len(b[0]))
+    A_list = _to_int_lists(A, n, n)
+    b_list = _to_int_lists(b, n, c)
+    lognum, logden = _hadamard_bits(A_list, b_list)
+    margin = 48
+    A_t = torch.tensor(A_list, dtype=torch.int64, device=dev)
+    primes_all = torch.tensor(_ODD_PRIMES, dtype=torch.int32, device=dev)
+    Bout_all, singular = _batched_modinv(A_t, primes_all)
+    good = [i for i in range(len(_ODD_PRIMES)) if int(singular[i]) == 0]
+    if not good:
+        return None
+    Q = [_ODD_PRIMES[i] for i in good]
+    P = 1
+    for q in Q:
+        P *= q
+    Pinv = pow(P, -1, 1 << 64)
+    Pinv_s = Pinv - (1 << 64) if Pinv >> 63 else Pinv
+    Lx = 1
+    while (1 << (64 * Lx - 1)) <= P:
+        Lx += 1
+    m = int(math.ceil((1.0 + lognum + logden + margin) / math.log2(P)))
+    Mmod = P ** m
+    Nb = 1 << int(math.ceil(lognum + margin / 2))
+    Db = 1 << int(math.ceil(logden + margin / 2))
+    primes, inv, mu, pow32 = _crt_tables(Q, dev)
+    Ecnt = len(Q)
+    Mp, Kp, Cp = max(n, 17), _pad8(n), max(8, ((c + 7) // 8) * 8)   # narrow RHS pad for small c
+    Blift = Bout_all[torch.tensor(good, device=dev)].contiguous()
+    Blift_p = torch.zeros(Ecnt, Mp, Kp, dtype=torch.int8, device=dev)
+    Blift_p[:, :n, :n] = Blift
+    r = torch.tensor(b_list, dtype=torch.int64, device=dev)
+    digits = torch.empty(m, Lx, n, c, dtype=torch.int64, device=dev)
+    for it in range(m):
+        rT = r.t().contiguous()
+        rT_planes = _extract(rT, Cp, Kp, primes, mu, pow32)
+        xi = _matmul(Blift_p, rT_planes, Lx, primes, inv, mu)[:, :n, :c].contiguous()
+        digits[it] = xi
+        r = _padic(r, _i64_matmul(A_t, xi[0].contiguous()), Pinv_s)
+    NC = n * c
+    D3 = np.ascontiguousarray(digits.permute(0, 2, 3, 1).reshape(m, NC, Lx).cpu().numpy())
+    # Reconstruct every entry independently to its exact (num, den). The bounds
+    # give Nb*Db <= Mmod/2, so rational reconstruction is unique and correct per
+    # entry; no shared common-denominator estimate is involved. Parallelize the
+    # per-entry Euclid across cores for large grids (fork or spawn pool).
+    rb = _recon()
+    use_pool = parallel and NC >= 256
+    pairs = None
+    if use_pool:
+        try:
+            nch = min(NC, min(int(os.cpu_count() or 8), 32) * 3)
+            bnd = [round(x * NC / nch) for x in range(nch + 1)]
+            tasks = [(D3[:, bnd[i]:bnd[i + 1], :].copy(), m, P, Mmod, Nb, Db)
+                     for i in range(nch) if bnd[i + 1] > bnd[i]]
+            pairs = []
+            for res in _pool().map(rb._recon_chunk, tasks):
+                pairs.extend(res)
+        except Exception:
+            pairs = None
+    if pairs is None:
+        pairs = rb._recon_chunk((D3, m, P, Mmod, Nb, Db))
+    if any(p is None for p in pairs):
+        return None
+    return [[Fraction(pairs[i * c + j][0], pairs[i * c + j][1]) for j in range(c)] for i in range(n)]
+__all__ = ["solve"]

build/torch213-cxx11-cu130-aarch64-linux/_exact_solve_cuda_31a8fa7_dirty.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:ab909fca772b391348b267060579307c2f63977d5cf3d1c8c2e1542d25dcf740
+size 1496576

build/torch213-cxx11-cu130-aarch64-linux/_exact_solve_recon.py ADDED Viewed

	@@ -0,0 +1,45 @@

+"""Torch-free rational reconstruction (Horner + extended-Euclid) for exact-solve;
+imported by the reconstruction process-pool workers (gmpy2/numpy only, no torch)."""
+import math
+try:
+    import gmpy2
+    from gmpy2 import mpz
+    _HAVE_GMP = True
+except Exception:                       # pragma: no cover - optional acceleration
+    _HAVE_GMP = False
+    def mpz(x):
+        return int(x)
+def _rat_recon(X, M, N, D):
+    """Rational reconstruction: a/b with a == b X (mod M), |a|<=N, 0<b<=D."""
+    r0, r1 = mpz(M), mpz(X) % mpz(M)
+    t0, t1 = mpz(0), mpz(1)
+    while r1 > N:
+        q = r0 // r1
+        r0, r1 = r1, r0 - q * r1
+        t0, t1 = t1, t0 - q * t1
+    a, b = r1, t1
+    if b < 0:
+        a, b = -a, -b
+    if b == 0 or b > D:
+        return None
+    g = math.gcd(int(abs(a)), int(b)) if not _HAVE_GMP else gmpy2.gcd(abs(a), b)
+    a, b = a // g, b // g
+    return (int(a), int(b)) if b <= D else None
+def _recon_chunk(args):
+    """Reconstruct every entry in the chunk: Horner the p-adic digits into X, then
+    extended-Euclid X (mod M) into an exact (num, den) pair, each independently."""
+    D3chunk, m, P, Mmod, Nb, Db = args
+    Pm = mpz(P); Mmz = mpz(Mmod)
+    out = []
+    for k in range(D3chunk.shape[1]):
+        X = mpz(0)
+        for it in range(m - 1, -1, -1):
+            X = X * Pm + int.from_bytes(D3chunk[it, k].tobytes(), "little", signed=True)
+        out.append(_rat_recon(X % Mmz, Mmz, Nb, Db))
+    return out

build/torch213-cxx11-cu130-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _exact_solve_cuda_31a8fa7_dirty
+ops = torch.ops._exact_solve_cuda_31a8fa7_dirty
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_exact_solve_cuda_31a8fa7_dirty::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/exact_solve/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch213-cxx11-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "exact-solve",
+  "id": "_exact_solve_cuda_31a8fa7_dirty",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ZG5Z4fcpkB0UCWfChZwPg4dTmNtHpTxDCEvKPLJc9PI=",
+      "_exact_solve_cuda_31a8fa7_dirty.abi3.so": "q5CfyncrORNIsmcGBXkwfC9jl31c89HIwuFULSXc90A=",
+      "_exact_solve_recon.py": "gU5Ywz+UPfUeXS2dyzJ8NOOxgj7pgUcbFV9uAOAVIzU=",
+      "_ops.py": "1ujza9fF0nYcw94Evi6lqaZ8wNTrVB+GsP8zCO4FVj8=",
+      "exact_solve/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "31a8fa724c4e4c412e04efba7ee534c01afb61cd",
+      "dirty": true
+    }
+  }
+}