Skip to content

Commit deb3648

Browse files
eymayeymay
committed
Strip compilation target from unnecessary builds with new target
1 parent d36a954 commit deb3648

5 files changed

Lines changed: 194 additions & 10 deletions

File tree

‎.github/workflows/release.yml‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,17 @@ jobs:
2121

2222
- name: "Run `bazel build`"
2323
run: |
24-
bazel build -c opt //tools:heir-opt //tools:heir-translate
24+
bazel build -c opt //tools:keymemrt-opt //tools:keymemrt-translate
2525
2626
- name: Ensure binaries can run
2727
shell: bash
2828
run: |
29-
bazel-bin/tools/heir-opt --help
30-
bazel-bin/tools/heir-translate --help
29+
bazel-bin/tools/keymemrt-opt --help
30+
bazel-bin/tools/keymemrt-translate --help
3131
3232
- name: GH Release
3333
uses: softprops/action-gh-release@a74c6b72af54cfa997e81df42d94703d6313a2d0 #pin@v2.0.6
3434
with:
3535
files: |
36-
bazel-bin/tools/heir-opt
37-
bazel-bin/tools/heir-translate
36+
bazel-bin/tools/keymemrt-opt
37+
bazel-bin/tools/keymemrt-translate

‎README.md‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,10 @@ source ~/.bashrc
3232
3. Clone and build the repo:
3333
```shell
3434
git clone https://github.com/eymay/KeyMemRT-Compiler.git && cd KeyMemRT-Compiler/
35-
bazel build @heir//tools:heir-opt
36-
bazel build @heir//tools:heir-translate
35+
bazel build //tools:keymemrt-opt
36+
bazel build //tools:keymemrt-translate
3737
```
38-
KeyMemRT Compiler is a series of passes that can be called from the newly built driver `./bazel-bin/tools/heir-opt`.
38+
KeyMemRT Compiler is a series of passes that can be called from the newly built driver `./bazel-bin/tools/keymemrt-opt`.
3939

4040
4. Run KeyMemRT tests:
4141
```shell
@@ -68,8 +68,8 @@ KeyMemRT reuses these HEIR and MLIR passes, some with modifications:
6868

6969
An example pipeline to lower the HEIR `ckks` dialect to `openfhe` dialect with KeyMemRT key management is:
7070
```
71-
./bazel-bin/tools/heir-opt
72-
--ckks-to-lwe \
71+
./bazel-bin/tools/keymemrt-opt \
72+
--ckks-to-lwe \
7373
--lwe-to-openfhe \
7474
--lower-linear-transform \
7575
--symbolic-bsgs-decomposition \

‎tools/BUILD‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,70 @@ cc_binary(
216216
],
217217
)
218218

