wip: Rotom layout alignment, assignment, and lowering - #2980
Draft
edwjchen wants to merge 86 commits into
Draft
Conversation
AlexanderViand
marked this pull request as draft
May 21, 2026 16:52
asraa
self-requested a review
May 21, 2026 16:52
Add withRolls() and enumerateSingleRolls() to LayoutAlignment. The latter ranges over ordered pairs of distinct, equal-extent slot-side traversal dims and returns the materializable single-roll variants of a layout, the building block for seeding rolled (diagonal) packings into the layout search. Rolls only affect the slot line of the materialized layout, so a roll whose 'from' dim is ciphertext-side lowers as a no-op; restrict both roll indices to the slot side so every variant is a genuinely distinct packing.
Two pipeline-execution tests proving rolled (diagonal) layouts compute correctly through layout assignment, materialization, the relation-driven lowering, and the interpreter: - a shared rolled layout on both addf operands, and - a rolled operand added to a row-major operand, which forces the lowering to remap between the rolled and unrolled packings. Validates the premise that the existing brute-force path already handles rolled layouts, de-risking rolled-variant injection into the search.
Add RotomTensorOpLowering::lowerMatmulByRotation, a diagonal matmul kernel that replaces the brute-force per-scalar mask+remap lowering with ciphertext rotations. For a single-ciphertext layout it groups every (i,j,k) contraction term by the rotation pair that realizes it -- (srcSlot - dstSlot) mod numSlots for each operand -- and emits one masked rotate-multiply per group, accumulated into the destination operand. The construction is correct for any such layout (each contribution is realized exactly once at its destination, with a shift-collision guard and single-point/single-ciphertext gates that fall back to the brute-force path), and collapses to few rotations exactly when the layout is rolled/diagonal: a Halevi-Shoup matvec with a diagonal-packed matrix uses 14 rotations vs 32 row-major. Also add tensor_ext.rotate support to the OpenFHE interpreter so the kernel is execution-verifiable, update the matmul lowering golden (remap -> rotate), and add an execution test asserting the rolled packing uses fewer rotations and still computes the correct matvec.
…NIST@32768
Implements the Halevi-Shoup ciphertext-axis diagonal matvec with a
baby-step/giant-step schedule, auto-discovered from a row-major seed, and
threads it through the whole seed -> assign -> materialize -> convert
pipeline so a full MNIST 784->512->10 MLP lowers at the reference
ciphertext size 32768 with the reference packing (16-ciphertext layer 1).
Kernels (RotomTensorOpLowering):
- lowerMatvecCtDiagonalBsgs: square/squat single-period + replicated BSGS.
- lowerMatvecDenseDiagonal: P=N/K diagonals per ciphertext (dense packing).
- Both handle non-power-of-two M and K by doing slot arithmetic over the
padded extents Dp=nextPow2(m), Kp=nextPow2(n) while verifying only the
real [0,m)x[0,n) domain (padding is zero); identical for power-of-two dims.
- Layout verification rebuilt the ISL relation per matrix element (O(m*n)
ISL builds, minutes at network scale); now indexes each relation once
(indexRangeByDomain) -> per-element lookups. Convert of a 512-hidden layer
drops from minutes to ~16s.
Search (LayoutAssignment):
- Generates the rolled diagonal/identity candidates and a rotation-aware
cost so the matvec kernel is discovered without a hand-written layout.
- Kernel gate checks the squat contract on padded extents.
- Fixes a back-propagation break where a bare operand candidate shadowed an
intermediate's candidate (chained matmuls losing their kernel).
Seed (SeedLayout): replication-to-fill -- a tensor smaller than the
ciphertext (e.g. a 10x512 weight at n=32768) now gets a seed (full slot
packing + a replication dim) so the search has a candidate to roll.
Secret path (MaterializeTensorExtLayout): copy a secret.generic operand's
materialized layout onto the corresponding function-argument attribute, so
the type converter converts the secret arg (fixes a block-arg/operand type
mismatch when a secret vector materializes wider than its data).
Utils/Layout: guard isRelation{Bicyclic,RowMajor,PerRow,SquatDiagonal}
against incompatible Presburger spaces (a replicated layout has extra local
vars) instead of asserting inside IntegerRelation::isEqual.
Pipeline: register --mlir-to-rotom-ciphertext (seed -> assign -> materialize
-> convert) so high-level tensor IR in secret.generic regions lowers in one
command with no hand-written layouts.
Tests: RotomPipelineExecutionTest gains end-to-end numerical coverage for
the square/squat/dense/replicated/under-filled/padded matvec regimes, a
two-layer MLP (matvec -> square -> matvec), and an MNIST-layer-scale
512x1024 -> 16x512 MLP; plus a lit test for the registered pipeline.
A neural-net layer y = x . W^T is a vector-times-matrix matmul x(1xK) * M(KxN),
but the ciphertext-axis diagonal matvec kernel only handles matrix * column-
vector. SeedLayout now rewrites such matmuls (gate: lhs.dim0==1 && rhs.dim1>1)
into transpose(matmul(M^T, x^T)) before seeding, reusing the existing kernel;
when M = transpose(W) the pre-transpose W is used directly to avoid a redundant
transpose-of-transpose. The matrix is then W (NxK) and must be squat (N<=K),
which the real MNIST layers (512x784, 10x512) satisfy.
Both matmuls of the real MNIST model now lower to RotomMatmul with no leftover
matmul. Adds a lit test for the rewrite.
Also guards two more IntegerRelation::isEqual call sites (ConvertTensor
{Collapse,Expand}Shape) against incompatible Presburger spaces -- exposed by the
rewrite's vector transposes lowering to reshapes over replicated layouts.
Seven fixes from the branch code review: - SeedLayout: only reuse a transpose producer's input when it is a true [1,0] permutation, else an identity transpose would build an invalid matmul(K x N, K x 1). - LayoutAssignment::visitMatmul: bail on dynamic shapes (matching the lowering) so the padded-extent cost model never divides by kDynamic. - MaterializeTensorExtLayout: error on a function argument feeding secret.generic operands with conflicting materialized layouts instead of silently keeping the first. - verifyLayoutRolls: reject a zero-extent rolled dim before the modulo (was a div-by-zero in the verifier). - lowerMatmulByRotation: index each relation once via indexRangeByDomain rather than rebuilding the IntegerRelation per (i,j,k) domain point. - lowerMatvecCtDiagonalBsgs: drop the dead 'n > Kp' gate (Kp = pow2(n)). - LayoutAssignment: use llvm::Log2_64_Ceil instead of a hand-rolled ceilLog2. heir-opt builds clean; all 28 Rotom tests pass.
Addresses two more code-review findings: google#4 layoutNumCiphertexts overcounted a dense straddling layout. Expose the straddle-aware inferCtPrefixLen from RotomAttributes (moved out of the anonymous namespace, declared in the header) and route the Rotom layout cost utilities through it, so a boundary dim that spans the ct/slot split contributes only its high (ciphertext) part. Adds a unit test pinning the 2-ciphertext count for a 4x8 straddling layout at n=16. google#11 the 'spaces compatible then isEqual' guard was duplicated at six sites. Factor it into relationsCompatibleAndEqual (lib/Utils/Layout) and route the four isRelation* pattern matchers and the two collapse/expand checks through it. Deliberately NOT isRelationEqual: that adds an ISL sampling (tryProveUnequal) pass that is wasted work on these one-off comparisons and measurably slowed layout propagation. heir-opt builds clean; layout_propagation, convert_to_ciphertext_semantics, and the Rotom suites all pass.
Code-review finding google#6. createMaskForPoints is atomic (validates all points, then builds one constant) but the three diagonal matvec kernels called it *after* emitting rotate/mul ops, so an (unreachable) mask failure returned failure() with a partially built chain left in the IR -- the caller treats that as 'kernel not applicable' and falls back, orphaning the dead ops. The masks are constants independent of the accumulation, so build them up front: the BSGS and dense kernels create the single gap mask right after the builder, and lowerMatmulByRotation pre-builds every group's mask before the emit loop. Now the only fallible step runs before any mutation, so a failure cannot orphan a rotation chain. Rotom execution + lit tests pass.
Code-review finding google#13. The func-arg layout propagation hand-rolled the block-arg -> func-arg attribute storage (dyn_cast<FunctionOpInterface> + getArgNumber + getArgAttr/setArgAttr). The shared association layer already implements exactly that routing (rule 1: a FunctionOpInterface block argument maps to its arg attr), so use findAttributeAssociatedWith to read the existing arg layout and setAttributeAssociatedWith to write it. Keeps the func-arg targeting (only function arguments need this) and the conflict diagnostic. No BUILD change (AttributeUtils already a dep). Rotom tests pass.
Code-review finding google#12. normalizeRowVectorMatmuls was a structural matmul rewrite (x.W^T -> transpose(matmul(W, x^T))) hardcoded at the top of the SeedLayout pass, which is otherwise about seeding layout attributes. Move it into a dedicated rotom-normalize-matmuls pass so the rewrite is independently registered, testable, and reorderable. - New pass NormalizeMatmuls (td/h/cpp) carrying the verbatim rewrite, with Linalg/Tensor as dependent dialects (it creates those ops). - SeedLayout drops the rewrite, its call, and the now-unused Linalg/Tensor includes + BUILD deps. - The mlir-to-rotom-ciphertext pipeline runs rotom-normalize-matmuls before rotom-seed-layout; the lit test invokes it explicitly. heir-opt builds clean; Rotom lit + execution tests pass.
Code-review finding google#9. lowerMatvecCtDiagonalBsgs and lowerMatvecDenseDiagonal duplicated ~120 lines of baby-step/giant-step accumulation, squat residual rotate-and-sum, and gap masking. Factor that into one emitDiagonalBsgs helper parameterized by the contraction vector, an extract-diagonal callback, the diagonal count, the rotation modulus, the residual limit, the gap mask, and an optional layout tag. The tag is a parameter rather than unified behavior on purpose: the single-period kernel works at the full ciphertext width (numSlots == output width) so its intermediates carry the output layout, whereas the dense kernel works K-wide (!= the N-wide output) and must not tag them. Each kernel keeps its own structure verification, gap-mask construction, and output placement. Behavior-preserving: the execution test (both single-period and dense matvecs) and the Rotom lit tests pass.
Replace the out-parameter with a small struct { size_t length; int64_t
straddleSlotExtent; } so the don't-care caller no longer declares a throwaway
local. Behavior unchanged; build + 28 tests pass.
…actor phase A)
Relation-preserving first step toward a strided layout model. The ct/slot
straddle was tracked as a StraddleRole{None,High,Low} per piece plus a global
straddleSlotExtent, special-cased in the ISL address emitter (High -> floorDiv,
Low -> mod). Replace that with a per-piece mixed-radix digit descriptor on
LayoutData: pieceDivBy/pieceModBy, where a piece reads digit = (i / divBy) mod
modBy of its tensor index. A whole-dim piece is (1, 0) => i; a straddle splits
into a ct piece (L, 0) => i/L and a slot piece (1, L) => i mod L.
This drops the StraddleRole enum and the straddle special-casing in
emitSegmentAddress/emitSplitCtSlotIsl in favor of one uniform digit rule that
also generalizes to N-way splits (groundwork for making the attribute stride the
address weight). Attribute syntax and seeding are unchanged; the emitted ISL
relations are byte-identical -- the lowering unit test, materialize lit tests,
and execution test all pass.
…ensor axis Refactor phase B. Make the dim attribute's stride meaningful and unify the old 'repeated dim id' tiling with the straddle's mixed-radix split: - getDim() is the tensor axis. Pieces sharing a getDim() are a mixed-radix split of that one axis and share its domain variable (preprocessing dedups the traversal entry per tensor dim). - For a multi-piece axis, the attribute stride is the within-axis digit divisor: digit = (i / stride) mod extent, with strides being the cumulative products of the lower extents (1, e0, e0*e1, ...). Placement and the ct/slot split stay positional (the cumulative product of extents crosses n at the boundary). - A single-piece axis keeps the legacy behavior (stride ignored, digit == i), so existing layouts -- including seeds with non-unit strides and the detected straddle -- are byte-identical. - LayoutAttr::verify rejects an axis whose pieces are not a valid mixed-radix decomposition (e.g. duplicate strides). - The ISL emitter now sums per piece, not per dim, so an axis can place several digits in one segment. The old position-based 'tiled duplicate dim' interpretation is replaced by this: the tiled-row-major tests keep the same (ct,slot) packing, now expressed with one domain variable per axis. Full Rotom suite (lit/unit/execution), convert-to-ciphertext-semantics, and layout_propagation pass.
…muls pass Removes all Rotom matmul/diagonal machinery, keeping elementwise and rolled-layout functionality: - lowering kernels (lowerMatmul, diagonal/BSGS matvec) in RotomTensorOpLowering - the convert-to-ciphertext matmul pattern branch (supportsRotomMatmul) - matmul cost/layout/kernel-selection paths in LayoutAssignment - the NormalizeMatmuls pass (.td/.h/.cpp + BUILD/Passes.h/heir-opt wiring) - the dead KernelName::RotomMatmul enum value across Kernel + LayoutOptimization - dependent matmul lit tests and the matmul lowering unit test
With the diagonal matvec layouts removed, the only producer of un-split straddling dims is gone: SeedLayout always splits a boundary-spanning axis into an explicit ct piece + slot piece. So inferCtPrefixLen no longer needs to detect a straddle and auto-insert a slot piece during preprocessing. - inferCtPrefixLen returns size_t (the CtPrefix struct collapses to its one remaining field); a dim that does not fit the slot budget simply stays on the ciphertext axis. - preprocessLayoutData drops the auto-split branch; boundary-spanning axes must be expressed as explicit mixed-radix splits. - layoutNumCiphertexts counts the ct prefix directly, no straddle divide. - Migrate the straddle unit test to an explicit mixed-radix split.
A roll(i, j) shifts one dim by the other modulo their shared extent, so the two rolled dims must have equal extents (only the extents must match; strides may differ). Tighten the verifier accordingly and drop the dead non-zero extent check (DimAttr::verify already guarantees size > 0).
Rename LayoutData::pieceDivBy -> pieceStride and pieceModBy -> pieceExtent to match Rotom's [dim:extent:stride] notation. pieceExtent now always holds the piece's true extent (never a 0 sentinel): the "drop the redundant modulus on the most-significant digit" optimization moves from a preprocessing sentinel to an emitter-side check (stride * extent < the axis full extent). Byte-identical ISL output; also simplifies preprocessLayoutData.
…ivalence Reverts the pieceStride/pieceExtent rename. Since the sibling field is the positional address coeff (not a Rotom dim), the digit descriptor keeps the implementation-oriented divBy/modBy names (and the 0 = no-modulus sentinel), with a comment noting they are the piece Rotom stride/extent.
Refactor phase A of the LayoutAssignment split. Each of the 12 per-op candidate visitors (visitFunc/visitElementwise/visitTranspose/... ) becomes a free static generate*(AssignmentContext&, XOp) function; visitOperation dispatches to them. The pass implements the new narrow AssignmentContext interface -- the seed/candidate/cost API the generators call back into -- so each generator depends only on that, not on the pass struct. This is the decoupling that lets the generators move to per-category files next. Also tighten the file-local anonymous namespace to wrap only the LayoutAssignment struct; the free helpers and generators are now plain static functions per the LLVM convention. Behavior-preserving: builds clean, 22/22 Rotom lit + execution tests pass.
Move getPlainValueType / isTensorLike / isLayoutCompatibleWithValue out of
LayoutAssignment.cpp into a small ValueUtils.{h,cpp}. These are the shared
value-type predicates the per-op generators rely on; giving them a header lets
the generators move to their own files (refactor phase C) without depending on
the pass translation unit. Behavior-preserving: 22/22 Rotom tests pass.
Refactor phase C. The 12 generate* functions now live in four op-family files under LayoutAssignment/gen/ -- Structural (func/secret.generic/yield/ passthrough), Elementwise (add/sub/mul + linalg.generic), ReduceTranspose, and Reshape (collapse/expand/extract/insert) -- all declared in Generators.h. The linalg-body predicates (isElementwiseGeneric/hasAddLikeBody) move with the elementwise generator. LayoutAssignment.cpp drops from ~680 to ~420 lines: it now holds the pass struct, the AssignmentContext implementation, the op dispatch, and the selection phase. Adding a new tensor op is now: a generator in the relevant gen/ file + a Generators.h declaration + one dispatch line. Behavior-preserving: builds clean, 22/22 Rotom lit + execution tests pass.
enumerateMatmulPlans derives, per (lhs, rhs) layout pairing in (i, j, k) iteration-space dims, the deterministic aligned placements for expand -> multiply -> sum-k matmul: compute/expanded/result layouts plus raw rotation/add counts for replication fill (free across ciphertexts, log2 within slots) and the k reduction. No kernel names; M2 will price these in generateMatmul and M3 will re-derive them in lowering.
linalg.matmul candidates now come from ContractionAlignment's deterministic plans: operand candidates are relabeled into the (i, j, k) iteration space, each plan is priced as operand alignment (convert-to-pre-fill + replication fill) + ciphertext multiplies + k-reduction, and the result candidate carries the plan's result layout with no kernel name -- the lowering (M3) re-derives the same plan from the assigned layouts.
ConvertRotomMatmul re-derives the deterministic ContractionAlignment plan from the (lhs, rhs, result) rotom layouts the assignment records under the rotom.matmul op attribute, then emits: one same-shape convert_layout per operand onto its inner expanded placement (the shift network fills slot replication), free outermost ciphertext copies via concat, one elementwise multiply, a log-tree rotate-and-reduce per k slot piece, and contiguous row-block adds for an outermost k ciphertext piece. No kernel attribute involved. ContractionAlignment now speaks each tensor's own dims at its boundary (computeLayout stays in iteration space) and exposes stripOuterCtReplication + isLowerableMatmulPlan; generateMatmul gates candidates on lowerability and prices operand alignment with the shift-network conversion cost, so the assignment never selects a plan the lowering cannot realize (4x4 matmul thus needs n=64; matvec works at n=16).
The deduped traversalDims kept first-appearance order, and the ISL lowering emits its domain variables in that order -- but every tensor_ext consumer reads domain variables positionally as tensor dims. Any layout whose pieces lead with a later dim (e.g. column-major [1:.][0:.]) thus materialized with silently permuted dims: row- and column-major 4x4 produced identical ISL, so conversions between them priced (and lowered) as free no-ops. preprocessLayoutData now sorts traversalDims by dim id and remaps pieceIndex, making the domain order a defined invariant. Re-blessed the tests that encoded the old convention: the column-major evaluate tests index the domain as [row, col] (expected packings unchanged), the column-major materialize tests now check genuinely column-major relations, and the assignment winners that flipped once conversion costs became honest (elementwise now emits the real transpose remap for mixed row/column-major seeds; the 4x4 matmul picks the lhs-hosted slot plan).
Execution tests run matvec (n=16) and 4x4 matmul (n=64) through seed -> assign -> materialize -> convert -> shift-network and compare against plaintext references, checking exactly the slots the result layout claims. Writing the expected packing exposed a wrong claim from M1: after the cyclic log-tree rotate-and-reduce, only the k=0 offset holds the true sum (other offsets hold window sums whose carries cross into the digit above k), so a summed slot-k piece now becomes a gap in resultLayout, not replication.
New rotom-outline-kernels pass (after layout assignment): each linalg.matmul carrying a rotom.matmul layout combination moves into a private kernel function keyed by (op, layouts, types) -- two matmuls with the same signature share one callee -- and the call site carries the result layout. The materializer and the ciphertext lowering process kernel bodies like any other layout-assigned function, keeping the Rotom kernel vocabulary contained per function and unit-testable in isolation. Pipelines (mlir-to-rotom-ciphertext and the execution test) inline the calls after the ciphertext lowering so cross-kernel optimization and backends without call support see a flat function; tensor_ext gains the trivial DialectInlinerInterface this requires, and heir-opt registers the func-dialect inliner extension.
OutlineKernels now covers every Rotom-lowered tensor operator, not just matmul: elementwise arith ops carrying a Rotom secret.kernel outline into per-signature functions too (operand layouts read from the per-value rotom.layout attributes, result layout carried on the call). The op-specific logic is reduced to a target predicate plus operand/ prologue selection over a shared outlineOp/buildKernel core, so future ops (reduce, transpose) are one dispatch entry each. Outlining is gated to plain-func bodies for now; secret.generic regions are a TODO with the secret-path work.
Per design rule: gate by cost, not capability. planLayoutExpansion (LayoutAlignment) decomposes any layout pair -- including different ciphertext counts, which tensor_ext.convert_layout cannot express -- into deterministic rotate/mask/accumulate steps grouped by (targetCt, sourceCt, shift). The matmul generator prices exactly those steps (pure ct-replication prices free), ConvertRotomMatmul emits them (per-target-ct extract/rotate/mask/add, free copies for full-row zero-shift steps), and the k-ciphertext reduction is generalized to digit-stripping row gathers for any k placement. isLowerableMatmulPlan and stripOuterCtReplication are deleted: every enumerated plan is now emittable, so search prunes expensive expansions purely by cost. The 4x4 matmul at n=16 -- previously unassignable -- now lowers and executes numerically exactly (new execution test).
roll(a, r) with r a replication dim of equal extent makes replica d hold dim a cyclically rotated by d: the layout materializes every rotation of the rolled data, so downstream alignment becomes replica selection instead of a shift network. The verifier now permits a replication dim in the roll-by position (the rolled dim must stay traversal; gaps stay rejected), and the ISL emitter substitutes the replica's existential variable into the roll expression. Evaluation test pins the packing (replica d = rotate-left-by-d). Also fixes a latent assert: makeCheckedLayout streamed verifier diagnostics into an inactive InFlightDiagnostic; it now uses a real, silenced emitter.
A roll-by dim may now be a gap of equal extent. The gap's block index is the shift, so block g holds the rolled dim cyclically shifted by g: a rolled-by gap claims its blocks with the rotations (a plain gap stays unclaimed space folded to zero). The ISL emitter existentially quantifies rolled-by gaps, emits their address term, and counts them in the ciphertext bound; layoutNumCiphertexts counts rolled ct-prefix gaps the same way. Rolling FROM a gap or replication dim stays rejected (the block index is existential, so the shift is a no-op).
…ering planLayoutExpansion now backs the general conversion path, not just matmul operand alignment: cachedConversionCost prices a ciphertext- count-changing conversion (expansion or compaction) by its explicit rotate/mask/accumulate steps, supportsRotomAlignmentLowering no longer requires equal ciphertext counts, and the elementwise lowering emits those steps via a shared RotomTensorOpLowering::convertToLayout helper (convert_layout when the count is unchanged). ConvertRotomMatmul's expandOperand delegates to the same helper. Compaction thus competes on cost anywhere a conversion is priced, per the gate-by-cost rule. New coverage: a planner unit test for compacting a gapped 4-ct matmul result into one column-major ciphertext, and a numeric pipeline test adding operands seeded at different ciphertext counts.
Extends enumerateMatmulPlans with rolled variants: a unit-stride k piece in the ciphertext prefix rolled by a same-extent unit-stride i/j piece in the slot region. The roll is positional metadata on the compute footprint, inherited piece-for-piece by both operand expansions (the operand that does not own the partner dim rolls by the replication piece that subsumed it), so the multiply and ciphertext-add counts are unchanged and the result is unrolled. Hosts also gain a reverse-subsumption variant -- the free dim placed over an existing same-extent replication piece -- so operands already at an expanded placement (rolled or not) enumerate the compute placement they came from at zero conversion cost. Since a rolled plan and its roll-free sibling share a result layout, the layout assignment now records the priced winner's computeLayout as a fourth element of rotom.matmul (via selectMatmulPlan, sharing the exact generator cost formula), and ConvertRotomMatmul matches plans by that unique identity, falling back to result-layout matching for 3-element attrs. The seeded ct-diagonal pair a = roll(0,2)[1:4];[0:4][R:4], b = roll(0,2)[0:4];[R:4][1:4] now lowers to 4 ciphertext multiplies and 3 adds with zero rotations, leaving one row-major result ciphertext; verified numerically end to end.
A source candidate (a value straight from seeding: a secret function argument or cleartext feeding secret compute) is data packed at encode time, so a matmul plan's expanded operand placement is reachable by packing the source there directly instead of converting in ciphertext space. generateMatmul now offers this repack option per operand alongside the priced conversion: the candidate's assignment overwrites the source's layout with the expanded placement at a small epsilon cost (carried on the source's assignment entry so multiple consumers charge it once, and large enough that a plan whose operands already sit at their seeded layouts wins ties -- the user's packing stands when nothing beats it). Shared-source consistency is preserved by the existing assignment merge: a source consumed elsewhere at a different layout rejects the combination and the conversion-priced candidates take over. Winners improve across the board: the n=64 matmul and n=16 matvec lit tests now pack operands at their compute placements outright (the matvec kernel drops its operand remap entirely -- multiply plus two rotate-adds), and the n=16 row-major-seeded matmul repacks both sources onto the 4-ciphertext ct-k placement, summing k with plain ciphertext adds. A new numeric test pins the conversion path: a matmul of an intermediate (a + a) still converts that operand in ciphertext space, since only sources can repack.
…R2.5) Two search goals, per design discussion: (1) cheap conversion/op costs, (2) compact ciphertext counts (future bootstrap placement wants compact candidates). Goal 2: RotomCostModel gains a ciphertextCount carrying weight, charged once per assigned value on its layout's ciphertext count (in setCandidates; block arguments and yields alias an already-charged value and are exempt; repacked sources re-charge at the expanded placement). This keeps encode-time repacking honest: blowing a source up to four ciphertexts is no longer epsilon-cheap, so compact packings plus rotation-only ciphertext-space expansions win when they should. Goal 1: a unit test pins that planLayoutExpansion already prices the replicate-then-roll route -- expanding compact column-major onto the rolled-by-replication placement (roll-by partner outermost in slots) is one full-ciphertext rotation per target ciphertext, no masks, since ciphertext replication is free and the roll's mod-extent wrap coincides with the whole-ciphertext cyclic rotation. A numeric test runs the route end to end: a column-major-seeded lhs stays at ONE ciphertext (the carrying cost outweighs the fat repack) and is expanded by pure rotations in ciphertext space, against a pre-rolled rhs; result matches the plaintext reference. The downstream_selects_intermediate lit fixture now uses equal-ciphertext-count seed orderings so it isolates the downstream-preference mechanism from the new compactness pressure.
[R:4:1] and [G:4:1] instead of [-1:4:1] and [-2:4:1] -- easier to read. The numeric ids are still accepted on input and round-trip to the letter forms (covered in syntax.mlir). Test spellings updated; the n=64 matmul lit test's cost-tied winner shifted because the structural tie key orders candidates by their printed form.
Generalizes the R1 roll decoration: a unit-stride k piece anywhere in the compute footprint may roll by any same-extent unit-stride traversal piece elsewhere in it (pairs entirely inside the ciphertext prefix are skipped -- that only permutes ciphertext contents). This adds the slot-diagonal family to the existing ct-diagonal one: a slot k rolled by a slot piece is the classic Halevi-Shoup diagonal packing, and by a ciphertext piece the replicate-then-roll form. The k reduction keeps its footprint shape -- per remaining coordinate the rolled index is a bijection of k, so the slot rotate-and-reduce still sums k and the result gaps the piece as usual -- and the lowering needs no changes at all. Also fixes a pricing bug this exposed: conversionMoves atomizes dims only and is blind to rolls, so its empty-moves fast path claimed conversions between rolled and unrolled layouts of the same footprint (e.g. row-major onto the diagonal packing) were free. Rolled layouts now return the sentinel and defer to the relation-based pricing. Numeric test: a matmul whose lhs is an intermediate held at the one-ciphertext diagonal packing -- unable to repack, its only cheap expansion is onto the slot-diagonal plan's replicated-diagonal placement (four zero-shift copies, no rotations or masks). The search picks the Halevi-Shoup plan (pinned by the four-ciphertext gapped result), the compact diagonal source stays at one ciphertext, and the lowered kernel is one multiply plus a two-rotation log-tree reduce; values match the plaintext reference.
A seeded value is packed at encode time, so its diagonal packings are available at the same zero packing cost. seedValue now widens each seed with enumerateSingleRollVariants: every verifier-legal, materializable single roll (from a unit-stride traversal piece, by a same-extent unit-stride traversal or replication piece, skipping pairs entirely inside the ciphertext prefix). A rolled placement materializes every rotation of the rolled piece across its partner's blocks, so any consumer -- the elementwise alignment, conversion pricing, and matmul plan enumeration all already handle rolled layouts -- can align against it by block selection instead of slot permutation; whether a variant wins stays a cost decision. Observable wins: row-major and column-major seeded elementwise operands now each adopt their diagonal variant, leaving a smaller residual conversion between the two diagonal packings than the direct slot swap (lit test re-blessed); a new numeric test pins the zero-conversion case, where one operand arrives packed only diagonally and the other source adopts the matching diagonal variant, lowering to a single add with no rotations or convert_layouts. Variants are filtered by isMaterializableRotomLayout up front: candidates become assigned layouts, which must materialize -- an unfiltered variant reached the same-count conversion pricing assert.
Measurement showed the VVE shift network quotes n-1 rotations for rolled-replication fills (15/63/255/1023 at D=4..32) where D-1 rotations suffice. Fixes: - planLayoutExpansion gains replica-aware greedy grouping: per target point, candidate source keys are ranked by (existing-group, zero-shift, smallest key), so replicated sources coalesce into D-1 rotation groups. - New chooseSameCountConversion(from, to, n) plans explicit steps AND quotes the VVE shift network, returning whichever needs fewer rotations (ties favor VVE). Both cachedConversionCost (pricing) and convertToLayout (emission) call it, so the priced route is the emitted route. - conversionMoves stays roll-blind and is not consulted here. - Regression test SameCountConversionPrefersStepsForRolledFill pins useSteps with D-1 rotations; the elementwise lowering lit test is re-blessed from tensor_ext.remap to 3 explicit rotates + masks.
enumerateMatmulPlans previously built footprints from host dims only (host rolls vanished) and always emitted an unrolled result. Now: - Footprints carry the host's own rolls (positions shifted for the ct-prepended free piece), minus reduction-incompatible pairs (a non-k dim rolled BY k does not commute with the k-sum). A rolled operand hosts its diagonal form at zero conversion, including multi-roll hosts the single-roll decorations cannot re-derive. - Roll decorations generalize beyond k: an i/j piece rolled by the other free/host traversal dim or by a replication piece (the free-swap diagonal). These commute with the k-sum, so the RESULT inherits them at its surviving positions -- a matmul can produce a diagonal result directly, the operand form a downstream Halevi-Shoup matmul consumes. Decorations stack on either base (roll-free or host-rolled) unless the composition would be order-dependent. - Expansions drop rolls whose FROM piece they subsume into replication (the operand does not own the dim, so the roll is a no-op for it); a subsumed BY piece keeps the roll on the replication piece as before. - Plans whose result layout cannot materialize are skipped: the result becomes a value's assigned layout (expansions were already priced). Row-major pair grows from 9 to 15 deduped plans (two free-swap variants per footprint). New unit tests pin the free-swap plan's expansions and rolled result, the non-seeding of incompatible host rolls, and rolled hosting of a chained diagonal result; a numeric (A x B) x C pipeline test executes the chain, where the search hosts the intermediate on a rolled ct-diagonal plan for the second matmul at zero conversion.
roll(i, j) rewrites dims[i]'s index to (idx_i - idx_j) mod size(dims[i]), which is well-defined for any partner extent: a smaller partner covers a prefix of the rotations, a larger one wraps. The materializer already reduced mod the from extent, so only the verifier's equal-extent check and the enumeration filters (matmul roll decorations, single-roll seed variants) blocked unequal pairs. Size-1 pieces stay excluded from enumeration on either side of a pair: a mod-1 roll is an attr-distinct identity that would pollute plan dedup. This unlocks rolled plans for rectangular matmuls (k extent different from i/j) and is a prerequisite for BSGS: the baby-step expansion is a layout whose k:D piece rolls by a sqrt(D)-extent partner. Re-blessed: the mismatched-extent verifier lit case is now positive; single-roll variants of [0:4][R:4][1:2] grow 1 -> 4; the non-power-of-two plan count grows 5 -> 10; layout_assignment's generic_assign seeds adopt unequal-extent diagonal variants (rolls = [(1, 0)] on both args). New coverage: ISL materialization of both unequal orientations, the rectangular replicate-then-roll plan family, and numeric rectangular matmul plus rectangular elementwise diagonal adoption.
Two roll extensions that together express baby-step/giant-step diagonal packings as ordinary layouts: - Scales: a roll tuple takes an optional third element (from, by, s), shifting by s times the partner index -- (idx_from - s*idx_by) mod extent. Stored in a parallel optional rollScales parameter whose canonical form is omitted when every scale is 1, so scale-free layouts unique to the same attribute as before; the printer emits pairs for unit scales. - Split-dim rolls: a roll FROM any piece of a mixed-radix split dim rewrites the WHOLE dim's index mod its full extent (each piece then takes its digit of the rolled index), and a roll BY a split piece shifts by that piece's digit of the dim's current (possibly already rolled) expression. The materializer now resolves roll pieces to traversal variables by dim id (preprocessing dedupes one variable per logical axis), where it previously matched (dim, size, stride) and failed on any split-dim roll. The BSGS diagonal packing of a 16x16 matrix materializes and evaluates correctly: dims [[k:4:4],[k:4:1],[i:16:1]] with rolls [(0, 2), (2, 0, -4)] puts iteration point ((a-4g) mod 16, (a+b) mod 16) in ciphertext (g, b) slot a -- giant-step partial sums pre-rotation. The baby-step vector expansion needs no new machinery (an unequal-extent roll by replication). Plan enumeration and the kernel reduce mode land separately; enumeration is unaffected (it builds unit-scale rolls only).
Enumeration: a ciphertext k piece of composite extent D = G * B splits into giant/baby digits [k:G:B][k:B:1]; the whole k is rolled by the outermost slot partner (same-extent i/j) and the partner is rolled back by the giant digit with scale -B. The operand owning the partner inherits the placement positionally -- the BSGS diagonal packing, free for a repackable source -- while the other operand collapses to the BABY form: replication over both k digits with its k placed whole at the partner's slot position, rolled by the baby replication. The result is fully unrolled (the giant rotations consume the partner roll). Kernel reduce prices numCtResult * (G - 1) giant rotations. Lowering: the ct-k sum reads giant structure straight off the compute layout (a scaled roll from a slot piece by a ct k digit); source blocks group per result ciphertext by their giant shift, sum plainly within a group, and only the G per-shift partial sums rotate into place. Conversion pricing and emission now share rotated rows across steps with the same (source ciphertext, shift): the lowering emits one extract/ rotate per distinct pair and reuses it for every target it feeds, and both the ct-changing pricing and chooseSameCountConversion count rotations the same way -- this is what lets the baby expansion price its B-1 rotations instead of one per replicated block. A 16x16 matvec with an intermediate vector operand now picks the BSGS plan end to end: 3 baby + 3 giant = 6 rotations against 15 for the full rolled expansion, verified numerically against the plaintext reference.
When a footprint's slot suffix leaves slack (an implicit front gap of s = n / slot extent), a unit-stride ciphertext piece can straddle the ct/slot boundary instead: enumerateMatmulPlans adds densified footprint variants splitting such a piece into a ct digit [d:E/s:s] and a slot digit [d:s:1] that fills the gap, dividing the ciphertext count by s. Splitting k this way and letting the existing roll decorations roll its slot digit by i (a roll FROM either digit rewrites the whole k index) yields the dense Halevi-Shoup packing: s diagonals per ciphertext. All downstream machinery -- decorations, expansions, result construction, counts, the kernel lowering, and the BSGS generator (whose giant/baby split lands its baby digit in the slack automatically at larger n) -- treats densified footprints like any other. A 16x16 matvec at n = 64 now repacks its matrix at 4 dense ciphertexts instead of 16 three-quarters-empty ones (the search picks the dense-BSGS hybrid), with the vector and result staying single ciphertexts, verified numerically.
This was referenced Jul 7, 2026
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.
do not merge yet!