Skip to content

Commit a52cab5

Browse files
j2kuncopybara-github
j2kun
authored and
copybara-github
committed
explicitly pass slotCount to MakeCKKSPackedPlaintext
PiperOrigin-RevId: 926877435
1 parent 5ec3d51 commit a52cab5

43 files changed

Lines changed: 496 additions & 125 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎lib/Dialect/LWE/Transforms/AddDebugPort.cpp‎

Lines changed: 72 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
#include <cassert>
44
#include <string>
5-
#include <utility>
65

76
#include "lib/Dialect/Debug/IR/DebugOps.h"
87
#include "lib/Dialect/LWE/IR/LWETypes.h"
@@ -63,9 +62,39 @@ FailureOr<Type> getPrivateKeyType(func::FuncOp op) {
6362
return lwePrivateKeyType;
6463
}
6564

65+
void populateDebugFuncCache(ModuleOp module,
66+
llvm::DenseMap<Type, func::FuncOp>& typeToDebugFunc,
67+
llvm::DenseSet<StringRef>& debugFuncNames) {
68+
for (auto funcOp : module.getOps<func::FuncOp>()) {
69+
if (!funcOp.isExternal()) continue;
70+
if (!funcOp.getName().starts_with("__heir_debug_")) continue;
71+
if (funcOp.getArgumentTypes().size() != 2) continue;
72+
if (!funcOp.getResultTypes().empty()) continue;
73+
74+
typeToDebugFunc[funcOp.getFunctionType()] = funcOp;
75+
debugFuncNames.insert(funcOp.getName());
76+
}
77+
}
78+
79+
static bool isAlreadyDebugged(Value value,
80+
const llvm::DenseSet<StringRef>& debugFuncNames) {
81+
for (auto& use : value.getUses()) {
82+
Operation* user = use.getOwner();
83+
if (isa<debug::ValidateOp>(user)) {
84+
return true;
85+
}
86+
if (auto callOp = dyn_cast<func::CallOp>(user)) {
87+
if (debugFuncNames.contains(callOp.getCallee())) {
88+
return true;
89+
}
90+
}
91+
}
92+
return false;
93+
}
94+
6695
func::FuncOp getOrCreateExternalDebugFunc(
67-
ModuleOp module, Type lwePrivateKeyType, Type valueType,
68-
llvm::DenseMap<Type, func::FuncOp>& typeToDebugFunc) {
96+
ModuleOp module, SymbolTable& symbolTable, Type lwePrivateKeyType,
97+
Type valueType, llvm::DenseMap<Type, func::FuncOp>& typeToDebugFunc) {
6998
auto* context = module.getContext();
7099
auto debugFuncType =
71100
FunctionType::get(context, {lwePrivateKeyType, valueType}, {});
@@ -75,28 +104,29 @@ func::FuncOp getOrCreateExternalDebugFunc(
75104
return it->second;
76105
}
77106

78-
int counter = typeToDebugFunc.size();
79-
std::string funcName = "__heir_debug_" + std::to_string(counter);
107+
unsigned uniquingCounter = typeToDebugFunc.size();
108+
SmallString<128> funcName = SymbolTable::generateSymbolName<128>(
109+
"__heir_debug",
110+
[&](StringRef name) { return symbolTable.lookup(name) != nullptr; },
111+
uniquingCounter);
80112

81-
// Assert that this name is not already in use.
82-
assert(!module.lookupSymbol<func::FuncOp>(funcName) &&
83-
"Symbol already exists");
84-
85-
ImplicitLocOpBuilder b =
86-
ImplicitLocOpBuilder::atBlockBegin(module.getLoc(), module.getBody());
87-
auto funcOp = func::FuncOp::create(b, funcName, debugFuncType);
113+
auto funcOp = func::FuncOp::create(module.getLoc(), funcName, debugFuncType);
88114
// required for external func call
89115
funcOp.setPrivate();
90116

117+
symbolTable.insert(funcOp, module.getBody()->begin());
118+
91119
typeToDebugFunc[debugFuncType] = funcOp;
92120
return funcOp;
93121
}
94122

