secret: add-debug-port/import-result for secret-arithmetic op

[ZenithalHourlyRate]

Should be rebased after #1513 is in.

The idea behind this PR is that, with plaintext backend, we can compare the plaintext backend execution result with the CKKS execution result to know the precision loss.

Knowing the precision loss is the first step towards implementing/validating CKKS noise model.

However, as these two pipelines are distinct, and we have two executable (plaintext executable and CKKS executable), we can not run them in one-line to get the result. Lets first see the design

Design

There are two passes added --secret-add-debug-port and --secret-import-execution-result, the first is similar to --lwe-add-debug-port, while the second one is responsible for importing the execution result of the executable back to IR. I will show why this design is needed.

For the plaintext backend, now it can has debug handler like __heir_debug_tensor_8xf32_ to print the intermediate value during execution.

Then such execution result is imported back to the IR via --secret-import-execution-result=file-name=trace.log, and the resulting IR is

    %0 = secret.generic ins(%arg0, %arg1 : !secret.secret<tensor<8xf32>>, !secret.secret<tensor<8xf32>>) attrs = {__argattrs = [{plaintext.result = [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.40000000596046448, 5.000000e-01, 0.60000002384185791, 0.69999998807907104, 0.80000001192092896]}, {plaintext.result = [0.20000000298023224, 0.30000001192092896, 0.40000000596046448, 5.000000e-01, 0.60000002384185791, 0.69999998807907104, 0.80000001192092896, 0.89999997615814208]}]} {
    ^body(%input0: tensor<8xf32>, %input1: tensor<8xf32>):
      %1 = arith.mulf %input0, %input1 {plaintext.result = [0.020000001415610313, 0.060000002384185791, 0.12000000476837158, 0.20000000298023224, 0.30000001192092896, 0.42000001668930054, 0.56000000238418579, 0.71999996900558472]} : tensor<8xf32>
      %2 = arith.addf %1, %cst {plaintext.result = [0.12000000476837158, 0.15999999642372131, 0.2199999988079071, 0.30000001192092896, 0.40000000596046448, 0.52000004053115845, 0.6600000262260437, 0.81999999284744263]} : tensor<8xf32>

Then such information can be used by the CKKS backend to determine precision.

      auto err = log2(std::abs(result[i] - plaintextResult[i]));
      maxError = std::max(maxError, err);
    }

    std::cout << "  Precision lost: 2^" << std::setprecision(3) << maxError
              << std::endl;

Why

A natural question is, why do we need to import the execution result in the first place. The reason is that further passes beyond secret-arithmetic will change the IR with additional operations (like mgmt op and tensor.extract lowering)

For example, in plaintext backend for the following IR you will only see 3 result

%1 = arith.muli %arg0, %arg1
%2 = arith.addi %1, %arg2
%3 = tensor.extract %2[%c0]

But when lowered, it will become

%1 = arith.muli
%2 = mgmt.relinearize
%3 = arith.addi
// tensor.extract lowering expressed in secret-arithmetic level
// actually it is lowered at scheme level
// oneHot pt
%cst = arith.constant [0, 0, 0, 1]
%4 = arith.mul %3.cst
%5 = tensor_ext.rotate
%7 = lwe.reinterpret_application_data

Then in lwe-add-debug-port there will be 7 sentenses. How do we know what line corresponds to what line when comparing? Only the compiler know. So compiler should annotate the plaintext result and any further transformation should inherit the attribute.

Current known issues

• secret-insert-mgmt wont inherit the plaintext.result attribute from previous op
• optimize-relinearization will not preserve the attribute in it as it just removes the relin op and re-insert it.
• The lowering for tensor.extract is tooo late, and we can not simulate it in plaintext backend.

About the plaintext backend

Currently, a.out in plaintext backend is built by lit in sandbox so it can not persist. I tried to build it by bazel (then we will get rid of the annoying cc issue), but then we do not have mlir_translate.bzl and llc.bzl macro to use. So the dumping of result in plaintext backend is ugly now.

Example

# run the plaintext backend to get trace
# this test intentially fail to dump the result
bazel test //tests/Examples/plaintext/dot_product_f_debug:dot_product_8f.mlir.test > dot_product_f_debug.log

# manually clean .log to make it only contains the stdout of a.out

