Skip to content

Unify rotation entry points and accumulate rotation folds lazily in the extended basis - #1225

Open
pascoec wants to merge 6 commits into
devfrom
gpu
Open

Unify rotation entry points and accumulate rotation folds lazily in the extended basis#1225
pascoec wants to merge 6 commits into
devfrom
gpu

Conversation

@pascoec

@pascoec pascoec commented Jul 18, 2026

Copy link
Copy Markdown
Collaborator

Summary
This PR makes OpenFHE's CKKS rotation-heavy paths — bootstrap EvalPartialSumInPlace, the Horner giant-step in CtS/StC, and the scheme-switching linear transforms — deterministic across entry points and cheaper, by (1) routing every automorphism through one shared core and (2) keeping rotation-fold accumulators in the extended Q·P basis with a single deferred ApproxModDown per fold. The fold radix is a build-time parameter (PARTIAL_SUM_RADIX), making the performance/rotation-key-count trade-off explicit, and the same lazy radix-fold treatment is extended to the scheme-generic EvalSum/EvalSumRows/EvalSumCols family. Along the way it deduplicates the multiparty keygen index computations, fixes a result-aliasing bug in the EvalAddMany/EvalMultMany wrappers, and removes the redundant copies inside the *Many operations.

Motivation: GPU CKKS backends (specifically FIDESlib) implement rotations in hoisted, lazily-accumulated form. For a GPU backend to produce ciphertexts bit-identical to the CPU reference, the CPU paths must apply the same operation sequence with the same rounding. The changes here achieve that while strictly reducing CPU-side mod-down counts — they are not GPU-specific and stand on their own as a CPU improvement.

Commits

  1. Unification of EvalFastRotation/EvalRotate/EvalAutomorphism/Conjugate
    Adds LeveledSHEBase::EvalAutomorphismCore(ciphertext, autoIndex, digits, evalKey) — the automorphism-plus-key-switch step evaluated from precomputed digits — and delegates EvalFastRotation, EvalAutomorphism (hence EvalRotate/EvalAtIndex), and CKKS Conjugate to it. Previously EvalAutomorphism used a non-hoisted sequence (KeySwitchInPlace, then permute both components) whose rounding differs from the hoisted path; the two entry points produced bit-different results for the same rotation. Now every rotation entry point is bit-identical to EvalFastRotation with fresh digits.

  2. Lazy extended-basis accumulation (deferred ApproxModDown)
    Rotation folds previously settled each key-switched rotation back to the base ring before adding it to the accumulator. Two exact identities make that unnecessary:

adding a base-ring operand into an extended-basis (Q·P) accumulator lifts it by P mod q exactly, and ApproxModDown(P·x + ks) = x + ApproxModDown(ks) — no rounding difference;
automorphisms are coefficient permutations (exact, linear), so they commute with the deferred settle.
So a fold of k rotations can keep the running sum extended and perform one ApproxModDown at the end, bit-identical to the eager form. Applied to:

EvalPartialSumInPlace (bootstrap partial sums) — the accumulator settles once per fold; the c1 component still settles per step, since the next step's digit decomposition consumes it in the base ring;
the Horner giant-step in CoeffsToSlots/SlotsToCoeffs (EvalHornerGiantRotate), including folding the delta corrections in before the final settle;
scheme-switching EvalLTWithPrecomputeSwitch/EvalLTRectWithPrecomputeSwitch (new defaulted ext flag; the sparse SlotsToCoeffs doubling and the wide-matrix log-fold now accumulate extended and settle once).
ApproxModDown is the dominant cost of a rotation after the key-switch products, so per fold this removes all but one of them on the accumulator path.

3/4. Radix-parameterized EvalPartialSumInPlace, configured at build time
Generalizes the partial-sum fold into a radix-r rotation fold and adds two overloads:

EvalPartialSumInPlace(ct, stride, size) — radix-2 specialization;
EvalPartialSumInPlace(ct, stride, size, radix) — general form, evaluated in radix-r levels so each level shares one digit decomposition across its up to r−1 rotations (delegates to the 3-arg version when radix == 2).
The radix used throughout the library is a new CMake cache variable, PARTIAL_SUM_RADIX (default 4, any power of two), emitted into config_core.h (introduced as CKKS_PARTIAL_SUM_RADIX; renamed in commit 5 when it ceased to be CKKS-specific). All in-tree callers fold at the configured radix — the three bootstrapping PartialSum sites and the FHEW→CKKS sparse re-encode (previously an inline eager doubling loop, now the one remaining occurrence of the pattern converted to the subroutine). Key generation produces exactly the fold's rotation-index set {i·stride·rˡᵉᵛᵉˡ : i ∈ [1, r)} at the three bootstrap Find*RotationIndices helpers and the two scheme-switching setup sites, so the configured radix always has its keys.

