Add external debug port during execution

[ZenithalHourlyRate]

Partially address #1266. This PR illustrate we can print decryption result after each step, for debugging. Also, I use such interface for evaluating noise magnititude after each step, so this is part of the param selection support.

This PR does the following things

• Implement lwe-add-debug-port pass that
  • Add secret key arg to main func
  • call __heir_debug after each HE op
• Modify lwe-to-openfhe to support these func declacation and func.call
• Add e2e demo

Example

For dot_product_8, the execution result is

$ ./bazel-out/k8-dbg/bin/tests/Examples/openfhe/dot_product_8_debug_test
Running main() from gmock_main.cc
[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from DotProduct8Test
[ RUN      ] DotProduct8Test.RunTest
( 2 6 12 20 30 42 56 72 ... )
( 2 6 12 20 30 42 56 72 ... )
( 30 42 56 72 2 6 12 20 ... )
( 32 48 68 92 32 48 68 92 ... )
( 68 92 32 48 68 92 32 48 ... )
( 100 140 100 140 100 140 100 140 ... )
( 140 100 140 100 140 100 140 100 ... )
( 240 240 240 240 240 240 240 240 ... )
( 240 240 240 240 240 240 240 240 ... )
( 0 0 0 0 0 0 0 240 ... )
( 240 ... )
( 240 ... )
( 240 ... )
[       OK ] DotProduct8Test.RunTest (4923 ms)
[----------] 1 test from DotProduct8Test (4923 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test suite ran. (4923 ms total)
[  PASSED  ] 1 test.

IR after lwe-add-debug-port

  // Notice that due to type issue we have multiple debug func
  func.func private @__heir_debug_4(!skey_L2_, !ct_L0_)
  func.func private @__heir_debug_3(!skey_L2_, !ct_L1_)
  func.func private @__heir_debug_2(!skey_L2_, !ct_L1_1)
  func.func private @__heir_debug_1(!skey_L2_, !ct_L2_)
  func.func private @__heir_debug_0(!skey_L2_, !ct_L2_D3_)
  func.func @dot_product(%sk: !skey_L2_, %ct: !ct_L2_, %ct_0: !ct_L2_) -> !ct_L0_ {
    %ct_1 = lwe.rmul %ct, %ct_0 : (!ct_L2_, !ct_L2_) -> !ct_L2_D3_
    call @__heir_debug_0(%sk, %ct_1) : (!skey_L2_, !ct_L2_D3_) -> ()
    %ct_2 = bgv.relinearize %ct_1 {from_basis = array<i32: 0, 1, 2>, to_basis = array<i32: 0, 1>} : !ct_L2_D3_ -> !ct_L2_
    call @__heir_debug_1(%sk, %ct_2) : (!skey_L2_, !ct_L2_) -> ()
...

IR after lwe-to-openfhe

  func.func private @__heir_debug_0(!openfhe.crypto_context, !openfhe.private_key, !ct_L2_D3_)
  func.func @dot_product(%cc: !openfhe.crypto_context, %sk: !openfhe.private_key, %ct: !ct_L2_, %ct_0: !ct_L2_) -> !ct_L0_ {
      %ct_1 = openfhe.mul_no_relin %cc, %ct, %ct_0 : (!openfhe.crypto_context, !ct_L2_, !ct_L2_) -> !ct_L2_D3_
    call @__heir_debug_0(%cc, %sk, %ct_1) : (!openfhe.crypto_context, !openfhe.private_key, !ct_L2_D3_) -> ()
    %ct_2 = openfhe.relin %cc, %ct_1 : (!openfhe.crypto_context, !ct_L2_D3_) -> !ct_L2_
    call @__heir_debug_1(%cc, %sk, %ct_2) : (!openfhe.crypto_context, !openfhe.private_key, !ct_L2_) -> ()

Then heir-translate would canonicalize all those debug call

void __heir_debug(CryptoContextT, PrivateKeyT, CiphertextT);
CiphertextT dot_product(CryptoContextT cc, PrivateKeyT sk, CiphertextT ct, CiphertextT ct1) {
  const auto& ct2 = cc->EvalMultNoRelin(ct, ct1);
  __heir_debug(cc, sk, ct2);
  const auto& ct3 = cc->Relinearize(ct2);
  __heir_debug(cc, sk, ct3);
...

In user file, user can specify what they want to debug, for example, printing decryption result

void __heir_debug(CryptoContextT cc, PrivateKeyT sk, CiphertextT ct) {
  PlaintextT ptxt;
  cc->Decrypt(sk, ct, &ptxt);
  ptxt->SetLength(8);
  std::cout << ptxt << std::endl;
}

Discussion

• Since .mlir does not have AnyType, we have to use __heir_debug_#number to address various LWE types. I could not think of better solution for now.
• Currently we can not tell which output corresponds to which variable, because we can not get variable name in mlir. This harms readablity. In my branch I hacked to get the following debug output, containing operation information and the "bound" attr from IR. I did so by passing all them as string arg to the debug port (at heir-translate stage), but in MLIR we cannot do such call.

ct=Encrypt pk, pt       cv 2 Ql 4 budget 123.592 noise: 42.4077 bound 26.5 gap -15.9077
ct=Encrypt pk, pt       cv 2 Ql 4 budget 123.493 noise: 42.5071 bound 26.5 gap -16.0071
ct2=EvalMultNoRelin ct, ct1     cv 3 Ql 3 budget 66.7256 noise: 83.2744 bound 26.5 gap -56.7744

[j2kun]

Currently we can not tell which output corresponds to which variable

One hack might be to pretty-print the MLIR op (op.print(stream)) as a string and emit that as an argument to printf in the OpenFHE emitter. This could be done generically just after the type switch in translate, with special cases to avoid printing all of func/module. Or it could be in its own function like

void __heir_debug(CryptoContextT cc, PrivateKeyT, sk, CiphertextT, ct);
void __heir_debug_with_print(CryptoContextT cc, PrivateKeyT, sk, CiphertextT, ct) {
  printf("<op_str>");
  __heir_debug(cc, sk, ct);
}

[j2kun]

So this is awesome work!

I have a few nitpicks about the code, but otherwise I'd be happy to submit this as-is and leave any improvements for followup work.

One idea that comes to mind: have a default implementation that decrypts and prints. In particular, the helper function is expected to be provided in C++, but people using the frontend won't be able to provide that. Long term we may want a user-provided debugger callback in python using the OpenFHE python bindings, but that seems beyond the scope of this PR. As a middle-ground, we could have a pass option for add-debug-port that specifies the (prefix of the) name of the debug function. If unset, then the compiler will emit a default implementation instead of just the external declaration.

For the numbered debug functions, I think this is an acceptable solution. Brainstorming on how to avoid this if we were forced to, while also keeping the definition at the lwe level, I think we would need to add something like an OpaqueCiphertext type to lwe, as well as an lwe.cast op that casts to and from these opaque types to regular ciphertext types. Then lowering to OpenFHE would trivially drop the casts and use OpenFHE's opaque type for everything. We'd have to be careful of how we use it, though, because casting from an opaque ciphertext type to a non-opaque type would be error-prone.

[AlexanderViand-Intel]

This is awesome, thanks!!!

[AlexanderViand-Intel]

As Jeremy already said, it'd be amazing if there was a "default" debug function that gets emitted, toggled by some pass option.

Looking at this one, it seems like the only thing that's somewhat application dependent is the length, but I'd be happy to have a debug that just dumps the entire ptxt for now. Once we have layout encodings, the compiler will be able to know how many cleartext elements are actually in there anyway, then we can consider revisiting this.

[ZenithalHourlyRate]

Looking at this one, it seems like the only thing that's somewhat application dependent is the length, but I'd be happy to have a debug that just dumps the entire ptxt for now.

If we delete the ptxt->SetLength(8);, the debug output would be of size 315K (lots of output) and is quite unreadable.

As a middle-ground, we could have a pass option for add-debug-port that specifies the (prefix of the) name of the debug function. If unset, then the compiler will emit a default implementation instead of just the external declaration.

I think I am able to generate a default impl at heir-translate stage (at lwe stage, we do not have the printf op abstraction), and also use lwe underlying type to infer the pt length. One subtle thing is that, for dot_product we have two underlying type (tensor<8xi16> and i16) and we need use the "correct" underlying type to make the debug func work.

[j2kun]

I think it would be fair to merge this PR and work on a default improvement in followup, so as to ensure we get the other useful parts of this PR in (e.g., emitting func.call and private func.func)

[ZenithalHourlyRate]

Comments addressed.

[ZenithalHourlyRate]

CI failure unrelated to this PR.

//tests/Dialect/CGGI/Transforms:boolean_vectorizer.mlir.test             FAILED in 3.2s

This failure can also be seen in https://github.com/google/heir/actions/runs/12955968820/job/36141258147

[ZenithalHourlyRate]

Dug around MLIR source and found funcOp.isDeclaration more suitable for detecting whether it has a body than funcOp.isPrivate, as we can have private funcOp with body, (i.e. static func in C term), see https://github.com/llvm/llvm-project/blob/de5ff8ad07ae824b86c5cefcba63f4b66607b759/mlir/lib/Conversion/FuncToEmitC/FuncToEmitC.cpp#L76-L81

[asraa]

@ZenithalHourlyRate - I had cherry-picked this PR and was using it to work on a demo, and then realized that this doesn't print out the type. I had to replace with emitAutoAssignPrefix(result);

[j2kun]

Merged as 9afa3e8