Skip to content

Commit 9525b3f

Browse files
copybara-githubcopybara-github
committed
Merge pull request #2031 from j2kun:more-pass-examples
PiperOrigin-RevId: 785450062
2 parents 62f8e08 + c8480b7 commit 9525b3f

File tree

11 files changed

+112
-110
lines changed

11 files changed

+112
-110
lines changed

‎lib/Transforms/AnnotateModule/AnnotateModule.td

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,7 @@ def AnnotateModule : Pass<"annotate-module"> {
2121
- `openfhe`
2222
- `lattigo`
2323

24-
Example: `--annotate-module="backend=openfhe scheme=ckks"`
25-
26-
```
27-
module attributes {backend.openfhe, scheme.ckks}
28-
...
29-
}
30-
```
24+
(* example filepath=tests/Transforms/annotate_module/pass.mlir *)
3125
}];
3226

3327
let options = [

‎lib/Transforms/GenerateParam/GenerateParam.td

Lines changed: 4 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -36,35 +36,12 @@ def GenerateParamBGV : Pass<"generate-param-bgv"> {
3636
This pass relies on the presence of the `mgmt` dialect ops to model
3737
relinearize/modreduce, and it relies on `mgmt.mgmt` attribute to determine
3838
the ciphertext level/dimension. These ops and attributes can be added by
39-
a pass like `--secret-insert-mgmt-bgv` and `--annotate-mgmt`.
39+
a pass like `--secret-insert-mgmt-bgv`.
4040

4141
User can provide custom scheme parameters by annotating bgv::SchemeParamAttr
4242
at the module level.
4343

44-
Example:
45-
46-
```mlir
47-
module {
48-
func.func @add(%arg0: !secret.secret<i16>) -> !secret.secret<i16> {
49-
%0 = secret.generic(%arg0 : !secret.secret<i16>) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {
50-
^body(%input0: i16):
51-
%1 = arith.addi %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : i16
52-
secret.yield %1 : i16
53-
} -> !secret.secret<i16>
54-
return %0 : !secret.secret<i16>
55-
}
56-
}
57-
```
58-
59-
After applying the pass, the module will be updated with the scheme parameters:
60-
61-
```mlir
62-
module attributes {bgv.schemeParam = #bgv.scheme_param<logN = 12, Q = [4294991873], P = [4295049217], plaintextModulus = 65537>} {
63-
func.func @add(%arg0: !secret.secret<i16>) -> !secret.secret<i16> {
64-
// same as above
65-
}
66-
}
67-
```
44+
(* example filepath=tests/Transforms/generate_param_bgv/doctest.mlir *)
6845
}];
6946

7047
let dependentDialects = [
@@ -128,30 +105,7 @@ def GenerateParamBFV : Pass<"generate-param-bfv"> {
128105
User can provide custom scheme parameters by annotating bgv::SchemeParamAttr
129106
at the module level. Note that we reuse bgv::SchemeParamAttr for BFV.
130107

131-
Example:
132-
133-
```mlir
134-
module {
135-
func.func @add(%arg0: !secret.secret<i16>) -> !secret.secret<i16> {
136-
%0 = secret.generic(%arg0 : !secret.secret<i16>) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {
137-
^body(%input0: i16):
138-
%1 = arith.addi %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : i16
139-
secret.yield %1 : i16
140-
} -> !secret.secret<i16>
141-
return %0 : !secret.secret<i16>
142-
}
143-
}
144-
```
145-
146-
After applying the pass, the module will be updated with the scheme parameters:
147-
148-
```mlir
149-
module attributes {bgv.schemeParam = #bgv.scheme_param<logN = 12, Q = [1152921504606994433], P = [1152921504607191041], plaintextModulus = 65537>} {
150-
func.func @add(%arg0: !secret.secret<i16>) -> !secret.secret<i16> {
151-
// same as above
152-
}
153-
}
154-
```
108+
(* example filepath=tests/Transforms/generate_param_bfv/doctest.mlir *)
155109
}];
156110