95-
void insertValidationOps(func::FuncOp op) {
123+
void insertValidationOps(func::FuncOp op,
124+
const llvm::DenseSet<StringRef>& debugFuncNames) {
96125
int count = 0;
97126
auto insertValidate = [&](Value value, OpBuilder& b) {
98127
Type valueType = value.getType();
99128
if (isa<LWECiphertextType>(getElementTypeOrSelf(valueType))) {
129+
if (isAlreadyDebugged(value, debugFuncNames)) return;
100130
debug::ValidateOp::create(b, value.getLoc(), value,
101131
"heir_debug_" + std::to_string(count++),
102132
nullptr);
@@ -121,8 +151,8 @@ void insertValidationOps(func::FuncOp op) {
121151
}
122152

123153
LogicalResult lowerValidationOps(
124-
func::FuncOp op, Value privateKey, int messageSize,
125-
llvm::DenseMap<Type, func::FuncOp>& typeToDebugFunc) {
154+
func::FuncOp op, SymbolTable& symbolTable, Value privateKey,
155+
int messageSize, llvm::DenseMap<Type, func::FuncOp>& typeToDebugFunc) {
126156
auto module = op->getParentOfType<ModuleOp>();
127157
Type lwePrivateKeyType = privateKey.getType();
128158

@@ -141,10 +171,10 @@ LogicalResult lowerValidationOps(
141171
attrs.push_back(b.getNamedAttr(
142172
"message.size", b.getStringAttr(std::to_string(messageSize))));
143173

144-
auto debugFunc = getOrCreateExternalDebugFunc(module, lwePrivateKeyType,
145-
valueType, typeToDebugFunc);
146-
auto callOp =
147-
b.create<func::CallOp>(debugFunc, ArrayRef<Value>{privateKey, value});
174+
auto debugFunc = getOrCreateExternalDebugFunc(
175+
module, symbolTable, lwePrivateKeyType, valueType, typeToDebugFunc);
176+
auto callOp = func::CallOp::create(b, b.getLoc(), debugFunc,
177+
ArrayRef<Value>{privateKey, value});
148178
callOp->setDialectAttrs(attrs);
149179

150180
validateOp.erase();
@@ -166,6 +196,10 @@ struct AddDebugPort : impl::AddDebugPortBase<AddDebugPort> {
166196
ModuleOp module = cast<ModuleOp>(getOperation());
167197
SymbolTable symbolTable(module);
168198

199+
llvm::DenseMap<Type, func::FuncOp> typeToDebugFunc;
200+
llvm::DenseSet<StringRef> debugFuncNames;
201+
populateDebugFuncCache(module, typeToDebugFunc, debugFuncNames);
202+
169203
SmallVector<func::FuncOp, 16> worklist;
170204
llvm::DenseMap<func::FuncOp, Type> funcToKeyType;
171205
if (failed(identifyInitialTargets(module, symbolTable, funcToKeyType,
@@ -183,7 +217,7 @@ struct AddDebugPort : impl::AddDebugPortBase<AddDebugPort> {
183217

184218
if (insertDebugAfterEveryOp) {
185219
for (auto& [func, _] : funcToKeyType) {
186-
insertValidationOps(func);
220+
insertValidationOps(func, debugFuncNames);
187221
}
188222
}
189223

@@ -198,8 +232,8 @@ struct AddDebugPort : impl::AddDebugPortBase<AddDebugPort> {
198232
return;
199233
}
200234

201-
llvm::DenseMap<Type, func::FuncOp> typeToDebugFunc;
202-
if (failed(lowerAllValidationOps(module, funcToKeyType, typeToDebugFunc))) {
235+
if (failed(lowerAllValidationOps(module, symbolTable, funcToKeyType,
236+
typeToDebugFunc))) {
203237
signalPassFailure();
204238
return;
205239
}
@@ -229,18 +263,21 @@ struct AddDebugPort : impl::AddDebugPortBase<AddDebugPort> {
229263
}
230264

231265
if (entryFunc) {
232-
auto type = getPrivateKeyType(entryFunc);
233-
if (succeeded(type)) {
266+
bool shouldProcess =
267+
containsAnyOperations<debug::ValidateOp>(entryFunc) ||
268+
insertDebugAfterEveryOp;
269+
270+
if (shouldProcess) {
271+
auto type = getPrivateKeyType(entryFunc);
272+
if (failed(type)) {
273+
entryFunc.emitError(
274+
"Cannot infer LWE private key type for entry function");
275+
return failure();
276+
}
234277
funcToKeyType[entryFunc] = *type;
235278
worklist.push_back(entryFunc);
236-
return success();
237-
}
238-
239-
if (containsAnyOperations<debug::ValidateOp>(entryFunc)) {
240-
entryFunc.emitError(
241-
"Cannot infer LWE private key type for entry function");
242-
return failure();
243279
}
280+
return success();
244281
}
245282

246283
for (auto funcOp : module.getOps<func::FuncOp>()) {
@@ -404,7 +441,8 @@ struct AddDebugPort : impl::AddDebugPortBase<AddDebugPort> {
404441
/// \param typePairToInt Map to track generated debug function names.
405442
/// \return success() if successful, failure() otherwise.
406443
LogicalResult lowerAllValidationOps(
407-
ModuleOp module, const llvm::DenseMap<func::FuncOp, Type>& funcToKeyType,
444+
ModuleOp module, SymbolTable& symbolTable,
445+
const llvm::DenseMap<func::FuncOp, Type>& funcToKeyType,
408446
llvm::DenseMap<Type, func::FuncOp>& typeToDebugFunc) {
409447
for (auto funcOp : module.getOps<func::FuncOp>()) {
410448
if (funcOp.isExternal()) continue;
@@ -428,8 +466,8 @@ struct AddDebugPort : impl::AddDebugPortBase<AddDebugPort> {
428466
}
429467

430468
if (privateKey) {
431-
if (failed(lowerValidationOps(funcOp, privateKey, messageSize,
432-
typeToDebugFunc))) {
469+
if (failed(lowerValidationOps(funcOp, symbolTable, privateKey,
470+
messageSize, typeToDebugFunc))) {
433471
funcOp.emitError("failed to lower validation ops");
434472
return failure();
435473
}

‎lib/Dialect/LWE/Transforms/BUILD‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ cc_library(
2828
],
2929
deps = [
3030
":pass_inc_gen",
31-
"@heir//lib/Dialect:FuncUtils",
3231
"@heir//lib/Dialect/Debug/IR:Dialect",
3332
"@heir//lib/Dialect/LWE/IR:Dialect",
3433
"@heir//lib/Utils",

‎lib/Pipelines/ArithmeticPipelineRegistration.cpp‎

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -219,12 +219,9 @@ void mlirToPlaintextPipelineBuilder(OpPassManager& pm,
219219
mlirToRLWEPipelineOptions.ciphertextDegree = options.plaintextSize;
220220
mlirToSecretArithmeticPipelineBuilder(pm, mlirToRLWEPipelineOptions);
221221

222-
if (options.debug) {
223-
// Insert debug handler calls
224-
secret::SecretAddDebugPortOptions debugOptions;
225-
debugOptions.insertDebugAfterEveryOp = true;
226-
pm.addPass(secret::createSecretAddDebugPort(debugOptions));
227-
}
222+
// Insert debug handler calls and/or lower debug.validate
223+
pm.addPass(secret::createSecretAddDebugPort(secret::SecretAddDebugPortOptions{
224+
.insertDebugAfterEveryOp = options.debug}));
228225

229226
pm.addPass(secret::createSecretDistributeGeneric());
230227
pm.addPass(createCanonicalizerPass());
@@ -411,9 +408,15 @@ void mlirToRLWEPipeline(OpPassManager& pm,
411408
exit(EXIT_FAILURE);
412409
}
413410

411+
// Lower debug.validate ops to function calls with private key
412+
pm.addPass(lwe::createAddDebugPort(
413+
lwe::AddDebugPortOptions{.messageSize = (int)options.ciphertextDegree,
414+
.insertDebugAfterEveryOp = options.debug}));
415+
414416
pm.addPass(createForwardInsertToExtract());
415417
pm.addPass(createCanonicalizerPass());
416418
pm.addPass(createCSEPass());
419+
pm.addPass(createSymbolDCEPass());
417420

418421
// TODO(#2554): skip this pass if the backend supports trivial encryption
419422
pm.addPass(lwe::createImplementTrivialEncryptionAsAddition());
@@ -459,11 +462,11 @@ BackendPipelineBuilder toOpenFhePipelineBuilder() {
459462
pm.addPass(ckks::createCKKSToLWE());
460463

461464
// insert debug handler calls
462-
if (options.debug) {
463-
lwe::AddDebugPortOptions addDebugPortOptions;
464-
addDebugPortOptions.entryFunction = options.entryFunction;
465-
pm.addPass(lwe::createAddDebugPort(addDebugPortOptions));
466-
}
465+
lwe::AddDebugPortOptions addDebugPortOptions{
466+
.entryFunction = options.entryFunction,
467+
.insertDebugAfterEveryOp = options.debug,
468+
};
469+
pm.addPass(lwe::createAddDebugPort(addDebugPortOptions));
467470

468471
// Convert LWE (and scheme-specific CKKS/BGV ops) to OpenFHE
469472
pm.addPass(lwe::createLWEToOpenfhe());
@@ -501,11 +504,11 @@ BackendPipelineBuilder toLattigoPipelineBuilder() {
501504
pm.addPass(ckks::createCKKSToLWE());
502505

503506
// insert debug handler calls
504-
if (options.debug) {
505-
lwe::AddDebugPortOptions addDebugPortOptions;
506-
addDebugPortOptions.entryFunction = options.entryFunction;
507-
pm.addPass(lwe::createAddDebugPort(addDebugPortOptions));
508-
}
507+
lwe::AddDebugPortOptions addDebugPortOptions{
508+
.entryFunction = options.entryFunction,
509+
.insertDebugAfterEveryOp = options.debug,
510+
};
511+
pm.addPass(lwe::createAddDebugPort(addDebugPortOptions));
509512

510513
// Convert LWE (and scheme-specific BGV ops) to Lattigo
511514
pm.addPass(lwe::createLWEToLattigo());

‎lib/Pipelines/ArithmeticPipelineRegistration.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,10 @@ struct MlirToRLWEPipelineOptions : public LoopOptions {
9494
llvm::cl::desc(
9595
"The level budget excluding levels required for bootstrap"),
9696
llvm::cl::init(10)};
97+
PassOptions::Option<bool> debug{
98+
*this, "debug",
99+
llvm::cl::desc("Insert debug ports after every secret operation."),
100+
llvm::cl::init(false)};
97101
PassOptions::Option<std::string> plaintextExecutionResultFileName{
98102
*this, "plaintext-execution-result-file-name",
99103
llvm::cl::desc("File name to import execution result from (c.f. --secret-"

‎lib/Pipelines/BUILD‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ cc_library(
6060
"@heir//lib/Dialect/Debug/Transforms:ValidateNames",
6161
"@heir//lib/Dialect/LWE/Conversions/LWEToPolynomial",
6262
"@heir//lib/Dialect/Secret/Conversions/SecretToCGGI",
63+
"@heir//lib/Dialect/Secret/Transforms:AddDebugPort",
6364
"@heir//lib/Dialect/Secret/Transforms:DistributeGeneric",
6465
"@heir//lib/Transforms/AddClientInterface",
6566
"@heir//lib/Transforms/BooleanVectorizer",
@@ -102,7 +103,6 @@ cc_library(
102103
":PipelineRegistration",
103104
"@heir//lib/Dialect/BGV/Conversions/BGVToLWE",
104105
"@heir//lib/Dialect/CKKS/Transforms:CKKSToLWE",
105-
"@heir//lib/Dialect/Debug/Transforms",
106106
"@heir//lib/Dialect/Debug/Transforms:ValidateNames",
107107
"@heir//lib/Dialect/LWE/Conversions/LWEToLattigo",
108108
"@heir//lib/Dialect/LWE/Conversions/LWEToOpenfhe",

‎lib/Pipelines/BooleanPipelineRegistration.cpp‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,10 @@
99
#include "lib/Dialect/CGGI/Conversions/CGGIToSCIFRBool/CGGIToSCIFRBool.h"
1010
#include "lib/Dialect/CGGI/Conversions/CGGIToTfheRust/CGGIToTfheRust.h"
1111
#include "lib/Dialect/CGGI/Conversions/CGGIToTfheRustBool/CGGIToTfheRustBool.h"
12-
#include "lib/Dialect/Debug/Transforms/Passes.h"
1312
#include "lib/Dialect/Debug/Transforms/ValidateNames.h"
1413
#include "lib/Dialect/Secret/Conversions/SecretToCGGI/SecretToCGGI.h"
14+
#include "lib/Dialect/Secret/Transforms/AddDebugPort.h"
1515
#include "lib/Dialect/Secret/Transforms/DistributeGeneric.h"
16-
#include "lib/Pipelines/PipelineRegistration.h"
1716
#include "lib/Transforms/BooleanVectorizer/BooleanVectorizer.h"
1817
#include "lib/Transforms/FoldConstantTensors/FoldConstantTensors.h"
1918
#include "lib/Transforms/ForwardInsertToExtract/ForwardInsertToExtract.h"
@@ -23,7 +22,6 @@
2322
#include "lib/Transforms/MemrefToArith/MemrefToArith.h"
2423
#include "lib/Transforms/Secretize/Passes.h"
2524
#include "lib/Transforms/TensorLinalgToAffineLoops/TensorLinalgToAffineLoops.h"
26-
#include "lib/Transforms/UnusedMemRef/UnusedMemRef.h"
2725
#include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project
2826
#include "mlir/include/mlir/Conversion/TensorToLinalg/TensorToLinalgPass.h" // from @llvm-project
2927
#include "mlir/include/mlir/Dialect/Affine/Transforms/Passes.h" // from @llvm-project
@@ -64,6 +62,8 @@ void mlirToCGGIPipeline(OpPassManager& pm,
6462
const std::string& yosysFilesPath,
6563
const std::string& abcPath) {
6664
pm.addPass(debug::createDebugValidateNames());
65+
pm.addPass(secret::createSecretAddDebugPort(secret::SecretAddDebugPortOptions{
66+
.insertDebugAfterEveryOp = options.debug}));
6767
pm.addPass(createConvertTensorToLinalgPass());
6868
pm.addPass(createLinalgGeneralizeNamedOpsPass());
6969

@@ -154,6 +154,8 @@ CGGIPipelineBuilder mlirToCGGIPipelineBuilder() {
154154
void mlirToCGGIPipeline(OpPassManager& pm,
155155
const MLIRToCGGIPipelineOptions& options) {
156156
pm.addPass(debug::createDebugValidateNames());
157+
pm.addPass(secret::createSecretAddDebugPort(secret::SecretAddDebugPortOptions{
158+
.insertDebugAfterEveryOp = options.debug}));
157159
// Bufferize
158160
::mlir::heir::oneShotBufferize(pm);
159161

‎lib/Pipelines/BooleanPipelineRegistration.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ enum DataType { Bool, Integer };
2020
#ifndef HEIR_NO_YOSYS
2121
// If Yosys is enabled, also add all yosys optimizer pipeline options.
2222
struct MLIRToCGGIPipelineOptions : public YosysOptimizerPipelineOptions {
23+
PassOptions::Option<bool> debug{
24+
*this, "debug",
25+
llvm::cl::desc("Insert debug ports after every secret operation."),
26+
llvm::cl::init(false)};
2327
PassOptions::Option<enum DataType> dataType{
2428
*this, "data-type",
2529
llvm::cl::desc("Data type to use for arithmetization, yosys must be "
@@ -44,6 +48,10 @@ void mlirToCGGIPipeline(OpPassManager& pm,
4448
#else
4549
struct MLIRToCGGIPipelineOptions
4650
: public PassPipelineOptions<MLIRToCGGIPipelineOptions> {
51+
PassOptions::Option<bool> debug{
52+
*this, "debug",
53+
llvm::cl::desc("Insert debug ports after every secret operation."),
54+
llvm::cl::init(false)};
4755
PassOptions::Option<enum DataType> dataType{
4856
*this, "data-type",
4957
llvm::cl::desc("Data type to use for arithmetization."),

0 commit comments

Comments
 (0)