Ozaki Scheme II Is Fast on CPUs Too: Multiple-Precision Matrix Multiplication on Intel AMX-INT8 and Arm SVE2-i8mm

1 Introduction

• Direct MPFR-style multiplication performs O(MNK) expensive multiply-add-and-round operations, while Strassen reduces operation count and Ozaki methods replace the arithmetic with low-precision GEMMs.The paper frames Ozaki methods as decomposing inputs into low-precision quantities so optimized GEMMs do the work.Source: §1 · p. 1
• Ozaki Scheme I [1] sums products of significand slices and needs O(p^2) GEMMs, while Ozaki Scheme II [2] uses residues and CRT with O(p) GEMMs.Scheme I has also been applied to GPU Tensor Cores, INT8 engines, and multiple-precision linear algebra; Scheme II was extended to FP8 units via quantization [5].Source: §1 · pp. 1–2
• The paper implements Ozaki Scheme II for arbitrary-precision GEMM with MPFR input and output, sharing CRT reconstruction across Intel AMX INT8 and binary64 DGEMM backends.The code is also ported to AArch64, where its INT8 kernel uses the SVE2-i8mm SMMLA instruction.Source: §1 · p. 2
• The paper positions its tutorial and quantitative CPU-performance explanation before presenting the precise algorithm and implementation.The tutorial introduces Ozaki Scheme II, followed by the CPU-performance analysis, algorithm formulation, and implementation.Source: §1 · p. 2

2.1 The problem: multiple-precision multiplication is built from bricks

• A 1024-bit significand uses 16 limbs, and its schoolbook multiplication requires 256 limb multiplications plus carry handling, normalization, and rounding.The paper measures an MPFR multiply-add pair at about 80 ns per core for p = 212 bits and estimates several seconds for an N = 1024 matrix product on 20 cores.Source: §2.1 · p. 2

2.2 Step 1: make the inputs integers (shared exponents)

• Shared exponents convert each row of A and column of B into t-bit integers; choosing t above the target precision leaves error only at initial quantization and final rounding.The default is t = p + 16 + ⌈log2 K⌉, and reported errors stay below 1 ulp in every condition.Source: §2.2 · pp. 2–3

2.3 Step 2: split one big-integer computation into many small ones (residue number systems)

• Pairwise-coprime moduli encode a large integer as residues, allowing independent small-integer inner products to be computed as low-precision GEMMs and recovered uniquely by CRT.For the worked example, the residue triple (0, 3, 1) reconstructs the inner product 10.Source: §2.3 · p. 3

2.4 Step 3: when is a low-precision GEMM exact?

• Low-precision GEMMs are exact for these residues when rounding and overflow are avoided; approximation enters only during initial quantization and final rounding.For INT8, the stated worst-case accumulation bound is K · 127^2 < 2^31; binary64 exactness is bounded by the magnitude of the integer products and sums.Source: §2.4 · pp. 3–4

2.5 Step 4: CRT reconstruction

• After the L residue GEMMs, CRT reconstructs each integer output, which is scaled back into an MPFR value using a formulation that avoids multiple-precision division.The reconstruction cost is O(L) per output entry.Source: §2.5 · p. 4

2.6 Contrast with Ozaki scheme I: slicing

• Ozaki Scheme I’s pairwise slice products require O(p^2) GEMMs, whereas Scheme II’s independent moduli require O(p); at 1024 bits, the stated counts are 1431 versus 99 DGEMMs.The paper identifies this count difference as the main reason measured speed ratios grow with precision.Source: §2.6 · p. 4

3.1 Reason 1: per-operation cost differs by 3–4 orders of magnitude

• The measured effective throughput gap between MPFR multiply-adds and matrix instructions is four to five orders of magnitude.The comparison is per operation; MPFR processes p bits per multiply-add while INT8 covers only 8 bits, so the paper says a fair account must compare cost per precision bit.Source: §3.1 · pp. 4–5, Table 1 · p. 4

3.2 Reason 2: the GEMM shape unlocks arithmetic density

