Rotom: plaintext pipeline and the MNIST runner test - #3382
Open
edwjchen wants to merge 7 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.
Walks a function bottom up, generating layout candidates for every value and folding them into one function-wide assignment. Elementwise ops take the reference's gen_binop; matmul enumerates the replicate, sum-roll and alignment variants, prices each chain, and rewrites the op into explicit conversion and roll ops plus a rotom.matmul that names every layout of the plan. A public operand is packed at encode time, so it takes the layout its consumer wants (the reference's match_public_kernel): an elementwise op may pack a public side at the compute layout, or exactly like its partner -- the only way to compute at a layout with gaps, which every matmul result has. Each use packs its own copy, since CSE merges equal literals. A matmul operand whose chain ends in a roll by the ciphertext digit, facing a public operand, keeps its pre-roll placement; the op records bsgs_baby and bsgs_source and the lowering folds the roll into a baby-step/giant-step schedule (the reference's BSGS_ROT_ROLL + BSGS_MATMUL). An apply_roll keeps the piece multiset; a step that also moves pieces is a conversion. A conversion the planner cannot express is refused in the search, not priced. The iteration space is read off the operand ranks rather than fixed dim ids: each operand axis takes its id in the result [..batch, i, j], and the summed axis takes the K dim, so a batched contraction maps the way a rank-2 one does. A compute placement is the only layout that carries a K piece; the result holds a plain gap where the sum collapsed it.
Lowers rotom.convert_layout, rotom.apply_roll and rotom.matmul to explicit ciphertext operations. A conversion replays the planner's steps -- one rotated row per (source ciphertext, shift), masked and accumulated per target -- then doubles each replication fill in with log2(extent) rotate-and-adds. A matmul multiplies once at the compute placement and reduces the contraction axis: slot digits by a log-tree rotate-and-reduce, ciphertext digits by block adds. A matmul with bsgs_source folds the pending ciphertext-digit roll into the multiply: B - 1 baby rotations shared across the giant groups, the public rows pre-rotated by the giant offset at encode time, G - 1 giant rotations on the partial sums (the reference's _bsgs_arbitrary_offsets). The pending operand is found by re-planning the roll from bsgs_source, not by row count.
rotom-outline-kernels outlines each assigned kernel region into its own function so the ciphertext lowering rewrites one call site per kernel; tensor_ext gets a dialect inliner interface so the outlined bodies inline back after lowering. Kernel operands are rewired by position, so an op whose operand positions alias one value (x * x) keeps a body faithful to its signature. Registers the mlir-to-rotom-ciphertext pipeline in heir-opt: one ciphertext-size option drives rotom-normalize-contractions -> rotom-seed-layout -> rotom-assign-layout -> rotom-outline-kernels -> rotom-materialize-tensor-ext-layout -> convert-to-ciphertext-semantics -> inline -> implement-shift-network, so high-level tensor IR in secret.generic regions lowers to ciphertext-semantic tensors with no hand-written layouts.
The passes the torch-exported MNIST model needs on top of the matmul and elementwise kernels: - rotom-normalize-contractions rewrites linalg.matvec and vecmat into linalg.matmul with unit dims, so every contraction reaches the one matmul generator; the pipeline runs it first. - Rank-changing reshapes name the unit axes they create (addUnitAxisPieces), so relations materialized from two layouts of the same type always compose. - A layouted public value produced by cleartext compute -- a bias that flows through host arithmetic before entering the secret region -- is an encode-time packing boundary: the materializer emits an explicit tensor_ext.assign_layout so the ciphertext-semantics conversion has a packing to lower. - An elementwise candidate must have an emission path (the Rotom kernel, or all three layouts coinciding); one that needs an alignment the kernel cannot lower is not offered, so selection walks back to producer candidates instead. - Outlined kernel operands are rewired by position, so aliased operands (x * x) keep a faithful body.
mlir-to-rotom-plaintext runs the CKKS front end, the Rotom layout passes, and the ciphertext-semantics lowering, then executes the result in cleartext through the LLVM backend. The runner test compiles the MNIST model of tests/Examples/common/mnist at n = 32768, packs the public layers with the generated helpers, encrypts the image, and checks all ten outputs against a cleartext reference of the same polynomial network. The circuit it executes is Rotom's plan on both layers: 27 ciphertext rotations.
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.
Adds the pipeline that makes the Rotom stack executable, and the numeric tests that check it computes the right answers.
mlir-to-rotom-plaintextruns the CKKS front end, the Rotom layout passes, and the ciphertext-semantics lowering, then executes the result in cleartext through the LLVM backend. Where mlir-to-rotom-ciphertext stops at ciphertext-semantic tensors, this one produces something that runs so the circuit Rotom planned can be checked against real numbers.Stacked on #3381.