Add a polynomial attribute

[j2kun]

This PR adds a polynomial attribute definition to the poly dialect, which can be used to annotate types or ops with information about the polynomial ring in which the operation occurs.

An example attribute is added to poly.mul for demonstration, with a TODO left to remove it, and the poly_ops.mlir test updated to check for it. Update: I had time to add #poly.ring so I went ahead and did it.

Notable aspects of this PR:

• Polynomial is an internal representation of a polynomial as a list of Monomials, using the llvm::TrailingObjects storage paradigm. Polynomial is not meant to be used for polynomial arithmetic per se (i.e., we shouldn't be doing large computations with it). I had originally added Polynomial::operator+ expecting to do some simple operations with it, and then removed it when I realized it was unnecessary and we can just use fromMonomials and construct the monomials directly.
• The underlying coefficient/exponent types are all APInt with a 64-bit width, the idea being that if we need to increase it later we can do so easily. I also saw MPInt in the Presburger library, and I think that might be a better type to use for this because it doesn't assert that the bit width of the two operands to be the same. But since it's hidden inside Presburger I wasn't sure if it is sensible/stable to depend on. I asked @Groverkss in the mlir chat about this.

[AlexanderViand]

Looking forward to digging into this properly next week! I think the decision to go with 64 bits is reasonable for now, with a proper answer making a good agenda item for Tuesday's meeting :) Kunwar (@Groverkss) also kindly agreed to join for the meeting so we can quiz him directly.

I already chatted with a few people about what the "correct" answer to this question is. So far, no one had a good use case for the non-RNS (i.e., multiprecision) variants of the schemes, but maybe they are relevant for some more esoteric/advanced applications of FHE. Even for RNS variants, I guess there could be (HW) implementations that use RNS limbs with more than 64 bits? However, again, I'm unaware of any of the top of my head. If we can do MPInt without too much overhead in the "common" 64 bit case, it's probably the safer/more future proof choice.

PS: @j2kun Could you please request the review/give any necessary repo permission to my work account (@AlexanderViand) instead? Thanks!

[j2kun]

PS: @j2kun Could you please request the review/give any necessary repo permission to my work account (@AlexanderViand) instead? Thanks!

Sent the invite to @AlexanderViand-Intel and will add that account as a reviewer after you accept.

[asraa]

Thank you! This is great

[asraa]

not sure if there is a good answer here but what is the difference / purpose of the attribute uniquer versus any other type?

[j2kun]

So I don't know why the MLIRContext has multiple uniquers (they're all the same type), but I figure they might be for conflict avoidance of some sort? Like if you want separate instances of a type de-duped for attributes vs types vs affine maps? Not really sure.

[j2kun]

@AlexanderViand-Intel any opposition to submitting this PR? If you're short on time, we can always chat about it post submission. I think Asra wants use of the polynomial attribute for her bgv dialect.

[j2kun]

I'm going to make a judgement call and merge this one in. Happy to follow up afterward with any additional comments and improvements.

[AlexanderViand-Intel]

I think the polynomial attribute itself looks great, I'm not so sure about the !poly.poly type, though. I don't think I agree with having this be completely parameter-less with a ring attribute on the operation makes sense here. This MLIR/LLVM idiom of "operations specify type semantics" seems to be rooted in a pretty low-level view of computation, where a lot of data is basically just 32/64/word size bits and individual CPU instructions will in fact happily take these bits and re-interpret them.

In the FHE world, however, there's a massive difference between a polynomial from a degree 64k, 600-bit cmod ring and one from an n=2k, 32-bit ring, and "instructions" (both software and hardware implementations of, e.g., NTT) can't just ignore these differences.

[AlexanderViand-Intel]

with an eye towards upstreaming, it probably makes sense to remove the heir namespace in the middle?

[AlexanderViand-Intel]

I'm a bit surprised to see the poly type with no ODS parameters. Is that just WIP, or an intentional design decision?

[j2kun]

It was mostly to keep the PR small enough. I think now that I'm deep in the dialect conversion framework, I see that, at the very least, we need the size of the underlying storage type and the coefficient bit width available in the type in order to use the TypeConverter.

[AlexanderViand-Intel]

Mh, I'm not sure we want details of the TypeStorage to leak into the IR, rather the coefficient bit width should determine (in some implementation defined way) how the TypeStorage stores this. My motivation for this is to have a clear separation between "MLIR the concept" and "MLIR the C++ code base", making sure the things we standardized depend only on the former (of course, with a reference implementation using the latter).

In order to make the TypeConverter work, you could add a storageType() method to the (C++ implementation of) the type. Alternatively, the type could just figure this out internally in it's create/build/constructor, especially if we end up using something like MPint anyway?

[j2kun]

Sorry, when I said "underlying storage type" what I mean was some sort of dimension signifier. I.e., poly.poly<1024xi32> to signify the coefficient bit width and the "degree" (of the underlying ring) which corresponds to the size of the tensor one would lower to.

[AlexanderViand-Intel]

I see! I guess in an ideal world we'd want some kind of Z_q type for the second part (i.e., instead of i32)? Sort of a baseQ (taking inspiration from the base2 dialect proposed by @KFAFSP)

[j2kun]

I think the polynomial attribute in this PR would be sufficient to put on the type. And then we just mandate that each polynomial type is associated with a specific ring, and if we want to switch rings we have to add some special ops for that.

[KFAFSP]

For example, base2.fixed<signed 32,16> (alias base2.si32_16) is only usable with a compile-time fixed precision. Is this desired here, or do you need a type to model $\mathbb{Q}$ instead of $\mathbb{Z}_N \cdot 2^{-E}$?

[AlexanderViand-Intel]

I think the polynomial attribute in this PR would be sufficient to put on the type. And then we just mandate that each polynomial type is associated with a specific ring, and if we want to switch rings we have to add some special ops for that.

I'm not sure I'm following. In my understanding, an attribute is a compile-time value (usually closely associated with a type, with a dialect.constant operation that "turns" an attribute into a value of that type). Having a #poly.polynomial attribute as a type parameter for the !poly.polynomial type would be parameterizing the type by a specific polynomial? Sorry if I'm misunderstanding things here!

Ahh, 🤦 we could parameterize the !poly.polynomial type by #poly.ring which I guess is what you meant?

[AlexanderViand-Intel]

Sorry for the late response, I was at USENIX all week and forgot to submit a review beforehand, I submitted it now (see above).