Re-enable OpenFHE bootstrap placement for CKKS

[j2kun]

No description provided.

[j2kun]

bazel run //tools:heir-opt -- \
  '--annotate-module=backend=openfhe scheme=ckks' \
  '--torch-linalg-to-ckks=ciphertext-degree=1024 ckks-bootstrap-waterline=8' \
  '--scheme-to-openfhe' $PWD/tests/Examples/common/lenet.mlir

[j2kun]

heir-opt: lib/Parameters/RLWESecurityParams.cpp:25: int mlir::heir::computeRingDim(int, int): Assertion `false && "Failed to find ring dimension, logT
otalPQ too large"' failed. 

 #0 0x000055e6a9d2958d llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:842:
11                                                                                                                                                    
 #1 0x000055e6a9d29abb PrintStackTraceSignalHandler(void*) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:924:1               
 #2 0x000055e6a9d26fc4 llvm::sys::RunSignalHandlers() /proc/self/cwd/external/llvm-project/llvm/lib/Support/Signals.cpp:108:5                         
 #3 0x000055e6a9d2a0a9 SignalHandler(int, siginfo_t*, void*) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:429:38            
 #4 0x00007fe6cd849df0 (/lib/x86_64-linux-gnu/libc.so.6+0x3fdf0)                                                                                      
 #5 0x00007fe6cd89e95c __pthread_kill_implementation ./nptl/pthread_kill.c:44:76                                                                      
 #6 0x00007fe6cd849cc2 raise ./signal/../sysdeps/posix/raise.c:27:6                                                                                   
 #7 0x00007fe6cd8324ac abort ./stdlib/abort.c:81:3                                                                                                    
 #8 0x00007fe6cd832420 __assert_perror_fail ./assert/assert-perr.c:31:1                                                                               
 #9 0x000055e6a4b831b9 (/usr/local/google/home/jkun/.cache/bazel/_bazel_jkun/886f3804b421b5442165b2f8eb57dad5/execroot/_main/bazel-out/k8-dbg/bin/tool
s/heir-opt+0x5f611b9)                                                                                                                                 
#10 0x000055e6a4af4bd6 mlir::heir::ckks::SchemeParam::getConcreteSchemeParam(int, int, int, int, bool, bool, bool) /proc/self/cwd/lib/Parameters/CKKS/
Params.cpp:187:8                                                                                                                                      
#11 0x000055e6a33ae545 mlir::heir::GenerateParamCKKS::runOnOperation() /proc/self/cwd/lib/Transforms/GenerateParam/GenerateParamCKKS.cpp:129:5

[j2kun]

I also noticed that layout-propagation is quite slow here

[j2kun]

One weird thing is that when I set the bootstrap waterline to something silly (e.g., 8), then the IR for lenet.mlir generates some 74 levels, and bootstraps appropriately at level 66, but then only ever gets down to level 37 before the function completes. So we have 37 completely unused limbs.

[j2kun]

The issue is that this snippet of code sets a baseLevel of 37 for the subsequent annotate-mgmt pass:

heir/lib/Transforms/OptimizeRelinearization/OptimizeRelinearization.cpp

Lines 88 to 109 in 4200115

	// temporary workaround for B/FV and all schemes of Openfhe
	auto baseLevel = 0;
	if (moduleIsBFV(getOperation()) || moduleIsOpenfhe(getOperation())) {
	// inherit mulDepth information from existing mgmt attr.
	mgmt::MgmtAttr mgmtAttr = nullptr;
	getOperation()->walk([&](secret::GenericOp op) {
	for (auto i = 0; i != op->getBlock()->getNumArguments(); ++i) {
	if ((mgmtAttr = dyn_cast<mgmt::MgmtAttr>(op.getOperandAttr(
	i, mgmt::MgmtDialect::kArgMgmtAttrName)))) {
	break;
	}
	}
	});
	if (!mgmtAttr) {
	getOperation()->emitError(
	"No mgmt attribute found in the module for B/FV");
	return signalPassFailure();
	}
	baseLevel = mgmtAttr.getLevel();
	}

Have to figure out what this is doing. I wonder if @ZenithalHourlyRate remembers?

[j2kun]

I'm now able to get the pipeline to complete, but then there are a few issues:

• Rot key generation blows through 200 GiB of RAM (likely because we have so many levels)
• Despite having inserted bootstraps, the mgmt pipeline produces the same total number of levels (43) as when we don't use bootstrapping. This smells like a bug.

[ZenithalHourlyRate]

The issue is that this snippet of code sets a baseLevel of 37 for the subsequent annotate-mgmt pass:

heir/lib/Transforms/OptimizeRelinearization/OptimizeRelinearization.cpp

Lines 88 to 109 in 4200115

	// temporary workaround for B/FV and all schemes of Openfhe
	auto baseLevel = 0;
	if (moduleIsBFV(getOperation()) || moduleIsOpenfhe(getOperation())) {
	// inherit mulDepth information from existing mgmt attr.
	mgmt::MgmtAttr mgmtAttr = nullptr;
	getOperation()->walk([&](secret::GenericOp op) {
	for (auto i = 0; i != op->getBlock()->getNumArguments(); ++i) {
	if ((mgmtAttr = dyn_cast<mgmt::MgmtAttr>(op.getOperandAttr(
	i, mgmt::MgmtDialect::kArgMgmtAttrName)))) {
	break;
	}
	}
	});
	if (!mgmtAttr) {
	getOperation()->emitError(
	"No mgmt attribute found in the module for B/FV");
	return signalPassFailure();
	}
	baseLevel = mgmtAttr.getLevel();
	}

Have to figure out what this is doing. I wonder if @ZenithalHourlyRate remembers?

This was used for annotate-module=backend=ckks to skip the entire mgmt pipeline. For this PR, the logic for BFV should be maintained while the logic for moduleIsOpenfhe can be removed.

[ZenithalHourlyRate]

To address the problem of bootstrap placement we first need to model bootstrap-depth, i.e. the number of levels consumed by bootstrapping itself. This problem has been long left since #1207.

The quick way to do so is ZenithalHourlyRate@47ef8ae . I have not polished it yet (and will not likely to get the time). The lenet.mlir example takes too long to run so I was not able to test it through (OK it fails, stilling having L58). The very_deep.mlir example seems to suggest the waterline does take effect as its maximal level is 28 (unlike 37 here). Maybe the re-running solver behavior I changed in BootstrapWaterline::matchAndRewrite changed the thing.

The things left are to match bootstrap-depth with a bunch of parameters in levelBudgetEncode/Decode in ConfigureCryptoContext and the match the key distribution in CKKS/Params.cpp (Does we always assume UNIFORM_TERNARY for CKKS? It is more common to use SPARSE_TERNARY/ENCAPSULATED for CKKS bootstrapping as uniform key is indeed unfriendly for bootstrapping.

[asraa]

The things left are to match bootstrap-depth with a bunch of parameters in levelBudgetEncode/Decode in ConfigureCryptoContext and the match the key distribution in CKKS/Params.cpp (Does we always assume UNIFORM_TERNARY for CKKS? It is more common to use SPARSE_TERNARY/ENCAPSULATED for CKKS bootstrapping as uniform key is indeed unfriendly for bootstrapping.

Thank you! I just wanted to add that when I was experimenting with parameters as well, here are my findings:

• SPARSE_TERNARY indeed yields a better # of levels consumed by bootstrap
• leve budgets have a size of keys needed for bootstrapping vs performance trade-off. the default is 3, 3. You can get the number of bootstrap levels consumed by the openfhe library API as a function of the secret key distribution and the level budgets
• one example parameter is setting encode/decode budget to 2, 2 with sparse keys and that yields a pretty small 13 levels consumed for bootstrapping
• You can set the numSlots in EvalBootstrapKeyGen - this is a HUGE performance uptick. Set that to the ciphertext-degree parameter. Otherwise the bootstrap key will be huge and the performance very bad