Rotom: Roll arguments can name a whole tensor axis - #3174
Conversation
|
@edwjchen Could you rebase this one? |
| extents need not match: the shift reduces modulo the rolled dim's extent, so | ||
| a smaller partner covers a prefix of the rotations and a larger one wraps. | ||
| Optional **rolls** encode `roll(from, by)` metadata objects, applied left | ||
| to right. Each endpoint is either a *piece* -- a position in the `dims` |
There was a problem hiding this comment.
"endpoint" is a bit of a weird term to use here, because the arguments to roll are not entities with a geometry (interval, segment) that has endpoints.
I see RollSpec above for the pair of from and by, so maybe call RollEndpoint instead RollArg and you can refer to it in prose as "the by argument" or "the from roll arg"?
| int64_t value; | ||
| if (parser.parseInteger(value)) return failure(); | ||
| rolls.push_back(value); | ||
| if (failed(parseRollEndpoint(parser, from)) || parser.parseComma() || |
There was a problem hiding this comment.
nit: I sort of missed this last time, but I think it would be better to just have one option for parsing and printing, and leaving the parens explicit seems better than leaving them out.
| for (Attribute a : rotomDims) { | ||
| auto d = cast<DimAttr>(a); |
There was a problem hiding this comment.
| for (Attribute a : rotomDims) { | |
| auto d = cast<DimAttr>(a); | |
| for (auto d : rotomDims.getAsRange<DimAttr>()) { |
| /// MULTIPLE of a digit, deliberately not layout vocabulary (layout rolls | ||
| /// shift by exactly the partner index). The plan that wants the packing | ||
| /// carries it; layouts stay unit-step and alignable endpoint-for-endpoint. | ||
| struct PreRotation { |
There was a problem hiding this comment.
I'm a little bit confused by this.
Is this essentially a workaround for a limitation of the rotom layout descriptor?
Is the rotom-produced plan accounting for this, i.e., is the operation that comes before this BSGS producing output that pre-computes the non-unit-step roll, so that the layout is already eagerly in this form? And the issue is just that we can't describe that with the rotom layout descriptor as-is?
What concerns me slightly: this deficiency would presumably propagate, so you need to track pre-rotations throughout the IR, and then the layout descriptor attribute needs to include one or more of these pre-rotations as well. But I don't really see where this is going: this PR doesn't update the attribute to have this, and it's not clear who would be calling this new method overload lowerToTensorExtIsl with a preRotation.
There was a problem hiding this comment.
The original problem was regarding how to express the BSGS shift in the layout representation. The goal was to generalize this optimization in the layout representation for easier cost analysis. My first attempt was to put it in the layout representation as a per-roll step. However, that approach conflated the stride attribute in the layout representation, making it harder to understand. The follow-up fix was to introduce this PreRotation, keeping the layout representation clean and inject the BSGS shift into the ISL relation when lowering.
Upon further thought, BSGS really shouldn't be a part of the layout representation. The layout representation should strictly describe how tensor elements are laid out. Instead, BSGS should be incorporated as optimization pass(es) when constructing the kernel. resulting in lowering the BSGS shifts as rotations on the plaintext weights. Further optimization passes could then hoist these rotations into the plaintext packings.
tldr: The updated changes drop this PreRotation, and BSGS will be incorporated in later Rotom PRs.
edwjchen
changed the title
edwjchen
changed the title
0bb5a0c to
187f110
Compare
95116c9 to
886dccf
Compare
A roll argument is now either a piece -- a dims-list position, the
original Rotom reading -- or a whole tensor axis, spelled 'axis N' and
stored as -(axis+1) in the flat rolls array. An axis argument is legal
only when the axis is packed as more than one piece; the piece spelling
is canonical for an unsplit axis, where the two coincide.
A piece FROM rewrites only its own mixed-radix digit -- the original
per-piece semantics, now materialized correctly on split axes (no borrow
crosses digits). An axis FROM rewrites the whole axis index modulo its
full extent, each piece then taking its digit of the rolled index: the
borrow across digits is what diagonal packings over a split axis need
and no combination of piece rolls can express. A BY piece of a split
axis shifts by that piece's digit of the axis's current (possibly
already-rolled) expression.
Rolls stay a pure packing description: a roll shifts by exactly its BY
argument's index, and a layout describes its value's packed bytes in
full. Kernel schedules that shift by a MULTIPLE of a digit -- the
baby-step/giant-step giant shift -- are not layout vocabulary and are
not folded into any value's packing; the kernel emits them as rotations
of its coefficient operand, which for plaintext weights a backend folds
into the encoded constants.
The BSGS diagonal packing is then an ordinary layout:
#rotom.layout<n = 16, rolls = [(axis 1, 2)],
dims = [[1:4:4], [1:4:1] | [0:16:1]]>
where the roll diagonalizes the whole split k against i (ciphertext
(g, b) holds the digits of (k - i) mod 16).
2 participants
Part 3 of #2980 and stacked on #3170
A Rotom roll decorates a layout:
roll(from, by)rewrites thefromindex by subtracting thebyargument's index. Until now a roll argument could only be a piece, a position in thedimslist. This PR lets it also be a whole tensor axis, spelledaxis Nand stored as-(axis+1)in the flat rolls array.