157111
let dependentDialects = [
@@ -192,30 +146,7 @@ def GenerateParamCKKS : Pass<"generate-param-ckks"> {
192146
User can provide custom scheme parameters by annotating bgv::SchemeParamAttr
193147
at the module level.
194148

195-
Example:
196-
197-
```mlir
198-
module {
199-
func.func @add(%arg0: !secret.secret<f16>) -> !secret.secret<f16> {
200-
%0 = secret.generic(%arg0 : !secret.secret<f16>) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {
201-
^body(%input0: f16):
202-
%1 = arith.addf %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : f16
203-
secret.yield %1 : f16
204-
} -> !secret.secret<f16>
205-
return %0 : !secret.secret<f16>
206-
}
207-
}
208-
```
209-
210-
After applying the pass, the module will be updated with the scheme parameters:
211-
212-
```mlir
213-
module attributes {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [36028797019389953], P = [36028797019488257], logDefaultScale = 45>} {
214-
func.func @add(%arg0: !secret.secret<f16>) -> !secret.secret<f16> {
215-
// same as above
216-
}
217-
}
218-
```
149+
(* example filepath=tests/Transforms/generate_param_ckks/doctest.mlir *)
219150
}];
220151

221152
let dependentDialects = [

‎lib/Transforms/PopulateScale/PopulateScale.td

Lines changed: 2 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -20,21 +20,7 @@ def PopulateScaleBGV : Pass<"populate-scale-bgv"> {
2020
to determine the scale for each ciphertext. Such annotation can be
2121
generated by the `generate-param-bgv` pass.
2222

23-
Example:
24-
```mlir
25-
module attributes {bgv.schemeParam = #bgv.scheme_param<logN = 13, Q = [67239937, 8796093202433], P = [8796093349889], plaintextModulus = 65537>, scheme.bgv} {
26-
func.func @mul(%arg0: !secret.secret<i16>) -> !secret.secret<i16> {
27-
%0 = secret.generic(%arg0 : !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 1, scale = 1>}) {
28-
^body(%input0: i16):
29-
%1 = arith.muli %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 1, dimension = 3, scale = 1>} : i16
30-
%2 = mgmt.relinearize %1 {mgmt.mgmt = #mgmt.mgmt<level = 1, scale = 1>} : i16
31-
%3 = mgmt.modreduce %2 {mgmt.mgmt = #mgmt.mgmt<level = 0, scale = 42541>} : i16
32-
secret.yield %3 : i16
33-
} -> (!secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0, scale = 42541>})
34-
return %0 : !secret.secret<i16>
35-
}
36-
}
37-
```
23+
(* example filepath=tests/Transforms/populate_scale/bgv/doctest.mlir *)
3824
}];
3925
}
4026

@@ -55,21 +41,7 @@ def PopulateScaleCKKS : Pass<"populate-scale-ckks"> {
5541

5642
The scaling factor is expressed in logarithm form.
5743

58-
Example:
59-
```mlir
60-
module attributes {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [36028797019389953, 35184372121601], P = [36028797019488257], logDefaultScale = 45>, scheme.ckks} {
61-
func.func @mul(%arg0: !secret.secret<f32>) -> !secret.secret<f32> {
62-
%0 = secret.generic(%arg0 : !secret.secret<f32> {mgmt.mgmt = #mgmt.mgmt<level = 1, scale = 45>}) {
63-
^body(%input0: f32):
64-
%1 = arith.mulf %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 1, dimension = 3, scale = 90>} : f32
65-
%2 = mgmt.relinearize %1 {mgmt.mgmt = #mgmt.mgmt<level = 1, scale = 90>} : f32
66-
%3 = mgmt.modreduce %2 {mgmt.mgmt = #mgmt.mgmt<level = 0, scale = 45>} : f32
67-
secret.yield %3 : f32
68-
} -> (!secret.secret<f32> {mgmt.mgmt = #mgmt.mgmt<level = 0, scale = 45>})
69-
return %0 : !secret.secret<f32>
70-
}
71-
}
72-
```
44+
(* example filepath=tests/Transforms/populate_scale/ckks/doctest.mlir *)
7345
}];
7446

