Rotom: axis roll endpoints, roll steps, and per-piece roll semantics

A roll endpoint 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 endpoint 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 also take an optional step s, stored in a parallel rollSteps array
(canonically omitted when all 1): the FROM index moves by -s * shift(by)
mod its extent. A baby-step/giant-step pre-rotation's magnitude is
expressible structurally (a piece FROM shifts its axis in strides of the
piece), but its direction is not -- the kernel's shared baby rotations
need the pre-rotation to run against the roll's subtraction convention,
and the mirror-direction packing would need per-giant baby rotations,
forfeiting the BSGS saving -- so the signed step stays explicit.

Together these make the BSGS diagonal packing an ordinary layout:

  #rotom.layout<n = 16, rolls = [(axis 1, 2), (2, 0, -4)],
                dims = [[1:4:4], [1:4:1], [0:16:1]]>

where roll 1 diagonalizes the whole split k against i (ciphertext (g, b)
holds the digits of (k - i) mod 16) and roll 2 pre-rotates i by 4g,
giant digit read from the rolled k.

Author: edwjchen

lib/Dialect/Rotom/IR/RotomAttributes.cpp

@@ -15,6 +15,48 @@
 
 namespace mlir::heir::rotom {
 
+// One endpoint of a roll: either one piece of the layout (a position in its
+// dims list) or a whole tensor axis (spelled `axis N`; 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).
+struct RollEndpoint {
+  bool isAxis;
+  int64_t index;  // Dims-list position, or the tensor axis id when isAxis.
+  bool operator==(const RollEndpoint& other) const {
+    return isAxis == other.isAxis && index == other.index;
+  }
+};
+
+// Endpoint encoding in the flat rolls storage: a piece endpoint is its
+// non-negative dims-list position, an axis endpoint is -(axis + 1).
+inline int64_t encodeRollEndpoint(RollEndpoint e) {
+  return e.isAxis ? -(e.index + 1) : e.index;
+}
+inline RollEndpoint decodeRollEndpoint(int64_t encoded) {
+  return encoded < 0 ? RollEndpoint{true, -encoded - 1}
+                     : RollEndpoint{false, encoded};
+}
+
+// One roll of a layout: FROM's index is rewritten to
+// (idx_from - step * shift(by)) mod extent(from), where a piece FROM
+// rewrites only its own mixed-radix digit and an axis FROM rewrites the
+// whole axis index. `step` is 1 unless the layout carries an explicit
+// rollSteps entry.
+struct RollSpec {
+  RollEndpoint from;
+  RollEndpoint by;
+  int64_t step;
+};
+
+// The layout's rolls with both endpoints decoded and their steps attached
+// (step 1 where the rollSteps storage is absent). Reading rolls this way
+// keeps a caller from copying roll pairs while silently dropping their
+// steps.
+llvm::SmallVector<RollSpec> getRollSpecs(LayoutAttr layout);
+
+// Whether every roll of the layout has unit step (or there are no rolls).
+bool allRollStepsAreUnit(LayoutAttr layout);
+
 enum class LayoutPieceKind { Traversal, Replication, Gap };
 
 struct LayoutData {

lib/Dialect/Rotom/IR/RotomAttributes.h

@@ -57,26 +57,53 @@ def Rotom_LayoutAttr : Rotom_Attr<"Layout", "layout"> {
     ciphertext side of Rotom's `;` split (one piece); remaining entries are
     in-slot. See [Section 4.2 of the Rotom paper](https://eprint.iacr.org/2025/1319.pdf).
 
-    Optional **rolls** encode a `roll(i,j)` metadata object: each pair `(i, j)`
-    indexes into the `dims` array (the flattened `ct_dims + slot_dims` list) and
-    rewrites `dims[i]`'s index to `(idx_i - idx_j) mod size(dims[i])`. The two
-    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`
+    list, spelled as a bare integer -- or a whole tensor *axis*, spelled
+    `axis N` (legal only when axis `N` is packed as more than one piece; the
+    piece spelling is canonical for an unsplit axis, where the two coincide).
+
+    A piece FROM rewrites that piece's own mixed-radix digit in place,
+    `digit(from) <- (digit(from) - s * shift(by)) mod extent(from)`, leaving
+    the axis's other digits untouched. An `axis` FROM rewrites the whole axis
+    index modulo its full extent, and each piece then takes its digit of the
+    rolled index -- the shift borrows across digits, which no combination of
+    piece rolls can express. The shift is the by endpoint's index: a
+    traversal piece's digit (its whole index when the axis is unsplit), a
+    whole axis's index via `axis`, a replication piece's replica index, or a
+    gap piece's block index. The two extents need not match: the shift
+    reduces modulo the rolled extent, so a smaller partner covers a prefix of
+    the rotations and a larger one wraps.
+
+    The step `s` defaults to 1 and may be spelled as a third tuple element
+    `(from, by, s)`. A baby-step/giant-step pre-rotation runs AGAINST the
+    roll's subtraction direction (step `-B` shifts by `+B * shift(by)`);
+    the magnitude of a step is often expressible structurally instead (a
+    piece FROM shifts its axis in strides of the piece), but the direction
+    is not, so the step stays explicit.
+
+    Optional **rollSteps** stores one step per roll pair; it must be omitted
+    when every step is 1 (the canonical form the parser and builders produce).
   }];
 
   let parameters = (ins
     "::mlir::ArrayAttr":$dims,
     "int64_t":$n,
-    OptionalParameter<"::mlir::DenseI64ArrayAttr">:$rolls
+    OptionalParameter<"::mlir::DenseI64ArrayAttr">:$rolls,
+    OptionalParameter<"::mlir::DenseI64ArrayAttr">:$rollSteps
   );
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 
   let extraClassDeclaration = [{
-    /// Layout with no `roll(i,j)` metadata (empty rolls storage).
+    /// Layout with no `roll(from, by)` metadata (empty rolls storage).
     static ::mlir::heir::rotom::LayoutAttr get(::mlir::MLIRContext *context,
                                                ::mlir::ArrayAttr dims, int64_t n);
+    /// Layout with unit-step rolls (empty step storage).
+    static ::mlir::heir::rotom::LayoutAttr get(::mlir::MLIRContext *context,
+                                               ::mlir::ArrayAttr dims, int64_t n,
+                                               ::mlir::DenseI64ArrayAttr rolls);
   }];
 }
 

lib/Dialect/Rotom/IR/RotomAttributes.td

@@ -29,6 +29,8 @@ cc_test(
         "@heir//lib/Dialect/Rotom/IR:Dialect",
         "@heir//lib/Utils/Layout:Evaluate",
         "@heir//lib/Utils/Layout:IslConversion",
+        "@heir//lib/Utils/Layout:Utils",
+        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
     ],

lib/Dialect/Rotom/Utils/BUILD

@@ -18,20 +18,40 @@
 namespace mlir::heir::rotom {
 namespace {
 
-/// Maps a `#rotom.dim` from the layout's `dims` list to its iterator index `i*`
-/// after preprocessing (match logical axis, size, and stride).
-static FailureOr<int64_t> traversalIndexForRotomDim(
-    const SmallVector<DimAttr>& traversalDims, DimAttr want) {
+/// Maps a tensor axis id to its iterator index `i*` after preprocessing.
+/// Preprocessing dedupes traversal dims per logical axis (a mixed-radix split
+/// contributes one variable), so the axis id identifies the variable.
+static FailureOr<int64_t> traversalIndexForAxis(
+    const SmallVector<DimAttr>& traversalDims, int64_t axis) {
   for (int64_t i = 0; i < static_cast<int64_t>(traversalDims.size()); ++i) {
-    if (traversalDims[i].getDim() == want.getDim() &&
-        traversalDims[i].getSize() == want.getSize() &&
-        traversalDims[i].getStride() == want.getStride()) {
-      return i;
-    }
+    if (traversalDims[i].getDim() == axis) return i;
   }
   return failure();
 }
 
+/// Product of the extents of `dim`'s pieces in the layout's dims list (the
+/// logical axis extent; a single whole-dim piece gives its own size).
+static int64_t fullExtentOfDim(ArrayAttr rotomDims, int64_t dim) {
+  int64_t extent = 1;
+  for (Attribute a : rotomDims) {
+    auto d = cast<DimAttr>(a);
+    if (!d.isGap() && !d.isReplicate() && d.getDim() == dim) {
+      extent *= d.getSize();
+    }
+  }
+  return extent;
+}
+
+/// Whether `dim` is packed as more than one mixed-radix piece.
+static bool isSplitDim(ArrayAttr rotomDims, int64_t dim) {
+  int64_t count = 0;
+  for (Attribute a : rotomDims) {
+    auto d = cast<DimAttr>(a);
+    if (!d.isGap() && !d.isReplicate() && d.getDim() == dim) ++count;
+  }
+  return count > 1;
+}
+
 static std::string modExpr(llvm::StringRef expr, int64_t mod) {
   std::string out;
   llvm::raw_string_ostream os(out);
@@ -84,7 +104,7 @@ static LogicalResult emitSegmentAddress(
     const SmallVector<DimAttr>& replicationDims,
     int64_t numActiveTraversalComponents, size_t segStart, size_t segEnd,
     const llvm::DenseSet<int64_t>& rolledGapIndices, ArrayRef<int64_t> rolls,
-    ArrayAttr rotomDims) {
+    ArrayRef<int64_t> rollSteps, ArrayAttr rotomDims) {
   // Effective extent of a piece = the range of its mixed-radix digit: modBy
   // when the piece reads i mod L (the slot/low part), else sz/divBy (the whole
   // dim when divBy 1, or the ct/high part sz/L of a straddling dim).
@@ -166,56 +186,117 @@ static LogicalResult emitSegmentAddress(
     traversalExprs.push_back("i" + std::to_string(i));
   }
 
-  // Apply roll(a,b) transforms left-to-right:
-  // t_a <- (t_a - t_b) mod extent(a).
+  // Apply rolls left-to-right. A piece FROM rewrites that piece's own
+  // mixed-radix digit in place, (digit - step * shift) mod extent(piece),
+  // leaving its axis's other digits untouched; an `axis` FROM rewrites the
+  // whole axis index mod its full extent, so the shift borrows across digits
+  // (each piece then takes its digit of the rolled index). The shift is the
+  // by endpoint's index: a traversal piece's digit (its whole index when the
+  // axis is unsplit), a whole axis via `axis`, a slot replication's replica
+  // index, or a gap's block index. A step multiplies the shift -- the
+  // baby-step/giant-step pre-rotation runs against the roll's subtraction
+  // direction, so its sign cannot be folded into the endpoint structure.
   //
-  // The rewrite lands wherever the FROM dimension sits -- the ciphertext
-  // address, the slot address, or both when it straddles the boundary. BY is
-  // another traversal dimension on either axis, or a slot replication/gap
-  // block index, so a roll diagonalizes a ciphertext dimension against a slot
-  // one (one ciphertext per diagonal) or two slot dimensions (a Halevi-Shoup
-  // slot diagonal); a roll within the ciphertext axis is a free ciphertext
+  // The rewrite lands wherever the FROM pieces sit -- the ciphertext
+  // address, the slot address, or both when the axis straddles the boundary.
+  // A roll diagonalizes a ciphertext dimension against a slot one (one
+  // ciphertext per diagonal) or two slot dimensions (a Halevi-Shoup slot
+  // diagonal); a roll within the ciphertext axis is a free ciphertext
   // relabeling that enumeration never generates.
   if (!rolls.empty()) {
     if (!rotomDims || rolls.size() % 2 != 0) return failure();
     for (size_t i = 0; i < rolls.size(); i += 2) {
-      const int64_t fromIdx = rolls[i];
-      const int64_t toIdx = rolls[i + 1];
-      if (fromIdx < 0 || toIdx < 0 ||
-          fromIdx >= static_cast<int64_t>(rotomDims.size()) ||
-          toIdx >= static_cast<int64_t>(rotomDims.size())) {
+      const RollEndpoint from = decodeRollEndpoint(rolls[i]);
+      const RollEndpoint by = decodeRollEndpoint(rolls[i + 1]);
+      const int64_t step = i / 2 < rollSteps.size() ? rollSteps[i / 2] : 1;
+      if (!from.isAxis && from.index >= static_cast<int64_t>(rotomDims.size()))
         return failure();
-      }
-      auto fromDim = dyn_cast<DimAttr>(rotomDims[fromIdx]);
-      auto toDim = dyn_cast<DimAttr>(rotomDims[toIdx]);
-      if (!fromDim || !toDim) return failure();
+      if (!by.isAxis && by.index >= static_cast<int64_t>(rotomDims.size()))
+        return failure();
+      DimAttr fromPiece =
+          from.isAxis ? DimAttr() : dyn_cast<DimAttr>(rotomDims[from.index]);
+      if (!from.isAxis && !fromPiece) return failure();
+      const int64_t fromAxis = from.isAxis ? from.index : fromPiece.getDim();
       FailureOr<int64_t> maybeFromTrav =
-          traversalIndexForRotomDim(traversalDims, fromDim);
+          traversalIndexForAxis(traversalDims, fromAxis);
       if (failed(maybeFromTrav)) return failure();
       const int64_t fromTrav = *maybeFromTrav;
+
       std::string toExpr;
-      if (toDim.isGap()) {
-        // Rolling by a gap dim: the shift is the gap's existential block
-        // index, so block g holds the rolled dim cyclically shifted by g.
-        toExpr = "g" + std::to_string(gapIndexForPosition(rotomDims, toIdx));
-      } else if (toDim.isReplicate()) {
-        // Rolling by a replication dim: the shift is the replica index, whose
-        // existential variable is named after the piece's replication slot.
-        int64_t replicationIndex = 0;
-        for (int64_t k = 0; k < toIdx; ++k) {
-          if (cast<DimAttr>(rotomDims[k]).isReplicate()) ++replicationIndex;
-        }
-        toExpr = "d" + std::to_string(numActiveTraversalComponents +
-                                      replicationIndex);
-      } else {
+      if (by.isAxis) {
+        // Rolling by a whole axis: the shift is the axis's (possibly
+        // already-rolled) index expression.
         FailureOr<int64_t> maybeToTrav =
-            traversalIndexForRotomDim(traversalDims, toDim);
+            traversalIndexForAxis(traversalDims, by.index);
         if (failed(maybeToTrav)) return failure();
         toExpr = traversalExprs[*maybeToTrav];
+      } else {
+        auto toDim = dyn_cast<DimAttr>(rotomDims[by.index]);
+        if (!toDim) return failure();
+        if (toDim.isGap()) {
+          // Rolling by a gap dim: the shift is the gap's existential block
+          // index, so block g holds the rolled index cyclically shifted by g.
+          toExpr =
+              "g" + std::to_string(gapIndexForPosition(rotomDims, by.index));
+        } else if (toDim.isReplicate()) {
+          // Rolling by a replication dim: the shift is the replica index,
+          // whose existential variable is named after the piece's replication
+          // slot.
+          int64_t replicationIndex = 0;
+          for (int64_t k = 0; k < by.index; ++k) {
+            if (cast<DimAttr>(rotomDims[k]).isReplicate()) ++replicationIndex;
+          }
+          toExpr = "d" + std::to_string(numActiveTraversalComponents +
+                                        replicationIndex);
+        } else {
+          FailureOr<int64_t> maybeToTrav =
+              traversalIndexForAxis(traversalDims, toDim.getDim());
+          if (failed(maybeToTrav)) return failure();
+          toExpr = traversalExprs[*maybeToTrav];
+          if (isSplitDim(rotomDims, toDim.getDim())) {
+            // The by piece's digit: (i / stride) mod extent, with the modulus
+            // redundant on the most-significant digit.
+            const int64_t fullTo = fullExtentOfDim(rotomDims, toDim.getDim());
+            if (toDim.getStride() > 1) {
+              toExpr = floorDivExpr(toExpr, toDim.getStride());
+            }
+            if (toDim.getStride() * toDim.getSize() < fullTo) {
+              toExpr = modExpr(toExpr, toDim.getSize());
+            }
+          }
+        }
+      }
+
+      if (step != 1) {
+        toExpr = std::to_string(step) + " * (" + toExpr + ")";
+      }
+
+      if (from.isAxis || !isSplitDim(rotomDims, fromAxis)) {
+        // Whole-axis rewrite (the piece spelling of an unsplit axis is the
+        // same operation: its one piece's digit is the axis index).
+        std::string diffExpr =
+            "(" + traversalExprs[fromTrav] + " - " + toExpr + ")";
+        traversalExprs[fromTrav] =
+            modExpr(diffExpr, fullExtentOfDim(rotomDims, fromAxis));
+      } else {
+        // Piece rewrite on a split axis: replace this piece's digit with
+        // (digit - step * shift) mod extent(piece), i.e. add
+        // stride * (newDigit - digit) to the axis index. No borrow crosses
+        // into the other digits.
+        const int64_t stride = fromPiece.getStride();
+        const int64_t extent = fromPiece.getSize();
+        const std::string axisExpr = traversalExprs[fromTrav];
+        std::string digit = axisExpr;
+        if (stride > 1) digit = floorDivExpr(digit, stride);
+        if (stride * extent < fullExtentOfDim(rotomDims, fromAxis)) {
+          digit = modExpr(digit, extent);
+        }
+        const std::string newDigit =
+            modExpr("(" + digit + " - " + toExpr + ")", extent);
+        traversalExprs[fromTrav] =
+            "(" + axisExpr + " + " + std::to_string(stride) + " * (" +
+            newDigit + ") - " + std::to_string(stride) + " * (" + digit + "))";
       }
-      std::string diffExpr =
-          "(" + traversalExprs[fromTrav] + " - " + toExpr + ")";
-      traversalExprs[fromTrav] = modExpr(diffExpr, fromDim.getSize());
     }
   }
 
@@ -257,7 +338,7 @@ static FailureOr<std::string> emitSplitCtSlotIsl(
     const SmallVector<DimAttr>& replicationDims,
     const SmallVector<DimAttr>& gapDims, int64_t numTraversalComponents,
     int64_t numReplication, int64_t numGap, ArrayRef<int64_t> rolls,
-    ArrayAttr rotomDims) {
+    ArrayRef<int64_t> rollSteps, ArrayAttr rotomDims) {
   if (prefix > pieces.size()) return failure();
 
   const llvm::DenseSet<int64_t> rolledGapIndices =
@@ -332,20 +413,21 @@ static FailureOr<std::string> emitSplitCtSlotIsl(
   emitAnd();
   bool firstTerm = true;
   os << "ct = ";
-  if (failed(emitSegmentAddress(os, firstTerm, pieces, pieceIndex, pieceDivBy,
-                                pieceModBy, traversalDims, gapDims,
-                                replicationDims, numTraversalComponents, 0,
-                                prefix, rolledGapIndices, rolls, rotomDims)))
+  if (failed(emitSegmentAddress(
+          os, firstTerm, pieces, pieceIndex, pieceDivBy, pieceModBy,
+          traversalDims, gapDims, replicationDims, numTraversalComponents, 0,
+          prefix, rolledGapIndices, rolls, rollSteps, rotomDims)))
     return failure();
   if (firstTerm) os << "0";
 
   emitAnd();
   firstTerm = true;
   os << "slot = ";
-  if (failed(emitSegmentAddress(
-          os, firstTerm, pieces, pieceIndex, pieceDivBy, pieceModBy,
-          traversalDims, gapDims, replicationDims, numTraversalComponents,
-          prefix, pieces.size(), rolledGapIndices, rolls, rotomDims)))
+  if (failed(emitSegmentAddress(os, firstTerm, pieces, pieceIndex, pieceDivBy,
+                                pieceModBy, traversalDims, gapDims,
+                                replicationDims, numTraversalComponents, prefix,
+                                pieces.size(), rolledGapIndices, rolls,
+                                rollSteps, rotomDims)))
     return failure();
   if (firstTerm) os << "0";
 
@@ -388,10 +470,14 @@ static FailureOr<std::string> lowerToIslImpl(LayoutAttr layout) {
   DenseI64ArrayAttr rollsAttr = layout.getRolls();
   ArrayRef<int64_t> rolls =
       rollsAttr ? rollsAttr.asArrayRef() : ArrayRef<int64_t>{};
+  DenseI64ArrayAttr stepsAttr = layout.getRollSteps();
+  ArrayRef<int64_t> rollSteps =
+      stepsAttr ? stepsAttr.asArrayRef() : ArrayRef<int64_t>{};
   return emitSplitCtSlotIsl(
       data.n, data.ctPrefixLen, data.pieces, data.pieceIndex, data.pieceDivBy,
       data.pieceModBy, data.traversalDims, data.replicationDims, data.gapDims,
-      numTraversalComponents, numReplication, numGap, rolls, layout.getDims());
+      numTraversalComponents, numReplication, numGap, rolls, rollSteps,
+      layout.getDims());
 }
 
 }  // namespace