integrate CKKS with TPU kernels into HEIR

google / heir

A compiler for homomorphic encryption

https://heir.dev/

Apache License 2.0

318 stars 47 forks source link

integrate CKKS with TPU kernels into HEIR #1065

Open JianmingTONG opened 2 hours ago

JianmingTONG commented 2 hours ago

integrate CKKS into HEIR

AlexanderViand-Intel commented 2 hours ago

Welcome to HEIR! HEIR already has support for CKKS (there's a CKKS dialect, we have a --secret-to-ckks pass, and there's support for emitting OpenFHE code for CKKS. I think there's also an end-to-end example in the demo notebook). Feel free to edit & re-open this issue or create a new issue if there's a specific feature/capability that you're missing!

EDIT: re-opened as I just saw asraa update the title to be more specific about CKKS-on-TPU

asraa commented 2 hours ago

yes! we're in a meeting right now but the intent for this was to integrate TPU kernels into HEIR which has a few routes:

Lower through a shell API like jaxite (implement CKKS to jaxite and emit python code)
Lower CKKS to polynomial, and then polynomial to Mosaic MLIR which can be compiled and scheduled for TPU (https://github.com/jax-ml/jax/discussions/16855)

Will update more!