CKKS: basis conversion and modulus switching

[bencemali]

While having a look at the CKKS key switch lowering, I noticed that the last step (Step 4: Remove the key-switch primes) incorrectly performs an LWE level basis conversion op from QP to Q instead of a modulus switch from QP to Q. These operations would correspond to Conv (basis conversion) and ModDown (modulus switching) in Appendix A of Faster Bootstrapping for CKKS with Less Modulus Consumption.

While investigating how ModDown could be implemented I also found the following:

• The LWE level basis conversion op gets lowered to the RNS level basis conversion.
• The RNS level basis conversion op then gets lowered to further RNS and ModArith level operations here with Garner's algorithm.

I propose the following fixes:

• Remove the LWE level basis conversion op. The direct translation from LWE basis conversion to RNS basis conversion is an indication that basis conversion should live in the RNS dialect and not the LWE one.
• Define a ModUp (basis extension from Q to QP) and a ModDown (modulus switching from QP to Q) operation on the LWE dialect following the standard naming convention from Section 2.2 of HK19. These ops will use the RNS level base conversion when lowered.
• Replace the usages (basis extension and modulus switching) of the current LWE level basis conversion with the newly added ModUp and ModDown ops.
• Reimplement the RNS level basis conversion lowering to use the fast base conversion algorithm (Section 2.2), which is the standard algorithm used in CKKS (hybrid) key switching implementations. The Garner's algorithm implementation could optionally be kept as an exact basis conversion algorithm.

If you think this is a good direction I can start putting together PRs for these. I am not confident about the compiler/MLIR side of things, so any suggestions are welcome.

tagging @AlexanderViand, @crockeea, @j2kun

[j2kun]

Hi @bencemali welcome to HEIR!

@crockeea is on vacation this week, so I'd ask we defer any hard decisions until next week so he can provide input to the claim of incorrectness.

I will give some basic answers in the meantime.

The direct translation from LWE basis conversion to RNS basis conversion is an indication that basis conversion should live in the RNS dialect and not the LWE one.

The translation is not direct. LWE basis conversion operates on ring elements, whereas RNS basis conversion operates on RNS-typed values (which is one coefficient of a polynomial). The lowering goes through polynomial.apply_coefficientwise. The reason we have the ops at both levels is because there are planned hardware backends that abstract over the RNS type but not over the RLWE ring elements.

Define a ModUp (basis extension from Q to QP) and a ModDown (modulus switching from QP to Q) operation on the LWE dialect

What is the reason to replace one general op with two more specific ops? In particular, there are concepts like rational rescaling and grafting which I believe convert_basis can naturally express but ModUp and ModDown cannot. It would be possible, however, to define mod_up and mod_down as convenience ops, and have them canonicalize to the appropriate convert_basis op.

Reimplement the RNS level basis conversion lowering to use the fast base conversion algorithm

I support having multiple implementations living side by side (via a pass option to select which method or a separate pass entirely). And certainly there is some degree to which the best method depends on the hardware. That said, I know @crockeea has thought long and hard about the best implementation here.

[crockeea]

I did forget to do a subtraction and modular mul for ModDown, approximately here. This should be something like

// extract both coeffs (existing code)
Value constTerm = lwe::ExtractCoeffOp::create(b, sum, b.getIndexAttr(0));
Value linearTerm = lwe::ExtractCoeffOp::create(b, sum, b.getIndexAttr(1));

// extract the mod q components
Value constTermQ =
      lwe::ExtractSliceOp::create(b, constTerm , ...));
  Value linearTermQ =
      lwe::ExtractSliceOp::create(b, linearTerm ,...);

// extract the mod p components
Value constTermP =
      lwe::ExtractSliceOp::create(b, constTerm, ...));
  Value linearTermP =
      lwe::ExtractSliceOp::create(b, linearTerm ,...);

// basis-convert mod-p to mod-q
Value modDownConstTerm =
      lwe::ConvertBasisOp::create(b, constTermP , TypeAttr::get(inputRNSType));
  Value modDownLinearTerm =
      lwe::ConvertBasisOp::create(b, linearTermP , TypeAttr::get(inputRNSType));

// subtract
Value diffConstTerm = lwe::RSubOp(b, constTermQ, modDownConstTerm, ...)
Value diffLinearTerm = lwe::RSubOp(b, linearTermQ, modDownLinearTerm, ...)

// scale by a constant, which is P^{-1} mod Q
Value newConstTerm = lwe::RMulConst(b, diffConstTerm, <CONST>, ...)
Value newLinearTerm = lwe::RMulConst(b, diffLinearTerm, <CONST>, ...)

// resume existing code
rewriter.replaceOp(op, {newConstTerm , newLinearTerm });
return success();

Note that there isn't an lwe::RMulConst op right now.

The fact that I made an error here in itself suggests that it ModDown could stand to be its own op. Other reasons for making it its own op would be if there are multiple lowerings for it, though if you're using FBC for ModDown, you'd probably use it for ModUp, too, which would mean there's really just two lowerings for KS overall. That said, what we really need is a test that lowers to LLVM with real inputs to ensure correctness, and then this error would have been caught in a test.

As for using FBC vs Garner; I opted for Garner to reduce error as much as possible, and because it's convenient (and faster) for some hardware platforms. I'm not opposed to adding FBC, and also support having multiple lowering options.

I recently realized that the key-switching we have defined is not unique to CKKS, and so it should live in the LWE dialect, and the current CKKS transform (at least DecomposeKeySwitch) should also live in LWE.

[j2kun]

@bencemali I think at this point it's up to you what you want to try implementing. Adding a FBC lowering seems like it wouldn't interfere with ongoing work from Eric. You could also try fixing this immediate error and working toward an e2e LLVM test. Or adding ModUp/ModDown with appropriate canonicalization patterns.

[bencemali]

@j2kun Sounds good, I can do these in two stages: figure out how to add ModUp/ModDown appropriately along with an 2e2 test for correctness, then FBC. But first I'll have to read up on MLIR to know what "appropriate canonicalization pattern" even means :)
I still find it weird that basis conversion should be in the same dialect as ct mult for example, then be lowered to the same operation coefficient-wise. I feel like basis conversion is an RNS coefficient-level operation, then ModUp and ModDown are the LWE scheme-agnostic modulus manipulation operations using bconv. ModUp and ModDown would still be meaningful operations even if you didn't use RNS representation, basis conversion in the narrow sense would not.

[crockeea]

I created #3017 to move KeySwitchInner from CKKS to LWE. I did not address the bug; I'll leave it up to @bencemali for now to decide how to resolve it (i.e., by creating ModUp/ModDown, or by creating lwe::RMulConst).

[crockeea]

Given lack of movement on this, I'll tackle the fix. Hopefully can get that out tomorrow.