Rotom: Layout Representation updates for tiled layouts

lib/Dialect/Rotom/IR/RotomAttributes.h

@@ -17,20 +17,55 @@ namespace mlir::heir::rotom {
 
 enum class LayoutPieceKind { Traversal, Replication, Gap };
 
+struct LayoutPiece {
+  DimAttr dim;
+  LayoutPieceKind kind;
+  // axisIndex, divBy, and modBy lower a traversal piece into its term of the
+  // ISL relation. The emitter builds an address `[i0, i1, ...] -> [ct, slot]`
+  // with one variable per axis (LayoutData::axes); each piece contributes a
+  // term reading one mixed-radix digit of its axis's variable.
+  //
+  // axisIndex picks the variable: an index into LayoutData::axes, emitted as
+  // `i{axisIndex}`.
+  int64_t axisIndex = -1;
+  // divBy and modBy pick which digit of that variable the piece reads, as
+  // (i / divBy) mod modBy. divBy is the digit's place value: the piece's
+  // stride when the axis is split across pieces, else 1. modBy is the digit's
+  // extent or 0 to drop the modulus on the most-significant digit.
+  int64_t divBy = 1;
+  int64_t modBy = 0;
+};
+
 struct LayoutData {
   int64_t n;
+  // Pieces [0, ctPrefixLen) are the ciphertext dimensions.
+  // Pieces [ctPrefixLen, pieces.size()) are the slot dimensions.
+  // The split is shown with the `|` separator.
   int64_t ctPrefixLen;
-  llvm::SmallVector<DimAttr> originalDims;
-  llvm::SmallVector<DimAttr> traversalDims;
-  llvm::SmallVector<DimAttr> replicationDims;
-  llvm::SmallVector<DimAttr> gapDims;
-  llvm::SmallVector<LayoutPieceKind> pieces;
-  llvm::SmallVector<int64_t> pieceIndex;
+  // Logical tensor axes.
+  llvm::SmallVector<DimAttr> axes;
+  llvm::SmallVector<LayoutPiece> pieces;
+
+  bool isCiphertextPiece(size_t p) const {
+    return static_cast<int64_t>(p) < ctPrefixLen;
+  }
 };
 
 /// Preprocess a Rotom layout.
 FailureOr<LayoutData> preprocessLayoutAttr(LayoutAttr attr);
 
+/// Canonicalizes raw layout pieces to the stored form. When the slot side
+/// (the longest dims suffix fitting `n`) underfills the ciphertext, it inserts
+/// the explicit front gap piece at the ct/slot boundary.
+void canonicalizeLayoutDims(MLIRContext* ctx, llvm::SmallVector<DimAttr>& dims,
+                            int64_t n, llvm::SmallVector<int64_t>& rolls);
+
+/// Computes how many leading entries of `dims` (read left-to-right) fall on the
+/// ciphertext axis for a ciphertext of `n` slots: the prefix that does not fit
+/// into the remaining slot budget. Shared so attribute preprocessing and the
+/// layout cost utilities derive the ct/slot split identically.
+size_t inferCtPrefixLen(llvm::ArrayRef<DimAttr> dims, int64_t n);
+
 }  // namespace mlir::heir::rotom
 
 #endif  // LIB_DIALECT_ROTOM_IR_ROTOMATTRIBUTES_H_

lib/Dialect/Rotom/IR/RotomAttributes.td

@@ -23,7 +23,9 @@ def Rotom_DimAttr : Rotom_Attr<"Dim", "dim"> {
     * `-1`: replication (logical fill along this layout piece)
     * `-2`: gap (padding / unused slots; constrained to zero in materialization)
 
-    Non-negative `dim` values index into the logical tensor shape.
+    Non-negative `dim` values index into the logical tensor shape. In the
+    assembly form the sentinels are spelled `R` (replication) and `G` (gap),
+    e.g. `[R:4:1]`; the numeric ids are also accepted on input.
   }];
 
   let parameters = (ins
@@ -47,13 +49,11 @@ def Rotom_LayoutAttr : Rotom_Attr<"Layout", "layout"> {
     A Rotom layout is an ordered list of `rotom.dim` dimensions plus the slot
     count `n` (ciphertext slot capacity).
 
-    The verifier determines which dims map across ciphertexts vs within slots
-    and checks the slot-side invariant that sizes and strides are powers of
-    two (after splitting).
-
-    For tensor_ext materialization, the **first** entry in `dims` is the
-    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).
+    Rotom's split between ciphertext dims and slot dims is denoted by a
+    `|` inside the dims list (e.g., `dims = [[0:2:4] | [0:4:1]]`).
+    Slot dims must fill `n` exactly; unused capacity is denoted
+    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
@@ -71,7 +71,11 @@ def Rotom_LayoutAttr : Rotom_Attr<"Layout", "layout"> {
   let hasCustomAssemblyFormat = 1;
 
   let extraClassDeclaration = [{
-    /// Layout with no `roll(i,j)` metadata (empty rolls storage).
+    /// Layout built from raw traversal dimensions.
+    static ::mlir::heir::rotom::LayoutAttr getCanonical(
+        ::mlir::MLIRContext *context, ::llvm::ArrayRef<DimAttr> dims,
+        int64_t n, ::llvm::ArrayRef<int64_t> rolls = {});
+    /// Canonicalizing builder with no `roll(from, by)` metadata.
     static ::mlir::heir::rotom::LayoutAttr get(::mlir::MLIRContext *context,
                                                ::mlir::ArrayAttr dims, int64_t n);
   }];