Determine whether modreduce require linearized input in optimize-relinearization
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
For the following example, the `--optimize-relinearization` will produce a `modreduce + relinearize` sequence
```mlir
func.func @foo(%a: i64, %b: i64) -> (i64) {
%0 = arith.muli %a, %a : i64
%1 = arith.muli %b, %b : i64
%2 = arith.subi %0, %1 : i64
%3 = arith.muli %2, %2 : i64
func.return %3 : i64
}
```
```mlir
%1 = arith.muli %input0, %input0 {addCount = 1 : i64, mgmt.mgmt = #mgmt.mgmt} : i64
%2 = arith.muli %input1, %input1 {addCount = 1 : i64, mgmt.mgmt = #mgmt.mgmt} : i64
%3 = arith.subi %1, %2 {addCount = 2 : i64, keySwitchCount = 2 : i64, mgmt.mgmt = #mgmt.mgmt} : i64
%4 = mgmt.modreduce %3 {mgmt.mgmt = #mgmt.mgmt} : i64
%5 = mgmt.relinearize %4 {mgmt.mgmt = #mgmt.mgmt} : i64
%6 = arith.muli %5, %5 {addCount = 1 : i64, mgmt.mgmt = #mgmt.mgmt} : i64
%7 = mgmt.modreduce %6 {mgmt.mgmt = #mgmt.mgmt} : i64
%8 = mgmt.relinearize %7 {mgmt.mgmt = #mgmt.mgmt} : i64
secret.yield %8 : i64
```
However, the current noise model explicitly disallow such behavior by an assertion, see https://github.com/google/heir/blob/d7fbf060e312b6c65fc5af70a37d3085a4512bf5/lib/Analysis/NoiseAnalysis/BGV/NoiseByBoundCoeffModel.cpp#L184
The reason inside noise model is that, for higher degree modreduce the formula need some correction because `s^n` is dependent then the bound becomes `correctionFactor * expansionFactor^(n - 1) * ||s||^n`. For `n = 2` the correction factor is 2, but for higher degree this grows rapidly. See https://ia.cr/2023/600 for more details.
Another intuitive reason from the noise perspective is that, based on the noise model AND our current `secret-insert-mgmt` strategy where modreduce is just before next `arith.mul` or `secret.yield`, `modreduce + relinearize (+ mul / secret.yield)` would produce worse noise growth than `relinearize + modreduce (+ mul / secret.yield)`.
From the parameter generation optimization perspective: For `relinearize + modreduce`, we can make the noise of `relinearize` much bigger (as the same level as multiplication) by making the special moduli chain (i.e. `P`) samller, so the `logPQ` can be smaller, which in turn could make `logN` smaller, gaining performance.
There are also reasons not supporting `relinearize + modreduce`, like from the performance perspective, indeed `modreduce + relinearize` will be less computationally inexpensive than `relinearize + modreduce`.
Currently for the ease of #1379 I just relinearize the input before mod reduce. When we introduce the cost model we can come back to decide which is better. (Openfhe have a script outputing the cost model, see https://github.com/ZenithalHourlyRate/openfhe-development/blob/noise_budget/benchmark/src/lib-benchmark.cpp)
Contributor guide
Assessment
This issue has not been assessed yet.