• GEMM’s O(N^3) work on O(N^2) data permits cache reuse, while reducing element size from 8 bytes to 1 byte yields 8× larger cache tiles and 8× more operations per byte.Ozaki preprocessing costs O(LN^2), one order below the O(LN^3) GEMMs, so its investment is amortized as N grows.Source: §3.2 · p. 5

3.3 Reason 3: per-bit economics and the O(p) GEMM count

• Despite using 2.7× more moduli, INT8 retains a >7× advantage on GB10; the corresponding Xeon coefficient is about 8.The comparison is framed as per-bit economics, with one modulus carrying about 8 bits for INT8 versus about 21.5 bits for binary64; Scheme II’s GEMM count is O(p), compared with O(p^2) for Scheme I.Source: §3.3 · p. 5

4.1 Overall flow

• The five-stage flow converts MPFR inputs to shared-exponent integer significands, computes low-precision modular GEMMs, reconstructs signed integers with CRT, then rescales and rounds to MPFR outputs.The supplied flow also identifies integerization and residue generation before the modular products.Source: §4.1 · pp. 5–6

4.2 Integerization with shared exponents

• Each row of A and column of B uses a shared exponent to form integer significands; the default width is t = p + g + ⌈log2 K⌉, with g = 16 guard bits.The quantization error is bounded roughly by K 2^t, after which an exact integer inner product allows only one final rounding.Source: §4.2 · p. 6

4.3 Unique CRT reconstruction and the choice of moduli

• At K = 4096 and t = 240, the binary64 backend uses 23 moduli; backend-specific exactness bounds determine the moduli, and the AMX design adds split two-limb moduli beyond its one-byte range.For AMX, one-byte moduli fit signed INT8 and INT32 accumulation is bounded by K · 127^2 < 2^31; binary64 centered residues are bounded to keep DGEMM results exact. The hybrid uses three GEMMs per two-limb modulus.Source: §4.3 · pp. 6–7

4.4 Direct CRT without multiple-precision division

• Direct CRT estimates the quotient in binary64, then uses integer correction and at most two subtractions to obtain the exact residue without multiple-precision division.This reorganizes the original paper’s FP64 quotient-estimation and integer-correction approach so the quotient is determined exactly [2].Source: §4.4 · p. 7

5.1 Structure

• The roughly 2,900-line implementation uses GMP, MPFR, and OpenBLAS, with x86-64 AMX and AArch64 SVE2-i8mm kernels selected automatically by the Makefile.The two kernels share pack-buffer dimensions and interfaces but use different internal layouts.Source: §5.1 · p. 8

5.2 The AMX-INT8 GEMM kernel

• L3 or memory tile loads cost 6–8× more than L1 loads and reduce TMUL performance sixfold, so the kernel blocks data to preserve locality.It keeps A and B blocks in L2 and streams a shared B block across threads; this raised the 32-thread rate from 5.0 to 27 Top/s.Source: §5.2 · p. 8

5.3 Residue generation and packing

• Residues are generated from base-216 significand digits, then B is converted to VNNI layout by transposing four k-rows and writing each 64-byte destination as one cache line.The digit loop enables AVX-512 vectorization of the element loop; reciprocal multiplication with correction performs the reduction.Source: §5.3 · p. 9

5.4 AVX-512 batched CRT

• AVX-512 CRT vectorizes independent output elements using eight 32-bit digit lanes, while branch-free mask operations handle correction and signed conversion.Two extra 32-bit digits of headroom prevent overflow when the product’s bit length approaches the allocated digit boundary.Source: §5.4 · p. 9

• Inputs are uniform random values in [−1, 1), with some tests widening the exponent range by ±spread.Source: §6 · p. 9
• Accuracy is the maximum relative error in ulps of the target precision, measured on 32 randomly chosen outputs against references using p + 256-bit inner products.Source: §6 · p. 9

7.1 Correctness

• Every tested element agreed with a p + 160-bit MPFR reference rounded to p bits across precisions from 53 to 1024 bits and varied matrix shapes and exponent ranges.The standalone CRT test also confirmed identical exact-integer reconstruction by scalar and AVX-512 paths over t = 40–1120 bits.Source: §7.1 · p. 9

7.2 Standalone AMX-INT8 GEMM performance

