enclosure-gemm

enclosure-gemm returns a pair of IEEE matrices (lo, hi) bracketing the matrix product evaluated in exact real arithmetic, computed on INT8 tensor cores and loadable through kernels. When both operands scale losslessly into its fixed point, the bracket is the tightest the format admits and its width is independent of the contraction length: one representable step at K = 16 and one at K = 65536. It uses the residue engine of exact-gemm and fp64-emu. The CPU and Intel GPU backends return identical bytes: enclosure-gemm-cpu, enclosure-gemm-xpu.

An interval GEMM rounds outward on each partial sum, so the bracket widens by about two ulps per accumulation and reaches thousands of ulps at production contraction lengths. Here the inner product is accumulated as an exact integer, in residues modulo the primes below 256 recombined by balanced Garner, and rounded outward once.

Usage

import torch
from kernels import get_kernel

eg = get_kernel("phanerozoic/enclosure-gemm", version=1, trust_remote_code=True)

A = torch.randn(1024, 4096, dtype=torch.float64, device="cuda")
B = torch.randn(4096, 1024, dtype=torch.float64, device="cuda")

lo, hi = eg.mm_enclose(A, B)          # lo <= A @ B <= hi, entrywise, exactly
eg.width_ulps(lo, hi).max()           # 1
eg.digest(lo, hi)                     # identical on every backend

lo, hi = eg.mm_enclose_nt(A, Bt)      # A @ Bt.T, no transpose copy
lo, hi = eg.bmm_enclose(A3, B3)       # [b, M, K] @ [b, K, N]
lo, hi = eg.mm_interval(A_mid, A_rad, B_mid, B_rad)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. Operands are float32 or float64. torch._int_mm shape constraints (M > 16; K, N multiples of 8) are met by zero-padding, which is exact.

API

Method

Operands are scaled by a power of two and floored into two's-complement integers of up to four int64 limbs. The scale is the finest binary point holding every entry inside the 256-bit fixed point, coarsened until the product fits the residue range. Where the exponent spread permits, the scaling is lossless and the integer inner product is the real inner product.

The residue scheme recovers that integer: one torch._int_mm INT8 tensor-core GEMM per prime below 256, reconstructed by balanced-Garner mixed-radix digits in signed 384-bit arithmetic.

The epilogue maps the exact integer S and its binary exponent Q to the endpoint pair by bit manipulation. The leading bit of S locates the significand, the top 53 bits are the mantissa, and the remainder selects the pair: a zero remainder means S * 2^-Q is representable and lo == hi, otherwise truncation gives the endpoint toward zero and incrementing the encoding gives the endpoint away from it, the sign deciding which is which. Overflow saturates to the largest finite value inward and to infinity outward. The subnormal boundary needs no case, the encoding being contiguous across it.

No rounding mode is consulted and no floating-point instruction executes, in kernel or wrapper; IEEE values are read and written as bit patterns. Directed rounding is unavailable under SYCL, and a target may have no fp64 unit; neither is required.

Beyond the fixed point's reach the bracket stays valid and loses tightness. Flooring leaves a remainder below 2^-F per entry, the induced error in the product is bounded termwise, and the bound is added back outward. return_info reports the regime.

Interval operands use Rump's midpoint-radius bound, its three products and its radius formed as exact integers, so the overestimation is that of the bound alone.

Planning an operand reads three reductions over its exponent field, stacked and transferred together, so it costs one host synchronisation. A second guards a planner-kernel invariant and set_operand_checks removes it; inf and nan are rejected by the planner either way. A batch shares one plan, so it pays the planning cost once.

Measured

Bracket width against the same computation with directed rounding on each accumulation, random normal operands:

Time on an RTX 6000 Ada, float64, best of ten, against torch's uncertified products on the same device. The fp64 rate of that part is a sixty-fourth of its INT8 rate.

At 256x1024x256: 59,215 M mac/s here, 1,943 M mac/s on an Intel Iris Xe (Gen12LP, no matrix engine, no fp64 unit), 559 M mac/s on 24 x86_64 threads, 33.7 M mac/s on one.

A batch of five 20x256x12 products takes 2.74 ms through bmm_enclose and 5.00 ms as a loop of mm_enclose, the difference being one plan and one synchronisation for the batch instead of five of each.

Correctness

The reference is exact rational arithmetic. Every binary64 value is a rational, so the tests evaluate the true inner product with fractions.Fraction and require the returned pair to bracket it, and where the scaling is lossless to equal its exact floor and ceiling in the format. Cases: random operands at four shapes; K from 8 to 16384; integer operands, where the product is representable and lo == hi must equal the torch result; catastrophic cancellation; subnormal operands; overflow to infinity; float32; operands spanning 520 binades, where containment is required and tightness is not; interval operands against sampled corners.

(lo, hi) is bit-identical across the targets below, compared by SHA-256 over the raw bits on a frozen case set including subnormals, cancellation and mixed scale.

The aarch64 entries run tools/standalone_check.cpp from the CPU backend, which requires a C++11 compiler and touches no floating-point type.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ and torch with _int_mm.
• K <= 2^17, set by the int32 residue accumulator.
• bits_a + bits_b + ceil(log2 K) <= 332, set by the 335-bit product of the primes below 256. Every backend reports the same figure, so an input accepted by one is accepted by all and yields the same bits.
• The one-step bracket requires both operands to scale losslessly into a 256-bit fixed point, so their exponent spread must fit the budget above. Wider operands receive a valid bracket and return_info marks it.
• Working memory is about P * (M*K + N*K) int8 plus P * M*N int32 for P primes.
• Three host synchronisations per call, two to plan the operands and one to check the split status; set_operand_checks(False) removes the third.
• Non-finite operands raise rather than propagate.

References

Moore 1966 (Interval Analysis); Rump 1999 (INTLAB, midpoint-radius product bound); Kulisch and Miranker 1986 (exact dot product); Ozaki, Ogita, Oishi, Rump 2012 (Numerical Algorithms 59); Ozaki, Uchino, Imamura 2025 (Ozaki Scheme II, arXiv:2504.08009).

License

Apache-2.0.