Add LWE to JaxiteWord conversion pass#2964
Conversation
j2kun
left a comment
j2kun
left a comment
There was a problem hiding this comment.
Looks great! Only some minor nitpicks
| } | ||
| if (dynamicShift) { | ||
| return rewriter.notifyMatchFailure( | ||
| op, "jaxiteword rotation requires static shift"); |
There was a problem hiding this comment.
This will end up being a bit of an obstacle for supporting looped linalg kernels. @JianmingTONG is this a strict limitation of CROSS?
No need to make any changes to this PR, I just want to know for tracking future work.
There was a problem hiding this comment.
Hey Jeremy
Thank you for the comment and wish you a great memorial weekend!
IIUC, the computation in HE is preknown ahead of execution, which could be used to generate all rotation keys. And this dynamicShift seems telling that we need some rotation indice to be generated in the runtime. Might I get some insights on what cases would we need this runtime dynamity? And correct me if I miss understand it
There was a problem hiding this comment.
If we express a program using loops, then the SSA value representing the desired shift is statically inferrable, but dynamically dependent on the loop iteration in question. So even if we know all the rotation keys at compile time, the operation syntax may still need to express a non-attribute for the shift operand.
There was a problem hiding this comment.
Aha got it!
It's the actual shift used at runtime that depends on which loop iteration is executing
I could confirm that is not the limitation of CROSS to support dynamic rotation, as it'll be the same API as OpenFHE
j2kun
left a comment
j2kun
left a comment
There was a problem hiding this comment.
LGTM! Please squash the commits (removing my co-authorship or replacing my email with jkun@google.com, I'm not sure why it doesn't recognize my account in the CLA...)
Register the pass in heir-opt and add lit tests with FileCheck coverage for lowered jaxiteword operations.
Zohaib58
force-pushed
the
feature/lwe-to-jaxiteword-pass
branch
from
89a124c to
cf2ccbf
Compare
j2kun
added
the
pull_ready
3 participants
Adds the lwe-to-jaxiteword conversion pass, which lowers the lwe dialect (and related ckks / bgv ops handled by the pass) to the jaxiteword dialect for CROSS / JaxiteWord codegen.
This is a follow-up to #2946 (JaxiteWord dialect, emitter, and configure-crypto-context). That PR landed the target dialect and Python emission; this PR wires the LWE → JaxiteWord lowering so HEIR pipelines can reach CROSS backend ops from scheme-level IR.
Changes
Notes
Intended to mirror patterns from existing LWE lowering passes (lwe-to-openfhe, lwe-to-polynomial).