[do-not-merge] add Lattigo backend support and end-to-end CKKS execution fixes - #3400
Open
copybara-service[bot] wants to merge 1 commit into
Open
copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
copybara-service
Bot
force-pushed
the
test_976505644
branch
2 times, most recently
from
September 5, 2026 00:46
bdca9c0 to
b823cfb
Compare
…ion fixes
Introduce backend emitter features, parameter generation adjustments, and
pipeline controls required for compiling and executing deep CKKS Transformer
models (e.g. BERT-tiny SST-2) on the Lattigo backend:
- Strided subview lowering & memref copy codegen:
- Support non-contiguous memref.subview operations resulting from tensor.pad
bufferization by aliasing the whole source slice in Go and tracking
aliases.
- Lower memref.copy between strided and contiguous views into nested Go loops
with explicit flattened multi-dimensional affine index arithmetic.
- Wide-message bootstrapping & declared-scale bundles:
- Introduce --lattigo-bootstrap-declared-scale-log2 in LattigoEmitter to wrap
emitted bootstraps in declared-scale bundles for wide-range messages.
- Temporarily declare input ciphertext scale as Delta * 2^s so Lattigo's
ScaleDown brings the message into the Mod1 Chebyshev approximation domain
(|m| <= 2^s).
- Restore the true message scale post-bootstrap via exact integer scalar
multiplication by 2^s without consuming ciphertext levels or rescaling.
- Cache module parameter lookups in LattigoEmitter to validate first modulus
bit-depth and eliminate redundant AST walks.
- 61-bit prime generation & deep-circuit RNS decomposition:
- Implement lastPrimeCapAware with a 64-bit deterministic Miller-Rabin test
in CKKS/Params.cpp to find first primes up to 61 bits congruent to 1 mod
2N, bypassing OpenFHE's 60-bit MAX_MODULUS_SIZE limit.
- Update computeDnum in RLWEParams.cpp to select dnum=4 for deep circuits
(level >= 16), sizing auxiliary P primes compactly to satisfy 128-bit
quantum security standards.
- Trivial zero plaintext sizing:
- In ImplementTrivialEncryptionAsAddition, size trivially-encrypted zero
constants using the physical slot count (scheme.actual_slot_count) instead
of the logical requested slot count.
- Prevent sparse zero plaintexts on rings expanded for security constraints.
- Pipeline & build configuration:
- Add alloc-to-in-place option to ToLattigoPipelineOptions in
ArithmeticPipelineRegistration, enabling circuits with complex shift
networks to bypass in-place reuse hazards.
- Declare ctx.executable._heir_translate_binary in tools attribute of
heir-translate.bzl for correct Bazel tool tracking.
- Testing:
- Add lit tests for strided memref.subview copies with nested Go loops in
memref.mlir.
- Add lit tests validating default bootstrap and declared-scale bundle
generation under --lattigo-bootstrap-declared-scale-log2 in
bootstrap.mlir.
- Add regression lit test validating zero plaintext sizing against physical
actual_slot_count in
implement_trivial_encryption_as_addition_actual_slot_count.mlir.
PiperOrigin-RevId: 976505644
copybara-service
Bot
force-pushed
the
test_976505644
branch
from
September 5, 2026 01:38
b823cfb to
4dc9d5a
Compare
copybara-service Bot
pushed a commit
that referenced
this pull request
Sep 5, 2026
…ipeline Introduce the LowerPadToIdentityMatmul pass and pipeline controls to lower zero-padding operations on secret tensors into cleartext identity matrix multiplications, avoiding irregular shift networks during layout conversion. This change is stacked on top of all experimental pending commits: - cl/974067900 (or GitHub PR #3388): cyclic layout support for Transformers - cl/976505644 (or GitHub PR #3400): Lattigo backend support and CKKS fixes - cl/976542552 (or GitHub PR #3402): approximate non-linearities and management - LowerPadToIdentityMatmul pass: - Rewrite static zero tensor.pad ops on secret tensors to explicit linalg.matmul or linalg.batch_matmul with cleartext identity diagonal matrices (P_L * X for row padding, X * P_R for column padding). - Handle 2D tensors ([M, K] -> [M', K']) and 3D batched tensors ([H, M, K] -> [H, M', K']) across row padding, column padding, and simultaneous 2D padding via sequential Left and Right PCMM. - Leverage structured BSGS rotate-and-reduce schedules (2*sqrt(n) Galois keys) of existing Left and Right PCMM kernels, completely avoiding irregular shift networks and eliminating hundreds of Galois rotation keys. - CKKS pipeline integration: - Introduce --enable-identity-pcmm-padding in ArithmeticPipelineRegistration to wire LowerPadToIdentityMatmul into the mlirToRLWE CKKS pipeline prior to layout propagation. - Follow identity PCMM lowering with canonicalization and CSE to fold intermediate structures. - Register createLowerPadToIdentityMatmul pass in heir-opt. - Waterline analysis & forced bootstrapping: - Support mgmt.force_bootstrap attribute in BootstrapWaterlineAnalysis to force bootstrap placement on critical join operations where level consumption must be reset. PiperOrigin-RevId: 976603963
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
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.
[do-not-merge] add Lattigo backend support and end-to-end CKKS execution fixes
Introduce backend emitter features, parameter generation adjustments, and
pipeline controls required for compiling and executing deep CKKS Transformer
models (e.g. BERT-tiny SST-2) on the Lattigo backend:
Strided subview lowering & memref copy codegen:
bufferization by aliasing the whole source slice in Go and tracking
aliases.
with explicit flattened multi-dimensional affine index arithmetic.
Wide-message bootstrapping & declared-scale bundles:
emitted bootstraps in declared-scale bundles for wide-range messages.
ScaleDown brings the message into the Mod1 Chebyshev approximation domain
(|m| <= 2^s).
multiplication by 2^s without consuming ciphertext levels or rescaling.
bit-depth and eliminate redundant AST walks.
61-bit prime generation & deep-circuit RNS decomposition:
in CKKS/Params.cpp to find first primes up to 61 bits congruent to 1 mod
2N, bypassing OpenFHE's 60-bit MAX_MODULUS_SIZE limit.
(level >= 16), sizing auxiliary P primes compactly to satisfy 128-bit
quantum security standards.
Trivial zero plaintext sizing:
constants using the physical slot count (scheme.actual_slot_count) instead
of the logical requested slot count.
Pipeline & build configuration:
ArithmeticPipelineRegistration, enabling circuits with complex shift
networks to bypass in-place reuse hazards.
heir-translate.bzl for correct Bazel tool tracking.
Testing:
memref.mlir.
generation under --lattigo-bootstrap-declared-scale-log2 in
bootstrap.mlir.
actual_slot_count in
implement_trivial_encryption_as_addition_actual_slot_count.mlir.