Lowering for comparisons in arithmetic pipeline

[AlexanderViand]

This is work towards fixing #1929 (for loop with secret bound) which primarily hinges on the ability to actually lower a comparison to a polynomial approximation.

So far, this (draft) PR has been work to get to the point where we can rewrite the code from #1929 all the way down to just having math_ext.sign that needs to be arithmetized/lowered to a polynomial.

On the way there, I ended up touching a bunch of places, though:

• Various fixes to --convert-secret-to-static-for to make it more robust/handle more edge cases
• Made --add-client-interface capable of handling multiple func.func in the same module
  (purely to avoid needing to split into multiple teot files, not directly related to the comparison issue)
• added IndexType to the type conversions in --convert-to-ciphertext-semantics
  (this will appear to "fix" Unexpected segfault after compilation error in --convert-to-ciphertext-semantics  #1992 but really that issue isn't about the inability to lower arith.index_cast but about the compiler segfaulting after encountering an op it cannot lower)
• Adds a bit of an UGLY HACK to --layout-propagation:
  The pass previously assumed that index type'd operands don't need to have a layout assigned (i.e., index valued scalars don't need to be converted to ciphertext-sized tensors) which is true for most index type values (e.g., the offset in a rotation, the indices in a tensor.extract/insert, etc). However, in the case of index values appearing in an arith.cmpi, we do need to assign them a layout (and they do need to be turned into ciphertext-sized vectors) because we will be doing math to them as part of the lowering. There was already a comment to suggest defining an op interface that can be used to decide whether a given op operand needs to have a layout assigned, and I think this would indeed be the correct solution. My current ugly hack simply hard-codes an exception for arith.cmpi into the index type skipping code.
• Adds a math_ext dialect with a single op math_ext.sign, with plans to move this upstream to math
• Adds a new pass --comparison-to-sign-rewrite that rewrites arith.cmpi/arith.cmpf:
  Currently, the pass only supports a < b or b > a, no equality. It also assumes that you can multiply with 0.5, so really only makes sense in a CKKS pipeline, though I've currently just added it to --mlir-to-secret-arithmetic. For BGV/BFV, comparison usually involves approximating mod t rather than sign....

The next step would be to try and actually enable polynomial approximation for math_ext.sign and see if the resulting programs to anything useful :)

EDIT: This now works for the full --mlir-to-ckks pipeline 🎉 but of course I had to make a few more changes along the way:

• added --polynomial-approximation to the arithmetic pipeline and added math_ext.sign lowering
• added lowerings (to lwe.reinterpret_application_data, as with other conversion) for arith.sitofp/uitofp/fptosi/tptoui (int/float conversions) to --secret-to-ckks
• added --lower-polynomial-eval to the arithmetic pipeline and had to make a few changes there:
  I had to change the degreeThreshold interpretation to be "inclusive" (I guess, alternatively, I could have bumped up the max from 5 to 6?) because the previous pass would generate degree 5 polynomials and the threshold was 5 here.
  I also had to add a convertFloatToSemantics function (mostly AI generated, and 100% AI named 😉) to make sure the pass would generate Attributes compatible with f32 values if that's what the original program was using, as float attributes apparently get interpreted as f64 if parsed from text?
  • EDIT2: It also needed a tiny fix in the OpenFHE emitter, since the scf.if -> arith.select -> mul/add/etc rewrites would introduce an i1 type which OpenFHE would realize as std::vector<int64_t> but MakePackedCKKSPlaintext only accepts double, so we emit one more vector copy/conversion in this case.

Final task: run the generated openfhe code and see if it does anything even remotely related to the original program 😅

EDIT2: If you rebase this on top of the mlir_src python frontend branch, you can run the example/test directly from python 🎉 ....but it'll be horrendously slow since I guess it'll run OpenFHE single threaded? Also, the example includes several bootstrap calls, so it'll be slow no matter what...

[AlexanderViand]

Ok, so after re-enabling OpenMP support for OpenFHE and allowing the python frontend to actually use OpenFHE in multi-threaded mode, it was at least somewhat more feasible to try and run the generated program. Unfortunately, it'll eventually fail with an OpenFHE Error (see below). Nevertheless, I think this is a good point to start getting some reviews and try and to get this PR merged, since it's already doing a lot of things, and I assume the fix for whatever noise/scale/etc management thing is going wrong here will be its own issue and PR.

RuntimeError: external/openfhe/src/pke/lib/scheme/ckksrns/ckksrns-fhe.cpp:l.489:EvalBootstrap(): Degree [10] must be less than or equal to the correction factor [7].

EDIT: I think I'll actually pull the OpenMP thing out of this PR, as it seems like the least related and easiest to extract + I want to make sure it works on a variety of machines before we merge it and accidentally break the OpenFHE backend

[j2kun]

Also note I'm working on upstreaming the sign ops in j2kun/llvm-project#8

[kragall]

I've just noticed that there's another open issue regarding lowering of arith.cmpi #827. Would there be a benefit to having this as an alternative approach to the --comparison-to-sign-rewrite for BGV/BFV?

[AlexanderViand]

I've just noticed that there's another open issue regarding lowering of arith.cmpi #827. Would there be a benefit to having this as an alternative approach to the --comparison-to-sign-rewrite for BGV/BFV?

Right, now that we have a LattiGo backend, the interpolation based approaches might indeed be viable! I'd say it's a lower priority right now but I'd be happy to offer my assistance if someone wants to try :)

[kragall]

The comparisons have been bugging me for a while now. I'll have a look, but I can't commit much time for them right now. Why does the LattiGo backend make that approach more viable?

Meanwhile, I managed to crash the --compare-to-sign-rewrite pass with this example:

  func.func @scr01(%arg0: f32 {secret.secret}) -> f32 {
    %cst = arith.constant 0.000000e+00 : f32
    %cst_0 = arith.constant 1.000000e+00 : f32
    %cst_1 = arith.constant 3.000000e+00 : f32
    %cst_2 = arith.constant 1.000000e+01 : f32
    %0 = arith.cmpf ugt, %arg0, %cst_1 : f32
    %1 = arith.cmpf ult, %arg0, %cst_2 : f32
    %2 = arith.select %1, %cst_0, %cst : f32
    %3 = arith.select %0, %2, %cst : f32
    return %3 : f32
  }

[AlexanderViand]

The comparisons have been bugging me for a while now. I'll have a look, but I can't commit much time for them right now. Why does the LattiGo backend make that approach more viable?

For this, you need the plaintext modulus $p$ to be extremely small since you need to evaluate a depth-$p$ polynomial to interpolate, and OpenFHE doesn't let you set the plaintext sufficiently small, but iirc LattiGo does.

Meanwhile, I managed to crash the --compare-to-sign-rewrite pass with this example:

Thanks, good catch! I wasn't taking into account different float types when creating the 1.0 and 0.5, should be fixed now!

[asraa]

Thanks for adding the follow-up TODOs!

[asraa]

should the result just be SignlessIntegerLike? I would expect that the result shouldn't be a float

edit: after finishing my review i realize it's floatlike too because the result types probably need to be float for CKKS