Rotom: layout alignment and conversion planning - #3375
Open
edwjchen wants to merge 2 commits into
Open
Conversation
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 the part of the axis index that piece reads
-- the original per-piece semantics, now materialized correctly on split
axes, where nothing carries from one piece into the next. An axis FROM
rewrites the whole axis index modulo its full extent, each piece then
taking its part of the rolled index: that carry across pieces is what
diagonal packings over a split axis need, and no combination of piece
rolls can express it. A BY piece of a split axis shifts by that piece's
part 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 piece's offset --
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 parts of (k - i) mod 16).
Adds the two utilities the layout search is built on. The alignment engine brings a pair of layouts onto a shared placement. It replicates and rolls a side where the operator's map demands it, states the matmul alignment map in general form rather than special-casing it, and returns both sides restated at one granularity. The alignment check walks runs with two cursors, so [R:64] may face [R:4][i:16]. The conversion planner matches layouts axis by axis rather than piece by piece, so a conversion that re-splits an axis while adding a roll -- the repack every diagonal matmul needs -- plans instead of failing. A slot replication the target adds is a fill (log2(E) doublings), not one address step per copy; a replication the source already holds is neither a fill nor a coordinate, only a span of the step, since a rotated replicated row is correct in every block. One description both prices a conversion and emits it. estimateConversionCost counts the plan the lowering would emit -- its steps and its fills -- and reports a conversion the planner cannot express as unlowerable, so the search never chooses a step the lowering drops on the floor. The price and the plan disagreed by 15x before, and the silent drops hid it. The matmul result no longer keeps the summation dim's ciphertext piece as gaps: the reduction sums those ciphertexts into one, so the result is one ciphertext, as the reference's BSGS_MATMUL yields. bsgsScheduleOpt reads the baby-step/giant-step schedule off a layout pair, for the matmul fold. A public operand may repack at the layout its partner needs. Its data is known at encode time, so the packing costs nothing at run time. Vocabulary: a layout's parts are pieces, a piece's place value within its region is its offset, and the rewrite a roll performs is a roll -- not a shear. A roll names its arguments by piece position, so two layouts carry the same rolls when their roll lists match position for position.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adding Rotom's layout alignment engine, which brings a pair of layouts onto a shared placement, and the conversion planner, which turns a layout pair into the steps that move the data.
estimateConversionCostprices the lowering plan.Stacked on #3174.