Skip to content

[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
mainfrom
test_976505644
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_976505644

Conversation

@copybara-service

Copy link
Copy Markdown
Contributor

[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:

    • 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.

@copybara-service
copybara-service Bot force-pushed the test_976505644 branch 2 times, most recently from bdca9c0 to b823cfb Compare September 5, 2026 00:46
…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 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant