Skip to content

Thread JaxiteWord keys through context generation and fix emitted ciphertext ops#3100

Open
Zohaib58 wants to merge 2 commits into
google:mainfrom
Zohaib58:optimize/heir-cross-integration
Open

Thread JaxiteWord keys through context generation and fix emitted ciphertext ops#3100
Zohaib58 wants to merge 2 commits into
google:mainfrom
Zohaib58:optimize/heir-cross-integration

Conversation

@Zohaib58

Copy link
Copy Markdown
Contributor

Change jaxiteword.gen_params to take public, secret, and evaluation keys, and
emit them into the generated ckks.CKKSParameters. Update the crypto-context
configuration pass so __generate_crypto_context accepts those three keys and
__configure_crypto_context only takes the crypto context.

Teach LWE-to-JaxiteWord lowering to add crypto context/eval-key arguments to
functions with crypto-typed arguments, update func.call operands for converted
callees, and lower CKKS rotations when a dynamic shift is defined by an arith
constant.

Rework the JaxiteWord Python emitter around tiled Polynomial values:

  • emit func.call operations
  • add _ensure_poly and _assign_poly helpers
  • normalize ciphertext operands before add/sub/mul/square/rotate/rescale/decrypt
  • use he_mul[level].mul for relined multiplication
  • emit he_rescale for modulus reduction across levels
  • implement plaintext add/sub against the c0 limb
  • generate rotation keys via key_gen.gen_rotation_key
  • store key material in generated CKKS params instead of mutating program init

Update JaxiteWord emitter/configure-context FileCheck tests and add 8x8 matvec
MLIR/Python artifacts for validating the emitted JaxiteWord path.

Zohaib58 added 2 commits June 19, 2026 04:46
…hertext ops

Change jaxiteword.gen_params to take public, secret, and evaluation keys, and
emit them into the generated ckks.CKKSParameters. Update the crypto-context
configuration pass so __generate_crypto_context accepts those three keys and
__configure_crypto_context only takes the crypto context.

Teach LWE-to-JaxiteWord lowering to add crypto context/eval-key arguments to
functions with crypto-typed arguments, update func.call operands for converted
callees, and lower CKKS rotations when a dynamic shift is defined by an arith
constant.

Rework the JaxiteWord Python emitter around tiled Polynomial values:
- emit func.call operations
- add _ensure_poly and _assign_poly helpers
- normalize ciphertext operands before add/sub/mul/square/rotate/rescale/decrypt
- use he_mul[level].mul for relined multiplication
- emit he_rescale for modulus reduction across levels
- implement plaintext add/sub against the c0 limb
- generate rotation keys via key_gen.gen_rotation_key
- store key material in generated CKKS params instead of mutating program init

Update JaxiteWord emitter/configure-context FileCheck tests and add 8x8 matvec
MLIR/Python artifacts for validating the emitted JaxiteWord path.
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