7547
let options = [

‎tests/Transforms/generate_param_bfv/BUILD

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
load("//bazel:lit.bzl", "glob_lit_tests")
2+
3+
package(default_applicable_licenses = ["@heir//:license"])
4+
5+
glob_lit_tests(
6+
name = "all_tests",
7+
data = ["@heir//tests:test_utilities"],
8+
driver = "@heir//tests:run_lit.sh",
9+
test_file_exts = ["mlir"],
10+
)

‎tests/Transforms/generate_param_bfv/doctest.mlir

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// RUN: heir-opt --generate-param-bfv %s | FileCheck %s
2+
3+
// CHECK: module attributes {bgv.schemeParam = #bgv.scheme_param<logN = 12, Q = [2147565569], P = [2147573761], plaintextModulus = 65537>}
4+
module {
5+
func.func @add(%arg0: !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) -> (!secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) {
6+
%0 = secret.generic(%arg0: !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {
7+
^body(%input0: i16):
8+
%1 = arith.addi %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : i16
9+
secret.yield %1 : i16
10+
} -> (!secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>})
11+
return %0 : !secret.secret<i16>
12+
}
13+
}

‎tests/Transforms/generate_param_bgv/BUILD

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
load("//bazel:lit.bzl", "glob_lit_tests")
2+
3+
package(default_applicable_licenses = ["@heir//:license"])
4+
5+
glob_lit_tests(
6+
name = "all_tests",
7+
data = ["@heir//tests:test_utilities"],
8+
driver = "@heir//tests:run_lit.sh",
9+
test_file_exts = ["mlir"],
10+
)

‎tests/Transforms/generate_param_bgv/doctest.mlir

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// RUN: heir-opt --generate-param-bgv %s | FileCheck %s
2+
3+
// CHECK: module attributes {bgv.schemeParam = #bgv.scheme_param<logN = 12, Q = [4294991873], P = [4295049217], plaintextModulus = 65537>}
4+
module {
5+
func.func @add(%arg0: !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) -> (!secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) {
6+
%0 = secret.generic(%arg0: !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {
7+
^body(%input0: i16):
8+
%1 = arith.addi %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : i16
9+
secret.yield %1 : i16
10+
} -> (!secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>})
11+
return %0 : !secret.secret<i16>
12+
}
13+
}

‎tests/Transforms/generate_param_ckks/BUILD

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
load("//bazel:lit.bzl", "glob_lit_tests")
2+
3+
package(default_applicable_licenses = ["@heir//:license"])
4+
5+
glob_lit_tests(
6+
name = "all_tests",
7+
data = ["@heir//tests:test_utilities"],
8+
driver = "@heir//tests:run_lit.sh",
9+
test_file_exts = ["mlir"],
10+
)

‎tests/Transforms/generate_param_ckks/doctest.mlir

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// RUN: heir-opt --generate-param-ckks %s | FileCheck %s
2+
3+
// CHECK: {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [36028797019389953], P = [36028797019488257], logDefaultScale = 45>}
4+
module {
5+
func.func @add(%arg0: !secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) -> (!secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) {
6+
%0 = secret.generic(%arg0: !secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {
7+
^body(%input0: f16):
8+
%1 = arith.addf %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : f16
9+
secret.yield %1 : f16
10+
} -> (!secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>})
11+
return %0 : !secret.secret<f16>
12+
}
13+
}

‎tests/Transforms/populate_scale/bgv/doctest.mlir

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
// RUN: heir-opt %s --populate-scale-bgv | FileCheck %s
2+
3+
module attributes {bgv.schemeParam = #bgv.scheme_param<logN = 13, Q = [67239937, 8796093202433], P = [8796093349889], plaintextModulus = 65537>, scheme.bgv} {
4+
func.func @mul(%arg0: !secret.secret<i16>) -> !secret.secret<i16> {
5+
%0 = secret.generic(%arg0 : !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 1>}) {
6+
^body(%input0: i16):
7+
%1 = arith.muli %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 1, dimension = 3>} : i16
8+
// CHECK: mgmt.relinearize
9+
// CHECK-SAME: scale = 1
10+
%2 = mgmt.relinearize %1 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : i16
11+
// CHECK: mgmt.modreduce
12+
// CHECK-SAME: scale = 42541
13+
%3 = mgmt.modreduce %2 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : i16
14+
secret.yield %3 : i16
15+
} -> (!secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 0>})
16+
return %0 : !secret.secret<i16>
17+
}
18+
}

0 commit comments

Comments
 (0)