[JaxiteWord] Update lowering and emission for CROSS polynomial APIs - #3251
Open
Zohaib58 wants to merge 10 commits into
Open
[JaxiteWord] Update lowering and emission for CROSS polynomial APIs#3251Zohaib58 wants to merge 10 commits into
Zohaib58 wants to merge 10 commits into
Conversation
…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.
Use CROSS's level-indexed Polynomial operations directly for add, sub, multiply, rotation, rescale, relinearization, encryption, and decryption. Remove legacy layout-normalization templates and raw-array compatibility glue, and add a context-aware binary emitter helper. Lower CKKS multiplication to mul_no_relin, remove obsolete multiplication depth configuration, and emit output_scale in CKKS parameters. Remove unsupported plaintext add/sub and negate operations, and update tests for the new CROSS interface.
Zohaib58
changed the title
Zohaib58
changed the title
asraa
approved these changes
Jul 30, 2026
asraa
left a comment
asraa
left a comment
Collaborator
There was a problem hiding this comment.
thank you! just a few minor comments
| let summary = "Homomorphic addition of two ciphertexts"; | ||
| } | ||
|
|
||
| def AddPlainOp : JaxiteWord_Op<"add_plain", [Pure]> { |
Collaborator
There was a problem hiding this comment.
I know many of the ops in this file are also missing traits, but for the LWE ciphertext/plaintext types, adds should probably expect SameOperandsAndResultRings trait
| if (!dynamicShift) { | ||
| return failure(); | ||
| } | ||
| auto constOp = dynamicShift.getDefiningOp<arith::ConstantOp>(); |
Collaborator
There was a problem hiding this comment.
you could use a matcher:
heir/lib/Analysis/RotationAnalysis/DagBuilder.cpp
Lines 49 to 50 in e6bc99f
| } | ||
|
|
||
| bool funcNeedsCryptoContextAndKeys(func::FuncOp funcOp) { | ||
| return containsDialects<lwe::LWEDialect, ckks::CKKSDialect, bgv::BGVDialect>( |
Collaborator
There was a problem hiding this comment.
should you really have bgv dialect here?
Contributor
Author
Thanks for the review! I’ve addressed all your comments. |
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
2 participants
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.
This PR updates the JaxiteWord lowering pipeline, crypto-context configuration, and Python emitter to use CROSS’s current level-indexed polynomial APIs.
Changes