The trade-off the variable exposes: mod-ups per fold scale as log2(size)/log2(r), while dedicated rotation keys grow as (r−1)/log2(r) · log2(size). The default of 4 halves the digit decompositions at the raised level — the widest, most expensive mod-ups of the bootstrap — and measures ≈1.3× over radix 2 for the fold itself (single-threaded Ice Lake). Radix 2's indices are all powers of two — shared with keys bootstrapping generates anyway — so builds that prioritize key storage over bootstrapping runtime should set -DPARTIAL_SUM_RADIX=2, which reproduces the previous behavior exactly (for r > 2 the index set is a superset of the power-of-two set, so no existing consumer of those keys breaks). Documented in cmake_in_openfhe.rst and Best_Performance.md.

  1. Radix EvalSum family, multiparty keygen dedup, and *Many cleanups
    EvalSum/EvalSumRows/EvalSumCols fold at the same radix. Their doubling lives in the automorphism-index group (generator g ∈ {5, 5^rowSize, 5⁻¹}, squared per level); the radix-r generalization uses g, g², …, g^(r−1) within a level and advances g ← g^r. All four folds route through one shared EvalSumRadixFold, which under HYBRID key switching accumulates lazily in the extended basis exactly like EvalPartialSumInPlace; a hoisted eager fallback covers non-HYBRID via a new SchemeBase::EvalAutomorphismCore wrapper (EvalSum keys live in their own key map, so EvalFastRotation's key lookup does not apply). The fully-packed BFV/BGV case keeps the legacy doubling (its tail is the conjugation-like automorphism m−1). The lazy accumulation changes EvalSum output bits at every radix — the same class of change as commit 2. Measured vs baseline (single-threaded Ice Lake, ring 2¹⁴–2¹⁶): ≈1.15× at radix 2 from laziness alone, ≈1.5–1.65× at radix 4; radix 8 consistently trails radix 4. EvalSumKeyGen emits the radix's exact index set; EvalSum keys were never shared with rotation keys, so radix 4 costs +50% of an already-dedicated log-scale set.

Multiparty keygen deduplication. MultiEvalSumKeyGen had its own inline copy of the index-set loop (BFV-form only — a silent-drift risk under any radix change); it now calls AdvancedSHEBase::GenerateIndexListForEvalSum (shared via a friend declaration rather than widening the public API), so joint keys match single-party keygen at any radix. MultiEvalAtIndexKeyGen similarly replaced a hardcoded 2n/2nComplex dispatch with the scheme's virtual FindAutomorphismIndex.

EvalAddMany/EvalAddManyInPlace. Addition is exact and associative, so the pairing tree bought nothing: EvalAddMany is now a null-tolerant left fold with a single clone (≈1.5–2× vs baseline, growing with tower count), and EvalAddManyInPlace an actually-in-place serial fold with zero allocations (≈1.8–2.3×) whose result lands in slot 0 per the in-place contract; null-leading vectors work for the first time.

Size-1 aliasing fix. The cryptocontext.h wrappers for EvalAddMany/EvalMultMany returned the caller's input handle for single-element vectors, so mutating the "result" silently corrupted the input. Results never alias inputs now: EvalAddMany's early return is dropped (the scheme clone covers all sizes), EvalMultMany's becomes a keyless Clone().

EvalMultMany unification. The base tree picks levelsToDrop from GetCompositeDegree() (1 everywhere except composite-scaling CKKS), making the AdvancedSHECKKSRNS override redundant; it is removed. The unified fold keeps the balanced tree — depth bounds noise growth, unlike the adds — clones for size 1 (also fixing an empty-scratch-vector UB at the scheme layer), and frees each consumed partial immediately: live intermediates drop from n−1 to ≈n/2, measured −46% peak scratch memory at n=32 with unchanged runtime and pairing order (bit-identical).

@pascoec pascoec added this to the Release 1.6.0 milestone Jul 18, 2026
@pascoec pascoec self-assigned this Jul 18, 2026
@pascoec pascoec added the optimization Improves performance label Jul 18, 2026
@pascoec
pascoec requested a review from yspolyakov July 18, 2026 05:47
Comment thread src/pke/lib/scheme/ckksrns/ckksrns-fhe.cpp
pascoec added 2 commits July 31, 2026 03:08
…k the *Many operations

Renames CKKS_PARTIAL_SUM_RADIX to PARTIAL_SUM_RADIX: the knob is no
longer CKKS-specific because it now also drives the scheme-generic
EvalSum/EvalSumRows/EvalSumCols family. All four folds route through one
EvalSumRadixFold that, under HYBRID key switching, accumulates lazily in
the extended basis exactly like EvalPartialSumInPlace (one deferred
ApproxModDown for element 0, per-level settle for element 1, one digit
decomposition per level); a hoisted eager fallback covers non-HYBRID via
a new SchemeBase::EvalAutomorphismCore wrapper, needed because EvalSum
keys live in their own key map.

EvalAddMany becomes a null-tolerant left fold with a single clone
(addition is exact and associative, so the pairing tree bought nothing;
~1.5-2x vs dev, growing with tower count) and EvalAddManyInPlace an
actually-in-place serial fold with zero allocations (~1.8-2.3x); its
result lands in slot 0 per the in-place contract, and null-leading
vectors work for the first time. Fixes a size-1 aliasing bug in the
cryptocontext wrappers: EvalAddMany/EvalMultMany returned the caller's
input handle for single-element vectors, so mutating the result silently
corrupted the input; results never alias inputs now (EvalAddMany's early
return is dropped, EvalMultMany's becomes a keyless Clone).

Unifies EvalMultMany: the base tree picks levelsToDrop from
GetCompositeDegree() (1 everywhere except composite-scaling CKKS),
making the AdvancedSHECKKSRNS override redundant; it is removed. The
unified fold keeps the three-phase structure, clones for size 1 (also
fixing an empty-scratch-vector UB at the scheme layer), and frees each
consumed partial immediately: live intermediates drop from n-1 to ~n/2,
measured -46% peak scratch at n=32 with unchanged runtime and pairing
order (bit-identical).
The fully-packed BFV/BGV branch kept the legacy doubling index set, which
no longer covers EvalSum calls at smaller batch sizes once those fold at
the configured radix (radix-2 sets were nested across all batch sizes;
e.g. EvalInnerProduct with default keygen requested key 5^3). The
fully-packed path now radix-folds its first phase (covering batchSize/2)
with only the m-1 conjugation tail kept special, in both keygen and the
fold, making the fully-packed key set a superset of every smaller batch's
set again.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

optimization Improves performance

Projects

None yet

2 participants