bazel run //tools:heir-opt -- \
  $PWD/tests/Examples/plaintext/dot_product_f_debug/dot_product_8f.mlir \
  --mlir-to-ckks="ciphertext-degree=8 plaintext-execution-result-file-name=$PWD/tests/Examples/openfhe/ckks/dot_product_8f_debug/dot_product_f_debug.log" \
  --scheme-to-openfhe=insert-debug-handler-calls=true

# then openfhe-related compilation

The result would be

# bazel test --test_output=all //tests/Examples/openfhe/ckks/dot_product_8f_debug/...
Input
  Precision lost: 2^-47.9
Input
  Precision lost: 2^-47.3
openfhe.mul_no_relin
  Precision lost: 2^-25.4
openfhe.relin
openfhe.add_plain
  Precision lost: 2^-15.3
openfhe.rot
  Precision lost: 2^-15.3
openfhe.rot
  Precision lost: 2^-25.4
openfhe.add
  Precision lost: 2^-15.3
openfhe.add
  Precision lost: 2^-15.3
openfhe.rot
  Precision lost: 2^-15.3
openfhe.add
  Precision lost: 2^-15.3
openfhe.add
  Precision lost: 2^-15.3
openfhe.rot
  Precision lost: 2^-15.3
openfhe.add
  Precision lost: 2^-15.3
openfhe.add
  Precision lost: 2^-15.3
openfhe.rot
  Precision lost: 2^-15.3
openfhe.add
  Precision lost: 2^-15.3
openfhe.mod_reduce
openfhe.mul_plain
  Precision lost: 2^1.32
openfhe.rot
  Precision lost: 2^-15.3
lwe.reinterpret_application_data
openfhe.mod_reduce

You can notice that there is no comparison result for mgmt op like relin/reduce, and the comparison for tensor.extract is wrong.

[j2kun]

I suppose the symlinks are fine, but I suspect this may behave badly with bazel in the long run. If possible, it would be better to have the rule's srcs point directly to the file, and you may need to use https://bazel.build/reference/be/functions#exports_files to make the file visible from the BUILD file of the source files, or else wrap the files in a filegroup rule and use it as a data dependency.

If you can't get either to work relatively quickly, we can submit as-is and I can patch it later.

[ZenithalHourlyRate]

I'll leave it for another PR.

[j2kun]

secret-insert-mgmt wont inherit the plaintext.result attribute from previous op
optimize-relinearization will not preserve the attribute in it as it just removes the relin op and re-insert it.

For these I wonder if it makes sense to have a general "forward-propagate an attribute" pass that would do something like: given an attribute name, for each op that has that attribute attached to its results, propagate it forward to downstream user ops, unless the user op itself has that attribute.

Then you could call this pass to do cleanup whenever needed.

[ZenithalHourlyRate]

Comments addressed.

I am still a little concerned about the artifact tests/Examples/openfhe/ckks/dot_product_8f_debug/dot_product_8f_debug.log being added to the repo. Ideally this should be generated each time when we run the test as the secret-arithmetic IR might change in the long run. Still want proper bazel macro for mlir-translate and llc, is there any advice from LLVM bazel team?

EDIT:

For these I wonder if it makes sense to have a general "forward-propagate an attribute" pass that would do something like: given an attribute name, for each op that has that attribute attached to its results, propagate it forward to downstream user ops, unless the user op itself has that attribute.

I agree with such approach. Lets leave it for another PR when we need it.

[j2kun]

It is not hard to add a proper bazel rule for these. For one tool in isolation, it would be similar to tools/heir-translate.bzl, but where the binary is replaced with mlir-translate or llc (@llvm-project//mlir:mlir-translate and @llvm-project//llvm:llc) and the arguments are populated appropriately.

The slightly trickier part is that you want to run the generated binary, capture the stdout to a file, and then use that file as an input to another rule. The easiest way to do this is probably to wrap the execution of the generated binary in a genrule, though you could also make a custom rule. The main reason for needing a rule is that any time a file is produced that needs to be passed between rules, the file itself must be declared as an output by one rule and as an input by another rule.

Then you could chain all these rules together in a bazel macro. I'm fine if we merge this as-is and do this in a followup. If you can't figure it out, I will probably have time to do it in early April.