Rotation Keys & Composite Rotations
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
In RLWE-schemes (B/FV, BGV, CKKS), rotations over packed/batched (SIMD) ciphertexts require a key switching operation to make the resulting ciphertext decrypt correctly under a non-permuted key, which requires a key-switching-key (ksk) for each supported rotation. These keys are also commonly called "rotation keys" or "galois keys" and, for most applications, make up the vast majority of the "evaluation keys".
While it would be possible to generate a rotation key for every possible rotation -(N-1), ... ,-1,1, ... ,N-1 this would create in a prohibitively large amount of key material. Instead, existing FHE libraries generally default to the strategy of generating keys for rotations +/- 2^k and then assemble all rotations from those. For example, a rotation by 9 would be realized as "8, 1"; a rotation by 15 as "16, -1"; a rotation by 58 as "64, -8, 2". Note that the availability of the negative rotation offsets leads to significantly more efficient paths than if we merely had the positive rotations available. Clearly, there are multiple ways to get to each number, but (afaik) libraries determine a unique path by computing the [Non-Adjacent Form](https://en.wikipedia.org/wiki/Non-adjacent_form) of the desired rotation offset. More specifically, it seems like the [balanced NAF](https://eprint.iacr.org/2021/1161) is the best approach.
The +/-2^k approach is a great default solution, yet for many programs (e.g., kernels of a fixed size, as in the SIMD vectorizer examples) we end up needing a much smaller number of unique rotation offsets and can gain a lot by generating only the keys for those exact rotations. This is doubly advantageous, as it both (a) reduces the amount of key material, speeding up keygen and reducing memory load and (b) allows us to realize all rotations with a single "native" rotation (i.e., all rotation paths are length one), speeding up the homomorphic computation.
Targets such as SEAL/OpenFHE hide this complexity behind their "rotate" API and, afaik, are happy as long as the rotation requested can be realized from the keys available. However, for targets such as LLVM/x86, HW accelerators, and anything else that doesn't have this internal logic, the compiler needs to do this translation. In addition, we might want to manually control how rotations are realized even for OpenFHE/etc, for example when our compiler can provide a better solution than the built-in defaults.
Effectively, this requires the following:
- [ ] Add a lowering from arbitrary rotation offsets to a specific subset
- [ ] Add an analysis that collects all required rotation* keys
*(technically, in order to support multi-key/advanced use cases, the analysis should probably collect all key-switching-keys for a given base key)
- [ ] Add a heuristic /decision logic for which approach to use (individual keys vs powers-of-two approach) that can be overriden/configured via pass parameter(s)
- [ ] Implement a (balanced) NAF generator and associated lowering
- [ ] Tie this into the OpenFHE* context/keygen logic
*(it might be interesting to create something equivalent that works across various targets)
PS: Clearly, there is the usual circular dependency issue here w.r.t. optimization: given a program, we can easily find the optimal set of keys, and given a constrained set of keys, we can (probably?) find an optimal(-ish) program. However, given a high-level algorithm, it's not clear it's feasible to find the best program+keys for non-toy-sized prgorams. Given what we do/support now in terms of SIMD, following the program->rotation keys approach seems more reasonable right now.
PPS: This discussion was prompted by #742 (`bgv.rotate`'s rotation index is now an Attribute, i.e., must be statically known). While I agree with the change, there's a bit of a false dichotomy in the description: the opposite of a statically known rotation index isn't necessarily a blind rotate - it could still be a plaintext rotation index, just one that's only available at runtime. However, this case, if ever needed, should probably be handled by code-generating dynamic logic that uses a NAF-based approach anyway, as that's target independent.
Contributor guide
Assessment
This issue has not been assessed yet.