Add helper functions for server-side packing of plaintext function arguments
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
Given an input MLIR function that has secret and non-secret arguments, today our packing pipeline will insert `assign_layout` operations for all things that get packed (as plaintexts or ciphertexts), and then for secret operands those `assign_layout` operations will get lifted into func arg attributes (`tensor_ext.layout`) which are then read by `add-client-interface` and converted to concrete packing code before encryption.
We currently have no such support for packing a plaintext. I believe today the `assign_layout` operations for a plaintext function argument will be left in place, and then be lowered to something like `openfhe.make_ckks_packed_plaintext`, which will be recomputed on every call to the compiled function.
For ML model inference demos, this turns out to be a big pain because:
1. It makes inference slower by JIT-packing when it could pre-compute the packed weight matrices
2. It makes compile time slow by introduced a huge number of stack variables for weights (some MLP model @code-perspective was working with took some 2-3 hours to compile due to this)
So this issue is to insert plaintext packing helpers.
I think the easiest way to do this is to update `add-client-interface`, since the work is mostly done already there for secret func args. So the work in that case would be:
1. Add a new branch to `AddClientInterface::convertFunc` to support non-secret func arguments that have layouts. This should then call a new function similar to `AddClientInterface:: generateEncryptionFunc` but for the plaintext packing step. This function should be annotated similarly to the current encryption helpers so we can identify which plaintext argument the helper is for.
2. While `AddClientInterface:: generateEncryptionFunc` uses `secret.conceal` (after packing) which is later lowered to `encode + encrypt`, we need to have an op that will be lowered to just `encode` and represented as a plaintext type. We have #1875 which is to augment `secret.conceal` to support "trivial encryptions" (with zero noise). We should decide whether to use that here: have a "trivial encryption" be interpreted as a plaintext packing (and scrap #1875) or else have some other flag on `secret.conceal` that stops after encoding and doesn't encrypt. Either that, or have a new `secret` op that corresponds to just encoding.
3. Whatever happens in (2), update `lib/Dialect/Secret/Conversions/Patterns::ConvertClientConceal` to support the new option, or else add a new pattern to lower the new op to just encode. Note that in this case no secret key is required.
4. Update `heir_py` to generate new callables for the plaintext packing helpers, and call them and thread them through the compiled function. This would be happening mostly [here](https://github.com/google/heir/blob/2b1732d50c2c3e89c779de423ffd15cb87f8a91d/frontend/heir/backends/openfhe/backend.py#L303) and I think the rest of the pipeline should emit the new functions without any changes.
5. Some e2e tests may be manually packing plaintexts, so we should look to see if any can be simplified by using the newly compiled helper.
Contributor guide
Assessment
This issue has not been assessed yet.