219+
# KeyMemRT-specific mlir-opt with only required dependencies
220+
cc_binary(
221+
name = "keymemrt-opt",
222+
srcs = ["keymemrt-opt.cpp"],
223+
includes = ["include"],
224+
deps = [
225+
# Core HEIR interfaces
226+
"@heir//lib/Dialect:HEIRInterfaces",
227+
# KeyMemRT dialects
228+
"@heir//lib/Dialect/CKKS/IR:Dialect",
229+
"@heir//lib/Dialect/CKKS/Conversions/CKKSToLWE",
230+
"@heir//lib/Dialect/LWE/IR:Dialect",
231+
"@heir//lib/Dialect/LWE/Conversions/LWEToOpenfhe",
232+
"@heir//lib/Dialect/LWE/Transforms",
233+
"@heir//lib/Dialect/KMRT/IR:Dialect",
234+
"@heir//lib/Dialect/KMRT/Transforms",
235+
"@heir//lib/Dialect/Openfhe/IR:Dialect",
236+
"@heir//lib/Dialect/Openfhe/Transforms",
237+
"@heir//lib/Dialect/Openfhe/Transforms:ConfigureCryptoContext",
238+
# Supporting dialects
239+
"@heir//lib/Dialect/ModArith/IR:Dialect",
240+
"@heir//lib/Dialect/Polynomial/IR:Dialect",
241+
"@heir//lib/Dialect/RNS/IR:Dialect",
242+
"@heir//lib/Dialect/RNS/IR:RNSTypeInterfaces",
243+
"@heir//lib/Dialect/Random/IR:Dialect",
244+
"@heir//lib/Dialect/TensorExt/IR:Dialect",
245+
# KeyMemRT transforms
246+
"@heir//lib/Transforms/AddRotationKeys",
247+
"@heir//lib/Transforms/AnnotateModule",
248+
"@heir//lib/Transforms/BootstrapRotationAnalysis",
249+
"@heir//lib/Transforms/LowerLinearTransform",
250+
"@heir//lib/Transforms/SymbolicBSGSDecomposition",
251+
"@heir//lib/Transforms/ProfileAnnotator",
252+
"@heir//lib/Transforms/UnnecessaryBootstrapRemoval",
253+
# MLIR core
254+
"@llvm-project//llvm:Support",
255+
"@llvm-project//mlir:AffineDialect",
256+
"@llvm-project//mlir:AffineToStandard",
257+
"@llvm-project//mlir:ArithDialect",
258+
"@llvm-project//mlir:FuncDialect",
259+
"@llvm-project//mlir:IR",
260+
"@llvm-project//mlir:MemRefDialect",
261+
"@llvm-project//mlir:MlirOptLib",
262+
"@llvm-project//mlir:Pass",
263+
"@llvm-project//mlir:SCFDialect",
264+
"@llvm-project//mlir:Support",
265+
"@llvm-project//mlir:TensorDialect",
266+
"@llvm-project//mlir:Transforms",
267+
],
268+
)
269+
270+
# KeyMemRT-specific mlir-translate with only required dependencies
271+
cc_binary(
272+
name = "keymemrt-translate",
273+
srcs = ["keymemrt-translate.cpp"],
274+
includes = ["include"],
275+
deps = [
276+
"@heir//lib/Target/Metadata:MetadataEmitter",
277+
"@heir//lib/Target/OpenFhePke:OpenFheRegistration",
278+
"@llvm-project//llvm:Support",
279+
"@llvm-project//mlir:TranslateLib",
280+
],
281+
)
282+
219283
cc_binary(
220284
name = "heir-lsp",
221285
srcs = ["heir-lsp.cpp"],

‎tools/keymemrt-opt.cpp‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
#include <cstdlib>
2+
#include <functional>
3+
#include <memory>
4+
#include <string>
5+
6+
// KeyMemRT dialects
7+
#include "lib/Dialect/HEIRInterfaces.h"
8+
#include "lib/Dialect/CKKS/Conversions/CKKSToLWE/CKKSToLWE.h"
9+
#include "lib/Dialect/CKKS/IR/CKKSDialect.h"
10+
#include "lib/Dialect/KMRT/IR/KMRTDialect.h"
11+
#include "lib/Dialect/KMRT/Transforms/Passes.h"
12+
#include "lib/Dialect/LWE/Conversions/LWEToOpenfhe/LWEToOpenfhe.h"
13+
#include "lib/Dialect/LWE/IR/LWEDialect.h"
14+
#include "lib/Dialect/LWE/Transforms/Passes.h"
15+
#include "lib/Dialect/ModArith/IR/ModArithDialect.h"
16+
#include "lib/Dialect/Openfhe/IR/OpenfheDialect.h"
17+
#include "lib/Dialect/Openfhe/Transforms/Passes.h"
18+
#include "lib/Dialect/Polynomial/IR/PolynomialDialect.h"
19+
#include "lib/Dialect/RNS/IR/RNSDialect.h"
20+
#include "lib/Dialect/RNS/IR/RNSTypeInterfaces.h"
21+
#include "lib/Dialect/Random/IR/RandomDialect.h"
22+
#include "lib/Dialect/TensorExt/IR/TensorExtDialect.h"
23+
24+
// KeyMemRT transforms
25+
#include "lib/Transforms/AddRotationKeys/AddRotationKeys.h"
26+
#include "lib/Transforms/AnnotateModule/AnnotateModule.h"
27+
#include "lib/Transforms/BootstrapRotationAnalysis/BootstrapRotationAnalysis.h"
28+
#include "lib/Transforms/LowerLinearTransform/LowerLinearTransform.h"
29+
#include "lib/Transforms/ProfileAnnotator/ProfileAnnotator.h"
30+
#include "lib/Transforms/SymbolicBSGSDecomposition/SymbolicBSGSDecomposition.h"
31+
#include "lib/Transforms/UnnecessaryBootstrapRemoval/UnnecessaryBootstrapRemoval.h"
32+
33+
// MLIR core
34+
#include "mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h" // from @llvm-project
35+
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project
36+
#include "mlir/include/mlir/Dialect/Affine/Passes.h" // from @llvm-project
37+
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
38+
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
39+
#include "mlir/include/mlir/Dialect/MemRef/IR/MemRef.h" // from @llvm-project
40+
#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h" // from @llvm-project
41+
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
42+
#include "mlir/include/mlir/Pass/PassManager.h" // from @llvm-project
43+
#include "mlir/include/mlir/Pass/PassRegistry.h" // from @llvm-project
44+
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
45+
#include "mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h" // from @llvm-project
46+
#include "mlir/include/mlir/Transforms/Passes.h" // from @llvm-project
47+
48+
using namespace mlir;
49+
using namespace heir;
50+
51+
int main(int argc, char **argv) {
52+
mlir::DialectRegistry registry;
53+
54+
// Register KeyMemRT dialects
55+
registry.insert<ckks::CKKSDialect>();
56+
registry.insert<kmrt::KMRTDialect>();
57+
registry.insert<lwe::LWEDialect>();
58+
registry.insert<mod_arith::ModArithDialect>();
59+
registry.insert<openfhe::OpenfheDialect>();
60+
registry.insert<::mlir::heir::polynomial::PolynomialDialect>();
61+
registry.insert<random::RandomDialect>();
62+
registry.insert<rns::RNSDialect>();
63+
registry.insert<tensor_ext::TensorExtDialect>();
64+
65+
// Register MLIR dialects
66+
registry.insert<affine::AffineDialect>();
67+
registry.insert<mlir::arith::ArithDialect>();
68+
registry.insert<func::FuncDialect>();
69+
registry.insert<memref::MemRefDialect>();
70+
registry.insert<scf::SCFDialect>();
71+
registry.insert<tensor::TensorDialect>();
72+
73+
// Register MLIR passes
74+
registerTransformsPasses(); // canonicalize, cse, etc.
75+
affine::registerAffinePasses(); // loop unrolling, lower-affine
76+
77+
// Register affine-to-standard conversion
78+
registerPass(
79+
[]() -> std::unique_ptr<Pass> { return createLowerAffinePass(); });
80+
81+
// Register KeyMemRT transforms
82+
kmrt::registerKMRTPasses();
83+
lwe::registerLWEPasses();
84+
openfhe::registerOpenfhePasses();
85+
registerAddRotationKeysPasses();
86+
registerAnnotateModulePasses();
87+
registerBootstrapRotationAnalysisPasses();
88+
registerLowerLinearTransformPasses();
89+
registerProfileAnnotatorPasses();
90+
registerSymbolicBSGSDecompositionPasses();
91+
registerUnnecessaryBootstrapRemovalPasses();
92+
93+
// Register KeyMemRT conversions
94+
ckks::registerCKKSToLWEPasses();
95+
lwe::registerLWEToOpenfhePasses();
96+
97+
// Register KeyMemRT interfaces
98+
rns::registerExternalRNSTypeInterfaces(registry);
99+
registerOperandAndResultAttrInterface(registry);
100+
101+
return asMainReturnCode(
102+
MlirOptMain(argc, argv, "KeyMemRT Pass Driver", registry));
103+
}

‎tools/keymemrt-translate.cpp‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#include "lib/Target/Metadata/MetadataEmitter.h"
2+
#include "lib/Target/OpenFhePke/OpenFheTranslateRegistration.h"
3+
#include "llvm/include/llvm/Support/LogicalResult.h" // from @llvm-project
4+
#include "mlir/include/mlir/Tools/mlir-translate/MlirTranslateMain.h" // from @llvm-project
5+
6+
int main(int argc, char **argv) {
7+
// Metadata output
8+
mlir::heir::registerMetadataEmitter();
9+
10+
// OpenFHE output
11+
mlir::heir::openfhe::registerTranslateOptions();
12+
mlir::heir::openfhe::registerToOpenFhePkeTranslation();
13+
mlir::heir::openfhe::registerToOpenFhePkeHeaderTranslation();
14+
mlir::heir::openfhe::registerToOpenFhePkePybindTranslation();
15+
16+
return failed(mlir::mlirTranslateMain(argc, argv, "KeyMemRT Translation Tool"));
17+
}

0 commit comments

Comments
 (0)