Rotom: axis roll arguments and per-piece roll semantics

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).

Author: edwjchen

lib/Dialect/Rotom/IR/RotomAttributes.cpp

@@ -227,25 +227,40 @@ static ParseResult parseLayoutDims(AsmParser& parser,
   }
 }
 
+// One argument of a roll pair: a bare non-negative integer is a dims-list
+// position (one piece); `axis N` names the whole tensor axis N.
+static ParseResult parseRollArg(AsmParser& parser, int64_t& encoded) {
+  const bool isAxis = succeeded(parser.parseOptionalKeyword("axis"));
+  int64_t value;
+  if (parser.parseInteger(value)) return failure();
+  if (value < 0) {
+    return parser.emitError(parser.getNameLoc())
+           << (isAxis ? "an axis roll argument must name a non-negative "
+                        "tensor axis"
+                      : "a piece roll argument must be a non-negative dims "
+                        "position (spell a whole-axis argument as `axis N`)");
+  }
+  encoded = encodeRollArg({isAxis, value});
+  return success();
+}
+
 static ParseResult parseLayoutRolls(AsmParser& parser,
                                     SmallVector<int64_t>& rolls) {
   if (parser.parseLSquare()) return failure();
   if (succeeded(parser.parseOptionalRSquare())) return success();
 
+  // Each entry is one `(from, by)` pair, parenthesized -- the form the
+  // printer emits, so written and round-tripped layouts read alike.
   while (true) {
-    if (succeeded(parser.parseOptionalLParen())) {
-      int64_t from;
-      int64_t to;
-      if (parser.parseInteger(from) || parser.parseComma() ||
-          parser.parseInteger(to) || parser.parseRParen())
-        return failure();
-      rolls.push_back(from);
-      rolls.push_back(to);
-    } else {
-      int64_t value;
-      if (parser.parseInteger(value)) return failure();
-      rolls.push_back(value);
+    int64_t from;
+    int64_t by;
+    if (parser.parseLParen() || failed(parseRollArg(parser, from)) ||
+        parser.parseComma() || failed(parseRollArg(parser, by)) ||
+        parser.parseRParen()) {
+      return failure();
     }
+    rolls.push_back(from);
+    rolls.push_back(by);
 
     if (succeeded(parser.parseOptionalComma())) continue;
     return parser.parseRSquare();
@@ -255,48 +270,109 @@ static ParseResult parseLayoutRolls(AsmParser& parser,
 static LogicalResult verifyLayoutRolls(
     ArrayAttr dims, DenseI64ArrayAttr rolls,
     function_ref<InFlightDiagnostic()> emitError) {
-  if (!rolls) return success();
+  const bool noRolls = !rolls || rolls.empty();
+  if (noRolls) return success();
   ArrayRef<int64_t> r = rolls.asArrayRef();
-  if (r.empty()) return success();
   if (r.size() % 2 != 0) {
-    return emitError() << "rolls must contain an even number of integers "
-                          "(pairs of dim indices)";
+    return emitError() << "rolls must contain an even number of arguments "
+                          "(pairs)";
   }
 
+  // Traversal pieces of a tensor axis: their count decides whether an `axis`
+  // argument is legal (split axes only -- the piece spelling is canonical
+  // when the axis is one piece), and their extent product is the modulus a
+  // whole-axis rewrite reduces by.
+  auto piecesOfAxis = [&](int64_t axis) {
+    std::pair<int64_t, int64_t> countAndExtent{0, 1};
+    for (Attribute a : dims) {
+      auto d = dyn_cast<DimAttr>(a);
+      if (d && !d.isGap() && !d.isReplicate() && d.getDim() == axis) {
+        ++countAndExtent.first;
+        countAndExtent.second *= d.getSize();
+      }
+    }
+    return countAndExtent;
+  };
+
   for (size_t i = 0; i < r.size(); i += 2) {
-    const int64_t ti = r[i];
-    const int64_t tj = r[i + 1];
-    if (ti == tj) {
-      return emitError() << "each roll must use two distinct dim indices";
+    const RollArg from = decodeRollArg(r[i]);
+    const RollArg by = decodeRollArg(r[i + 1]);
+
+    // Resolve each argument: the piece it names (null for axis arguments)
+    // and the tensor axis it reads or rewrites (sentinel for gap/replication
+    // pieces).
+    DimAttr fromPiece;
+    DimAttr byPiece;
+    int64_t fromAxis = 0;
+    int64_t byAxis = 0;
+    auto checkArg = [&](const RollArg& e, DimAttr& piece,
+                        int64_t& axis) -> LogicalResult {
+      if (e.isAxis) {
+        auto [count, extent] = piecesOfAxis(e.index);
+        (void)extent;
+        if (count == 0) {
+          return emitError() << "an axis roll argument must name a tensor "
+                                "axis present in dims";
+        }
+        if (count == 1) {
+          return emitError() << "an axis roll argument requires a split "
+                                "axis; spell an unsplit axis's argument as "
+                                "its piece position";
+        }
+        axis = e.index;
+        return success();
+      }
+      if (e.index >= static_cast<int64_t>(dims.size())) {
+        return emitError() << "roll piece argument out of bounds for dims "
+                              "list";
+      }
+      piece = dyn_cast<DimAttr>(dims[e.index]);
+      if (!piece) {
+        return emitError() << "roll arguments must refer to #rotom.dim "
+                              "entries";
+      }
+      axis = piece.getDim();
+      return success();
+    };
+    if (failed(checkArg(from, fromPiece, fromAxis)) ||
+        failed(checkArg(by, byPiece, byAxis))) {
+      return failure();
     }
-    if (ti < 0 || tj < 0 || ti >= static_cast<int64_t>(dims.size()) ||
-        tj >= static_cast<int64_t>(dims.size())) {
-      return emitError() << "roll dim index out of bounds for dims list";
+
+    // The extents need not match: a roll rewrites the from index to
+    // (idx - shift) mod extent(from), well-defined for any partner extent (a
+    // smaller partner covers a prefix of the rotations, a larger one wraps).
+    // FROM is the index expression being rewritten, so it must be a
+    // traversal piece or a whole (traversal) axis. The by argument may be
+    // any kind: rolling by a replication or gap piece shifts by that piece's
+    // block index, so each block holds a distinct cyclic rotation of the
+    // rolled index -- the layout materializes every rotation and alignment
+    // becomes block selection. (A rolled-by gap thus claims its blocks,
+    // unlike a plain gap.)
+    if (!from.isAxis && (fromPiece.isGap() || fromPiece.isReplicate())) {
+      return emitError() << "the rolled dim must be a traversal dim (dim >= "
+                            "0)";
     }
-    auto di = dyn_cast<DimAttr>(dims[ti]);
-    auto dj = dyn_cast<DimAttr>(dims[tj]);
-    if (!di || !dj) {
-      return emitError() << "roll indices must refer to #rotom.dim entries";
+    // A roll may not shift an index by itself. Piece arguments must be
+    // distinct positions (two pieces of one axis are distinct digits); an
+    // axis argument overlaps every argument on the same axis, because a
+    // whole-axis rewrite touches all of its digits.
+    if (!from.isAxis && !by.isAxis && from.index == by.index) {
+      return emitError() << "each roll must use two distinct arguments";
     }
-    // The extents need not match: roll(i, j) rewrites dims[i]'s index to
-    // (i_i - i_j) mod size(dims[i]), well-defined for any partner extent (a
-    // smaller partner covers a prefix of the rotations, a larger one wraps).
-    // The rolled (from) dim must be a traversal dim -- it is the index
-    // expression being rewritten. The roll-by (second) dim may be any kind:
-    // rolling by a replication or gap dim shifts by that dim's block index,
-    // so each block holds a distinct cyclic rotation of the rolled dim -- the
-    // layout materializes every rotation and alignment becomes block
-    // selection. (A rolled-by gap thus claims its blocks, unlike a plain gap.)
-    if (di.isGap() || di.isReplicate()) {
-      return emitError() << "the rolled dim must be a traversal dim (dim >= 0)";
+    if ((from.isAxis || by.isAxis) && fromAxis == byAxis) {
+      return emitError() << "a roll may not shift an axis by one of its own "
+                            "pieces";
     }
     // A rolled-by GAP claims one ciphertext block per gap index, each holding
-    // a distinct rotation of the rolled dim. If the gap is larger than the
-    // rolled dim's extent the rotations repeat (period = the from extent),
-    // claiming blocks the conversion/kernel accounting was never audited for.
+    // a distinct rotation of the rolled index. If the gap is larger than the
+    // rolled extent the rotations repeat (period = the from extent), claiming
+    // blocks the conversion/kernel accounting was never audited for.
     // (Replication partners of larger extent are intended -- replicate-then-
     // roll -- so only gaps are bounded.)
-    if (dj.isGap() && dj.getSize() > di.getSize()) {
+    const int64_t fromExtent =
+        from.isAxis ? piecesOfAxis(fromAxis).second : fromPiece.getSize();
+    if (byPiece && byPiece.isGap() && byPiece.getSize() > fromExtent) {
       return emitError() << "a rolled-by gap dim must not exceed the rolled "
                             "dim's extent";
     }
@@ -345,10 +421,19 @@ void LayoutAttr::print(AsmPrinter& printer) const {
   DenseI64ArrayAttr rolls = getRolls();
   if (rolls && !rolls.asArrayRef().empty()) {
     ArrayRef<int64_t> values = rolls.asArrayRef();
+    auto printArg = [&](int64_t encoded) {
+      const RollArg e = decodeRollArg(encoded);
+      if (e.isAxis) printer << "axis ";
+      printer << e.index;
+    };
     printer << ", rolls = [";
     for (size_t i = 0; i < values.size(); i += 2) {
       if (i != 0) printer << ", ";
-      printer << "(" << values[i] << ", " << values[i + 1] << ")";
+      printer << "(";
+      printArg(values[i]);
+      printer << ", ";
+      printArg(values[i + 1]);
+      printer << ")";
     }
     printer << "]";
   }
@@ -436,6 +521,18 @@ Attribute LayoutAttr::parse(AsmParser& parser, Type type) {
       ArrayAttr::get(context, dims), n, DenseI64ArrayAttr::get(context, rolls));
 }
 
+SmallVector<RollSpec> getRollSpecs(LayoutAttr layout) {
+  SmallVector<RollSpec> specs;
+  DenseI64ArrayAttr rolls = layout.getRolls();
+  if (!rolls) return specs;
+  ArrayRef<int64_t> r = rolls.asArrayRef();
+
+  for (size_t i = 0; i + 1 < r.size(); i += 2) {
+    specs.push_back({decodeRollArg(r[i]), decodeRollArg(r[i + 1])});
+  }
+  return specs;
+}
+
 FailureOr<LayoutData> preprocessLayoutAttr(LayoutAttr layout) {
   return preprocessLayoutData(layout.getDims(), layout.getN(),
                               layout.getContext());
@@ -452,7 +549,9 @@ LogicalResult LayoutAttr::verify(function_ref<InFlightDiagnostic()> emitError,
     return emitError() << "`dims` must be an array of `#rotom.dim<...>`";
   }
 
-  if (failed(verifyLayoutRolls(dims, rolls, emitError))) return failure();
+  if (failed(verifyLayoutRolls(dims, rolls, emitError))) {
+    return failure();
+  }
 
   SmallVector<DimAttr> dimVec;
   dimVec.reserve(dims.size());
@@ -508,7 +607,8 @@ void canonicalizeLayoutDims(MLIRContext* ctx, SmallVector<DimAttr>& dims,
   if (fill <= 1) return;
   dims.insert(dims.begin() + ctLen,
               DimAttr::get(ctx, /*dim=*/-2, fill, /*stride=*/1));
-  // Roll endpoints at or past the insertion shift right.
+  // Piece arguments at or past the insertion shift right; axis arguments
+  // (encoded negative) name axes and do not move.
   for (int64_t& encoded : rolls) {
     if (encoded >= static_cast<int64_t>(ctLen)) ++encoded;
   }

lib/Dialect/Rotom/IR/RotomAttributes.h

@@ -15,6 +15,37 @@
 
 namespace mlir::heir::rotom {
 
+// One argument 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).
+struct RollArg {
+  bool isAxis;
+  int64_t index;  // Dims-list position, or the tensor axis id when isAxis.
+  bool operator==(const RollArg& other) const {
+    return isAxis == other.isAxis && index == other.index;
+  }
+};
+
+// Argument encoding in the flat rolls storage: a piece argument is its
+// non-negative dims-list position, an axis argument is -(axis + 1).
+inline int64_t encodeRollArg(RollArg e) {
+  return e.isAxis ? -(e.index + 1) : e.index;
+}
+inline RollArg decodeRollArg(int64_t encoded) {
+  return encoded < 0 ? RollArg{true, -encoded - 1} : RollArg{false, encoded};
+}
+
+// One roll of a layout: FROM's index is rewritten to
+// (idx_from - 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.
+struct RollSpec {
+  RollArg from;
+  RollArg by;
+};
+
+// The layout's rolls with both arguments decoded.
+llvm::SmallVector<RollSpec> getRollSpecs(LayoutAttr layout);
+
 enum class LayoutPieceKind { Traversal, Replication, Gap };
 
 struct LayoutPiece {

lib/Dialect/Rotom/IR/RotomAttributes.td

@@ -55,11 +55,54 @@ def Rotom_LayoutAttr : Rotom_Attr<"Layout", "layout"> {
     with an explicit gap piece (e.g. `[G:4:1]`).
     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 argument 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).
+
+    ```mlir
+    // Halevi-Shoup diagonal of a 4x4 matrix: slot j of ciphertext i holds
+    // A[i, (j - i) mod 4]. Argument 1 (the axis-1 piece) is rolled by
+    // argument 0 (the axis-0 piece).
+    #diag = #rotom.layout<n = 4, rolls = [(1, 0)], dims = [[0:4:1] | [1:4:1]]>
+
+    // Axis 1 is split across two pieces, so a whole-axis roll spells its
+    // argument `axis 1` rather than a piece position.
+    #split = #rotom.layout<n = 16, rolls = [(axis 1, 2)],
+                           dims = [[1:4:4], [1:4:1] | [0:16:1]]>
+    ```
+
+    A piece FROM rewrites that piece's own mixed-radix digit in place,
+    `digit(from) <- (digit(from) - 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:
+
+    ```mlir
+    // Axis 0 has extent 4, split into digits [0:2:2] (high) and [0:2:1]
+    // (low), and is rolled by a replication argument of extent 2. Writing
+    // the index as (high, low), replica d holds:
+    //
+    //   axis roll, `(axis 0, 0)`:   index (i - d) mod 4, so replica 1 maps
+    //                               (0,0) <- (0,1) <- (1,0) <- (1,1): the
+    //                               subtraction borrows out of the low
+    //                               digit into the high one.
+    //   piece roll, `(1, 0)`:       only the high digit moves, so replica 1
+    //                               maps (0,0) <- (1,0) and (0,1) <- (1,1);
+    //                               the low digit never carries.
+    //
+    // Rolling BOTH pieces by the same argument is still not the axis roll:
+    // each digit wraps inside itself, so (0,0) <- (1,1), not (1,1) <- (0,0).
+    ```
+
+    The shift is the by argument'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.
   }];
 
   let parameters = (ins