google / google/heir

Decouple `LevelAnalysis` and `MulResultAnalysis` from `secret.generic`

Open
#1,332 2 comments 0 reactions 0 assignees View on GitHub
dialect: mgmt
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

Currently, the `--secret-insert-mgmt-{bgv,ckks}` will only correctly insert modreduce operations into an IR if it has been previously `--wrap-generic`'d. Since I'd like to move `--wrap-generic` to a later point in the pipeline for readability/debuggability, I'd love to change this (but couldn't figure out how to do it in the 10min I gave myself, therefore this issue)

I think the culprit(s) here are `LevelAnalysis` (for the lack of modreduce) and `MulResultAnalysis` (not shown here, but `--optimize-relinearization` is a no-op if the function isn't wrapped in `secret.generic`). I assume they are also the reason `--annotate-mgmt` doesn't do anything for non-wrapped code.

Example:
```llvm
func.func @baz2(%x: i16 {secret.secret}, %y: i16 {secret.secret}, %z: i16) -> (i16) {
%0 = arith.muli %x, %z : i16
%1 = arith.muli %x, %x : i16
%2 = arith.muli %y, %y : i16
%3 = arith.addi %1, %2 : i16
%4 = arith.addi %0, %3 : i16
func.return %4 : i16
}
```
will be transformed into the followign (only relin's inserted, no modreduce):
```llvm
func.func @baz2(%arg0: i16 {secret.secret}, %arg1: i16 {secret.secret}, %arg2: i16) -> i16 {
%0 = arith.muli %arg0, %arg2 : i16
%1 = arith.muli %arg0, %arg0 : i16
%2 = mgmt.relinearize %1 : i16
%3 = arith.muli %arg1, %arg1 : i16
%4 = mgmt.relinearize %3 : i16
%5 = arith.addi %2, %4 : i16
%6 = arith.addi %0, %5 : i16
return %6 : i16
}
```
But
```llvm
func.func @baz2(%arg0: !secret.secret, %arg1: !secret.secret, %arg2: i16) -> !secret.secret {
%0 = secret.generic ins(%arg0, %arg1, %arg2 : !secret.secret, !secret.secret, i16) {
^body(%input0: i16, %input1: i16, %input2: i16):
%1 = arith.muli %input0, %input2 : i16
%2 = arith.muli %input0, %input0 : i16
%3 = arith.muli %input1, %input1 : i16
%4 = arith.addi %2, %3 : i16
%5 = arith.addi %1, %4 : i16
secret.yield %5 : i16
} -> !secret.secret
return %0 : !secret.secret
}
```
will be correctly converted to
```llvm
func.func @baz2(%arg0: !secret.secret, %arg1: !secret.secret, %arg2: i16) -> !secret.secret {
%0 = secret.generic ins(%arg0, %arg1, %arg2 : !secret.secret, !secret.secret, i16) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt}, arg1 = {mgmt.mgmt = #mgmt.mgmt}, arg2 = {mgmt.mgmt = #mgmt.mgmt}} {
^body(%input0: i16, %input1: i16, %input2: i16):
%1 = arith.muli %input0, %input2 {mgmt.mgmt = #mgmt.mgmt} : i16
%2 = arith.muli %input0, %input0 {mgmt.mgmt = #mgmt.mgmt} : i16
%3 = mgmt.relinearize %2 {mgmt.mgmt = #mgmt.mgmt} : i16
%4 = arith.muli %input1, %input1 {mgmt.mgmt = #mgmt.mgmt} : i16
%5 = mgmt.relinearize %4 {mgmt.mgmt = #mgmt.mgmt} : i16
%6 = arith.addi %3, %5 {mgmt.mgmt = #mgmt.mgmt} : i16
%7 = arith.addi %1, %6 {mgmt.mgmt = #mgmt.mgmt} : i16
%8 = mgmt.modreduce %7 {mgmt.mgmt = #mgmt.mgmt} : i16
secret.yield %8 : i16
} -> !secret.secret
return %0 : !secret.secret
}
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.