Design Proposal for ModArith centered representatives

[crockeea]

Basis conversion should normally (always?) use centered representatives in (-q/2, q/2). However, it currently relies on ModArith:ExtractOp, which yields the standard representative in [0,q) (assuming it was normalized with ModArith:ReduceOp first?). I'm proposing the following design to resolve this situtation:

I propose a new ModArith:LiftOp which is similar to ExtractOp, but parameterized by a new ModArith Enum with values centered or standard. There would be a transform that lowers LiftOp into other ModArith and Arith ops. If the parameter is standard, LiftOp lowers to ReduceOp followed by ExtractOp. If the parameter is centered, it lowers to ReduceOp, ExtractOp, and a conditional subtraction by q.

[AlexanderViand]

I'd be open to changing HEIR's assumptions/default behaviour here. I don't think there's any reason for us to default to [0,q) so we could just make (-q/2,q/2) the standard?

[crockeea]

It's not clear to me if the ModArithToArith lowering in particular really supports negative numbers. E.g., when I was passing data into HEIR, I had to pass the inputs in [0,q) because the LLVM code ended up with a urem which didn't work for negative inputs. Care needed at a minimum.

Re using centered always: there are a few places in my software FHE library where I specifically request standard representatives (in [0,q)). I haven't reviewed these carefully to make sure they're strictly necessary, and they are in some corner cases relating, e.g., to the NTT implementation, but I still like having the option to do either.

[j2kun]

The main consideration should be: does this provide what you need to lower to LLVM and test properly? From that perspective, I think this proposal is good.

In some distant future (where there are multiple different modular reduction strategies) I can imagine we will want a more generalized range analysis that can statically identify range bounds on the underlying representative, and/or add constraints on the needed representative ranges at different points in the program that a solver could then optimize for. But that is far more work (and of speculative benefit) than what is needed for the immediate task.

[crockeea]

I guess the question is whether it's worth keeping ExtractOp and adding a LiftOp, or if I should just make ExtractOp take the enum. As it is, ExtractOp is "lightweight" in that it doesn't do a reduction first, it just outputs whatever the current representative is. If we made ExtractOp always output a canonical representative (either standard or centered) then it potentially introduces additional unnecessary ReduceOps that could be removed by some hypothetical static analysis pass.

[j2kun]

I think I added extract mainly as a structural bridge between arith and mod_arith, without thinking too hard about it. I like the idea of having any op that is more intentional about choosing the representative, so I think renaming extract to lift and handling both cases would be a good design.

[crockeea]

Okay, that simplifies the design a bit then, because I won't have a dedicated transform that decomposes Lift into Extract. I'll start on that.

[crockeea]

Resolved by #3066