enclosure-gemm-cpu

enclosure-gemm-cpu returns a pair of IEEE matrices (lo, hi) bracketing the matrix product evaluated in exact real arithmetic. It is the CPU backend of enclosure-gemm, same contract and same bytes, and the reference against which the GPU backends are checked. 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.

An interval GEMM rounds outward on each partial sum, so the bracket widens by about two ulps per accumulation. Here the inner product is accumulated as an exact integer in 384-bit two's-complement arithmetic and rounded outward once.

Usage

import torch
from kernels import get_kernel

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

A = torch.randn(256, 1024, dtype=torch.float64)
B = torch.randn(1024, 256, dtype=torch.float64)

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 the CUDA and SYCL backends

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. Operands are float32 or float64.

API

Symbol Purpose
mm_enclose(A, B) (lo, hi) bracketing the exact real product A @ B
mm_enclose(A, B, return_info=True) adds the fixed-point plan and whether the scaling was lossless
mm_enclose_nt(A, Bt) A @ Bt.T for Bt [N, K], the orientation the kernels take
bmm_enclose(A, B) batched, [b, M, K] @ [b, K, N], one fixed point for the batch
bmm_enclose_nt(A, Bt) batched, transposed second operand
mm_interval(A_mid, A_rad, B_mid, B_rad) enclosure for midpoint-radius interval operands
mm_interval_nt(A_mid, A_rad, Bt_mid, Bt_rad) the same, transposed second operand
contains(lo, hi, X) containment check
width_ulps(lo, hi) bracket width per entry, in representable values
digest(lo, hi) SHA-256 over the raw bits
set_operand_checks(bool) the split status check and its synchronisation

Method

Operands are scaled by a power of two and floored into two's-complement integers of up to four int64 limbs, the scale taken from their exponent range. The inner product of the limb matrices is accumulated directly in 384-bit integer arithmetic; the CUDA backend recovers the same integer from residues on tensor cores. The outward-rounding epilogue is the shared core verbatim, so the enclosures agree byte for byte. The enclosure-gemm card gives the full account.

No floating-point instruction executes, in kernel or wrapper; IEEE values are read and written as bit patterns and the arithmetic between them is integer.

The inner product is register-blocked over a 4x4 tile of outputs. The accumulator is a six-word carry chain and therefore serial, so one output element at a time leaves the pipeline mostly idle; the tile runs sixteen independent chains and reads eight operand limbs per step instead of thirty-two. Limb counts and tile extents reach the loop as template parameters, so the loops in the shared core fold to straight-line code and the accumulators stay in L1. With runtime bounds the same source spills every step and runs about sixty times slower.

The extension carries its own thread pool: the threading configuration fixed at build time need not match the torch runtime the extension is loaded into, and at::parallel_for degrades to serial on mismatch. Pool width follows the hardware thread count; ENCLOSURE_GEMM_THREADS caps it.

Measured

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

K enclosure-gemm-cpu directed-rounding interval
16 1 ulp 52 ulp
128 1 ulp 189 ulp
1024 1 ulp 2504 ulp
8192 1 ulp 19771 ulp
65536 1 ulp 39266 ulp

Throughput at 256x1024x256 float64 on a 24-thread x86_64 host:

threads time rate
1 1990 ms 33.7 M mac/s
4 553 ms 121 M mac/s
12 218 ms 308 M mac/s
24 120 ms 559 M mac/s

The rate is set by 64-bit integer multiply and a dependent six-word carry chain. The CUDA backend is 106x faster on the same shape.

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 NVIDIA sm_89 under nvcc, Intel Iris Xe Gen12LP under icpx, x86_64 under MSVC, and Cortex-A76 and Cortex-A72 under GCC, compared by SHA-256 over the raw bits on a frozen case set including subnormals, cancellation and mixed scale.

tools/standalone_check.cpp verifies the arithmetic core with a C++11 compiler and no torch, BLAS or GPU, and touches no floating-point type, so it also runs on a target with no fp64 unit. It is how the two aarch64 parts were checked.

c++ -O2 -o standalone_check tools/standalone_check.cpp -I enclosure_gemm_cpu_csrc
./standalone_check cases.bin out.bin

Requirements and limits

  • K <= 2^17 and bits_a + bits_b + ceil(log2 K) <= 332, matched to the CUDA backend's residue range so an input accepted by one backend 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.
  • 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).

License

Apache-2.0.

Downloads last month
-
apache-2.0
OS
linux
Arch
x86_64aarch64
Kernel Builder
22e2aad