feat: release CKKS scheme caches and trim allocator#1179
Open
feat: release CKKS scheme caches and trim allocator#1179
Conversation
) FHECKKSRNS::m_bootPrecomMap is populated by EvalBootstrapSetup and EvalBootstrapPrecompute but never released by the existing ClearEval*Keys or ReleaseAllContexts APIs when a user still holds a CryptoContext reference, so the documented "clear" paths leave that memory alive for the lifetime of the process. Add a virtual FHEBase::ClearBootstrapPrecom() (default no-op), override it in FHECKKSRNS to clear m_bootPrecomMap, and expose matching forwarders on SchemeBase and CryptoContextImpl so callers can free the precomputations without destroying the context.
Extends the initial ClearBootstrapPrecom work so callers can actually drop every piece of scheme-level memory a CKKS CryptoContext retains, without destroying the context itself. - Add per-slot ClearBootstrapPrecom(slots) and ClearBootstrapPrecomExcept(keep) on FHEBase / SchemeBase / CryptoContextImpl so callers can evict a single transient precomputation or keep only a hot entry. - Add HasBootstrapPrecom() inspectors (whole-map and per-slot) for tests and memory-diagnostic code. - Introduce FHEBase::ClearSchemeSwitchPrecom() overridden by SWITCHCKKSRNS to drop the intermediate CKKS/BinFHE contexts, CKKS<->FHEW switching keys, auxiliary ciphertext, and the precomputed decoding-matrix plaintexts that EvalCKKStoFHEWSetup / EvalSchemeSwitchingSetup populate and which the existing Clear*Keys paths never released. - Add CryptoContextFactory::ReleaseAllContextsAndTrim() which extends ReleaseAllContexts() by also invoking the scheme-level clears on every tracked context and asking the allocator to return pages to the OS via the new TrimAllocator() helper (malloc_trim on glibc, no-op elsewhere). - Mark the cache-clear APIs noexcept so they are safe to call from destructors and shutdown paths. Adds UnitTestCKKSCacheClear (15 tests) covering per-slot / keep-except / full-release semantics, scheme-switch precom release, inspector agreement, and the ReleaseAllContextsAndTrim path. Adds ckks-release-memory example demonstrating the three common usage patterns. Closes #533
The test helper used modulus sizes and a scaling style that only work on 32/64-bit builds, so it threw on the 128-bit CI job. Pick larger moduli and FIXEDAUTO scaling when NATIVEINT == 128, matching what UnitTestSchemeSwitch.cpp already does. Production code is unchanged.
When the build links an alternative allocator like TCMalloc (WITH_TCM=ON), mallinfo2() and malloc_zone_statistics() return 0 because they only see the system malloc that TCMalloc replaced. The relative-drop check is unmeasurable in that case, so call GTEST_SKIP() instead of failing the assertion. The TrimAllocator contract test (TrimAllocatorIsSafe) still runs and pins the no-throw + capability-flag behavior on every backend.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #533. Lets callers drop scheme-level memory a CKKS
CryptoContextretains, without destroying the context itself.API additions
ClearBootstrapPrecom(slots)andClearBootstrapPrecomExcept(keep)onFHEBase/SchemeBase/CryptoContextImpl— evict a single slot's precomputation, or keep only one hot entryHasBootstrapPrecom()inspectors (whole-map and per-slot)FHEBase::ClearSchemeSwitchPrecom(), overridden inSWITCHCKKSRNS— drops the intermediate CKKS/BinFHE contexts, the CKKS<->FHEW switching keys, the auxiliary ciphertext, and the decoding-matrix plaintexts thatEvalCKKStoFHEWSetup/EvalSchemeSwitchingSetuppopulate (existingClear*Keyspaths never released these)CryptoContextFactory::ReleaseAllContextsAndTrim()— extendsReleaseAllContexts()by invoking the scheme-level clears on every tracked context, then calls a newTrimAllocator()helper (malloc_trimon glibc, no-op elsewhere) so the allocator can return pages to the OS