• At N = K = 4096, the AMX pipeline sustains 39–43 Top/s, rising to 51.6 Top/s at N = K = 8192.An isolated GEMM reaches about 20 Top/s at N = 4096; the pipeline achieves 23–28% of the per-core L1-resident peak on 32 cores.Source: §7.2 · p. 10

7.3 Phase breakdown

• At N = 4096, AMX GEMMs account for 13–17% of total time and residue conversion dominates, while binary64 GEMMs account for 61–66%.Source: §7.3 · p. 10, Table 7 · p. 10

7.4 Comparing the two backends

• The AMX GEMM is about 22× faster than DGEMM (51.6 Top/s vs. 2.3 TFlop/s), but its 3–4.5× higher modulus count adds pre/post-processing cost; backend break-even is N ≈2048.AMX takes over where GEMMs dominate, at N ≳4096.Source: §7.4 · p. 11, Table 8 · p. 11

7.5 Against a naive MPFR matrix product

• Table 9 reports AMX speedups against a naive 32-thread parallel GEMM built from mpfr_mul and mpfr_add.Source: §7.5 · p. 11, Table 9 · p. 11

7.6 Against BNCmatmul

• Table 11 compares speed ratios against BNCmatmul routines as shipped; extreme BNC-OzI-blas ratios reflect missing fairness corrections, not an algorithmic difference.The routines cover naive MPFR, Strassen, and Ozaki Scheme I variants.Source: §7.6 · p. 11, Table 11 · p. 12

7.7 Accuracy

• Ozaki Scheme II stays below 1 ulp in every reported condition, while Scheme I loses a few ulps to slice truncation and Strassen loses hundreds to thousands.Source: §7.7 · p. 11, Table 12 · p. 12

8 Porting to aarch64 (NVIDIA GB10)

• The port targets the GB10’s Cortex-X925 × 10 plus Cortex-A725 × 10 big.LITTLE configuration.The evaluation environment is identified in Table 13.Source: §8 · p. 13, Table 13 · p. 14

8.1 Instruction choice: no SME/SME2, so SVE2-i8mm

• Without SME/SME2, the port uses the SVE2-i8mm SMMLA instruction through NEON intrinsics for exact INT8×INT8→INT32 multiply-accumulate.Each SMMLA performs 32 multiply–adds; the implementation uses vmmlaq_s32.Source: §8.1 · p. 13

8.2 Kernel and pack layout

• The kernel accumulates an 8 × 8 output block in registers across the full K extent, writes C once, and reuses the AMX version’s API and driver.The port also replaces batched CRT with scalar GMP, dynamically distributes row tiles across unequal cores, and vectorizes packing in NEON.Source: §8.2 · p. 13

8.3 Correctness and standalone INT8 kernel performance

• The end-to-end results agree with the reference in every condition at 0 ulp, and the GB10 INT8 kernel sustains 6.4–6.5 TOPS at N = 4096.The sustained rate is 50–60% of the estimated SMMLA peak and roughly 20× the measured 0.33 TFlop/s of OpenBLAS DGEMM.Source: §8.3 · p. 14

8.4 Phase breakdown

• At 1024 bits, scalar two-limb packing and scalar CRT dominate the GB10 pipeline, leaving optimization headroom.At low and medium precision, GEMM and packing are balanced as on x86.Source: §8.4 · p. 14

8.5 Backend comparison: on the GB10, INT8 wins almost everywhere

• On the GB10, INT8 is fastest almost everywhere for N ≥512, whereas binary64 was preferable on x86 for N ≲2048.The GB10’s measured DGEMM performance is 0.33 TFlop/s; at N = 4096, p = 53, DGEMM accounts for about 3 s of the 3.46 s total.Source: §8.5 · p. 14

8.6 Against BNCmatmul’s Ozaki scheme I

• Ozaki Scheme II stays below 1 ulp (0.73–0.95 ulp) in every condition, while Ozaki Scheme I loses 1.8–12.4 ulp.The speedup over a naive MPFR product (20-thread parallel) is 19.8× at N = 1024, p = 212.Source: §8.6 · p. 15

8.7 Summary of the port, and outlook towards SME/SME2

• Across both AMX and SMMLA, the same algorithm outperformed each machine’s Ozaki Scheme I implementation by an order of magnitude or more.The engines differ 512-fold in per-instruction multiply–add granularity, supporting the paper’s account that the reduction to exact low-precision integer GEMMs is the common advantage.Source: §8.7 · p. 15

9.1 The OpenBLAS/OpenMP affinity clash

• The OpenBLAS/OpenMP affinity clash reduced an N = 2048 DGEMM to 114 GFlop/s versus about 2,000; row-blocking with single-threaded DGEMM restored 1.35–2.31 TFlop/s.The same fix made Ozaki Scheme I DGEMMs 8–11× faster; at N = 1024, 212 bits, throughput rose from 100 to 1,001 GFlop/s.Source: §9.1 · pp. 15–16

9.2 The serial section remaining in Ozaki scheme I

• After fixing DGEMM affinity, serial splitting remains 88% of OzI-best; idealizing that split narrows the gap to 6.6× at N = 512 and 7.1× at N = 1024.The extrapolated comparisons are closer to the intrinsic algorithmic difference than ratios between current implementations.Source: §9.2 · p. 16

9.3 Asymptotics of the GEMM count

• At K = 4096 and p = 1024, Scheme I needs 1,431 DGEMMs, versus 99 for Scheme II/binary64 and 435 for Scheme II/AMX; Scheme II won every measured condition.Scheme II’s GEMM count grows as O(p), compared with O(p²) for Scheme I, while its ratio over Strassen shrinks at high precision.Source: §9.3 · p. 16

9.4 Applicability

• At N ≤64, O(LN²) residue conversion and CRT cannot be amortized against O(N³) GEMMs, and the disadvantage grows with precision.Higher precision increases L and thus the preprocessing handicap.Source: §9.4 · p. 16

10 Limitations •

• INT8 backends require K ≤133,000 under the INT32 accumulation bound; in the binary64 backend, the modulus bound decreases as K grows.Source: §10 · p. 16
• The residue kernel uses 32-bit accumulation only for m ≤255 and at most 250 digits (t ≲4000 bits), falls back to a 64-bit path otherwise, and has a built-in modulus table supporting about 25,000 bits.Source: §10 · p. 17
• Inf/NaN do not propagate and are treated as 0.Source: §10 · p. 17
• With per-row/per-column shared exponents, small components are lost if the exponent range within a row or column exceeds t; guard_bits / t_bits are tunable.Source: §10 · p. 17
• The x86-64 backend requires AMX tile-state permission, Linux 5.18+, and an AMX-capable CPU.Source: §10 · p. 17
• The aarch64 backend requires the i8mm extension, generally available from Armv8.6-A; its batched CRT and two-limb packing remain scalar.The scalar operations leave headroom at high precision and large scale.Source: §10 · p. 17

11 Conclusions and Future Work

• Below 1 ulp, the CPU implementation runs up to 167× faster than naive MPFR, 588× faster than BNCmatmul’s Strassen, and 9–78× faster than Ozaki Scheme I.The AMX-INT8 and binary64-DGEMM backends share one CRT reconstruction stage.Source: §11 · p. 17
• The AArch64 port sustains 6.5 TOPS for exact INT8 GEMM and achieves 14–89× speedups over BNCmatmul’s Ozaki Scheme I without SME/SME2.It uses the SMMLA instruction from SVE2-i8mm; the authors attribute the reproduced speedup pattern to Ozaki Scheme II’s structure rather than a particular engine.Source: §11 · p. 17
• Residue conversion is the largest x86 INT8-backend bottleneck, accounting for 45–50% of total time; Montgomery representations and partial-product reuse are candidates.The authors identify faster residue conversion (pack) as a future improvement.Source: §11 · p. 17
• Future extensions include reusing A’s residues and packed planes, complex products via 3M/4M, mixed AMX/DGEMM moduli, and integration with BNCmatmul’s MPFMatrix.Reusing A can amortize its O(N^2) preprocessing across repeated products.Source: §11 · p. 17
• An SME/SME2 backend is a planned direction; the authors describe its kernel as structurally similar to AMX, with a pack-layout swap once suitable hardware is available.They also plan NEON/SVE2 vectorization of the two-limb pack and CRT on AArch64.Source: §11 · p. 17