From 056b09dc7e2cf4d7916414fcf14b76c5bd732330 Mon Sep 17 00:00:00 2001
From: Edward Chen <edwjchen@gmail.com>
Date: Wed, 8 Jul 2026 23:38:54 +0000
Subject: [PATCH] Rotom: mixed-radix digit model, canonical dims, ct/slot
 split, R/G spelling

Attribute-layer rework preparing the layout search for split and tiled
packings:

- preprocessLayoutData produces a LayoutPiece list (one per written dim,
  tagged Traversal/Replication/Gap) plus a deduped axes table -- one ISL
  domain variable per logical tensor dim. A traversal piece is a
  mixed-radix digit (i / divBy) mod modBy of its axis's index, so an axis
  split into several pieces (tiled or ct/slot-straddling layouts) shares
  one variable. Multi-piece axes are validated as complete mixed-radix
  decompositions. The ISL emitter reads the per-piece digit descriptors
  when building address terms.
- The axes table canonicalizes to ascending dim id. Consumers read it
  positionally as tensor dims (most importantly the ISL domain), so a
  layout whose pieces lead with a later dim (column-major) must not leak
  piece order into the domain: previously row- and column-major
  materialized identical relations, making conversions between them
  silently free.
- The dim assembly separates ciphertext dims from slot dims with `|`
  ([ct... | slot...], omitted when there are no ct dims). The split is
  derived (the longest dims suffix whose extents fill n) and the written
  `|` is validated against it, so a layout never reads differently than
  it packs. The slot side must fill n exactly; unused capacity is written
  as an explicit gap piece rather than synthesized.
- Replication and gap dims print as R and G ([R:4:1], [G:4:1]); the
  numeric ids -1/-2 are still accepted on input and round-trip to the
  letter forms.

Affected materialize/syntax/seed tests re-blessed.
---
 lib/Dialect/Rotom/IR/RotomAttributes.cpp      | 335 +++++++++++-------
 lib/Dialect/Rotom/IR/RotomAttributes.h        |  47 ++-
 lib/Dialect/Rotom/IR/RotomAttributes.td       |  22 +-
 .../Utils/RotomTensorExtLayoutLowering.cpp    | 211 ++++++-----
 .../RotomTensorExtLayoutLoweringTest.cpp      |  48 ++-
 tests/Dialect/Rotom/IR/layout.mlir            |  23 +-
 tests/Dialect/Rotom/IR/syntax.mlir            |  25 +-
 .../Rotom/Transforms/doctest_seed_layout.mlir |  16 +-
 .../Transforms/materialize_column_major.mlir  |   2 +-
 .../Transforms/materialize_ct_slot_split.mlir |  20 +-
 .../materialize_explicit_gap_dim.mlir         |   2 +-
 .../materialize_implicit_gap_dim.mlir         |   8 +-
 .../materialize_repeated_column_major.mlir    |   8 +-
 .../materialize_replication_dim.mlir          |   2 +-
 .../Transforms/materialize_slot_gap_dims.mlir |   4 +-
 .../materialize_tiled_duplicate_dim.mlir      |  20 +-
 .../Rotom/Transforms/materialize_vector.mlir  |   3 +-
 .../Dialect/Rotom/Transforms/seed_layout.mlir |  32 +-
 18 files changed, 503 insertions(+), 325 deletions(-)

diff --git a/lib/Dialect/Rotom/IR/RotomAttributes.cpp b/lib/Dialect/Rotom/IR/RotomAttributes.cpp
index 6703485d35..1f2a7423d9 100644
--- a/lib/Dialect/Rotom/IR/RotomAttributes.cpp
+++ b/lib/Dialect/Rotom/IR/RotomAttributes.cpp
@@ -2,9 +2,11 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <map>
 #include <utility>
 #include <vector>
 
+#include "llvm/include/llvm/ADT/DenseMap.h"           // from @llvm-project
 #include "llvm/include/llvm/ADT/DenseSet.h"           // from @llvm-project
 #include "llvm/include/llvm/ADT/STLExtras.h"          // from @llvm-project
 #include "llvm/include/llvm/ADT/SmallVector.h"        // from @llvm-project
@@ -20,9 +22,7 @@ namespace mlir {
 namespace heir {
 namespace rotom {
 
-namespace {
-
-static size_t inferCtPrefixLen(ArrayRef<DimAttr> dims, int64_t n) {
+size_t inferCtPrefixLen(ArrayRef<DimAttr> dims, int64_t n) {
   int64_t nRem = n;
   size_t i = dims.size();
   while (i > 0) {
@@ -41,76 +41,130 @@ static size_t inferCtPrefixLen(ArrayRef<DimAttr> dims, int64_t n) {
   return i;
 }
 
-static int64_t computeImplicitFrontGap(ArrayRef<DimAttr> dims, int64_t n) {
-  int64_t nRem = n;
-  for (auto it = dims.rbegin(); it != dims.rend(); ++it) {
-    DimAttr d = *it;
-    if (d.isGap()) return 1;
-    if (nRem <= 1) break;
-    const int64_t sz = d.getSize();
-    if (sz <= 0) return 1;
-    if (sz <= nRem && nRem % sz == 0) nRem /= sz;
-  }
-  return nRem;
-}
-
+// Preprocesses a layout (`dims`, slot count `n`) into the `LayoutData`
+// descriptor used to emit ciphertext addresses; also the validity check
+// behind `LayoutAttr::verify`.
+//
+// `pieces` describes the traversal dimensions, replication dimensions, and
+// gap dimensions. When lowering to ISL, a traversal piece maps to
+// `(i / divBy) mod modBy`, where i is the index of the piece's tensor axis;
+// replication and gap pieces map to their own ISL existential variables
+// instead.
 static FailureOr<LayoutData> preprocessLayoutData(ArrayAttr dims, int64_t n,
                                                   MLIRContext* ctx) {
   LayoutData data;
   data.n = n;
   if (data.n <= 0) return failure();
 
-  data.originalDims.reserve(dims.size());
+  std::map<int64_t, DimAttr> axisForDim;
+  SmallVector<DimAttr> writtenDims;
+  writtenDims.reserve(dims.size());
+  data.pieces.reserve(dims.size());
   for (Attribute a : dims) {
     auto d = dyn_cast<DimAttr>(a);
     if (!d) return failure();
-    data.originalDims.push_back(d);
+    writtenDims.push_back(d);
     if (d.isGap()) {
-      data.pieceIndex.push_back(static_cast<int64_t>(data.gapDims.size()));
-      data.gapDims.push_back(d);
-      data.pieces.push_back(LayoutPieceKind::Gap);
+      data.pieces.push_back({d, LayoutPieceKind::Gap});
       continue;
     }
     if (d.isReplicate()) {
-      data.pieceIndex.push_back(
-          static_cast<int64_t>(data.replicationDims.size()));
-      data.replicationDims.push_back(d);
-      data.pieces.push_back(LayoutPieceKind::Replication);
+      data.pieces.push_back({d, LayoutPieceKind::Replication});
       continue;
     }
-    if (d.getDim() >= 0) {
-      data.pieceIndex.push_back(
-          static_cast<int64_t>(data.traversalDims.size()));
-      data.traversalDims.push_back(d);
-      data.pieces.push_back(LayoutPieceKind::Traversal);
-      continue;
+    if (d.getDim() < 0) return failure();
+    axisForDim.try_emplace(d.getDim(), d);
+    // axisIndex is set below: a dim's rank isn't known until all dims are seen.
+    data.pieces.push_back({d, LayoutPieceKind::Traversal, /*axisIndex=*/-1,
+                           /*divBy=*/d.getStride()});
+  }
+
+  // Number the axes by ascending dim id (the order std::map gives).
+  llvm::DenseMap<int64_t, int64_t> axisIndexForDim;
+  for (auto& [dim, dimAttr] : axisForDim) {
+    axisIndexForDim[dim] = static_cast<int64_t>(data.axes.size());
+    data.axes.push_back(dimAttr);
+  }
+  for (LayoutPiece& piece : data.pieces) {
+    if (piece.kind == LayoutPieceKind::Traversal) {
+      piece.axisIndex = axisIndexForDim[piece.dim.getDim()];
     }
-    return failure();
   }
 
-  data.ctPrefixLen =
-      static_cast<int64_t>(inferCtPrefixLen(data.originalDims, data.n));
-
-  const int64_t implicitFrontGapSize =
-      computeImplicitFrontGap(data.originalDims, data.n);
-  if (implicitFrontGapSize > 1) {
-    const int64_t gapIdx = static_cast<int64_t>(data.gapDims.size());
-    data.gapDims.push_back(DimAttr::get(ctx, /*dim=*/-2,
-                                        /*size=*/implicitFrontGapSize,
-                                        /*stride=*/1));
-    data.pieces.insert(data.pieces.begin() + data.ctPrefixLen,
-                       LayoutPieceKind::Gap);
-    data.pieceIndex.insert(data.pieceIndex.begin() + data.ctPrefixLen, gapIdx);
+  // Count pieces per tensor dim and the dim's full extent.
+  llvm::DenseMap<int64_t, int64_t> pieceCount;
+  llvm::DenseMap<int64_t, int64_t> dimFullExtent;
+  for (const LayoutPiece& piece : data.pieces) {
+    if (piece.kind != LayoutPieceKind::Traversal) continue;
+    ++pieceCount[piece.dim.getDim()];
+    auto [it, inserted] = dimFullExtent.try_emplace(piece.dim.getDim(), 1);
+    it->second *= piece.dim.getSize();
   }
 
+  for (LayoutPiece& piece : data.pieces) {
+    if (piece.kind != LayoutPieceKind::Traversal) continue;
+    if (pieceCount[piece.dim.getDim()] == 1) {
+      // Lone piece: (i / 1) mod 0 = i.
+      piece.divBy = 1;
+      piece.modBy = 0;
+    } else {
+      // Split piece: digit = (i / stride) mod extent.
+      const int64_t extent = piece.dim.getSize();
+      const int64_t full = dimFullExtent[piece.dim.getDim()];
+      piece.modBy = (piece.divBy * extent < full) ? extent : 0;
+    }
+  }
+
+  // Each multi-piece tensor dim must be a valid mixed-radix decomposition:
+  // sorted by stride, the divisors are the cumulative products of the lower
+  // extents (1, e0, e0*e1, ...), and the extents multiply to the full extent.
+  for (auto& [dim, count] : pieceCount) {
+    if (count == 1) continue;
+    SmallVector<std::pair<int64_t, int64_t>> parts;  // (stride, extent)
+    for (const LayoutPiece& piece : data.pieces) {
+      if (piece.kind == LayoutPieceKind::Traversal &&
+          piece.dim.getDim() == dim) {
+        parts.push_back({piece.divBy, piece.dim.getSize()});
+      }
+    }
+    llvm::sort(parts);  // ascending stride
+    int64_t expected = 1;
+    for (auto [stride, extent] : parts) {
+      if (extent <= 0 || stride != expected) return failure();
+      expected *= extent;
+    }
+    if (expected != dimFullExtent[dim]) return failure();
+  }
+
+  for (size_t ti = 0; ti < data.axes.size(); ++ti) {
+    const int64_t dim = data.axes[ti].getDim();
+    data.axes[ti] = DimAttr::get(ctx, dim, dimFullExtent[dim], /*stride=*/1);
+  }
+
+  data.ctPrefixLen =
+      static_cast<int64_t>(inferCtPrefixLen(writtenDims, data.n));
+
   return data;
 }
 
+// The dim position is spelled `R` for replication and `G` for gap (the
+// readable forms, also how the printer emits them); the numeric ids -1 and
+// -2 are still accepted.
 static ParseResult parseDimTripleAfterLSquare(AsmParser& parser, int64_t& dim,
                                               int64_t& size, int64_t& stride) {
-  return failure(parser.parseInteger(dim) || parser.parseColon() ||
-                 parser.parseInteger(size) || parser.parseColon() ||
-                 parser.parseInteger(stride) || parser.parseRSquare());
+  if (succeeded(parser.parseOptionalKeyword("R"))) {
+    dim = -1;
+  } else if (succeeded(parser.parseOptionalKeyword("G"))) {
+    dim = -2;
+  } else if (parser.parseInteger(dim)) {
+    return failure();
+  }
+  if (parser.parseColon() || parser.parseInteger(size)) return failure();
+  stride = 1;
+  if (succeeded(parser.parseOptionalColon()) && parser.parseInteger(stride)) {
+    return failure();
+  }
+  return parser.parseRSquare();
 }
 
 static ParseResult parseDimTriple(AsmParser& parser, int64_t& dim,
@@ -120,12 +174,24 @@ static ParseResult parseDimTriple(AsmParser& parser, int64_t& dim,
 }
 
 static void printDimTriple(AsmPrinter& printer, DimAttr dim) {
-  printer << "[" << dim.getDim() << ":" << dim.getSize() << ":"
-          << dim.getStride() << "]";
+  printer << "[";
+  if (dim.isReplicate()) {
+    printer << "R";
+  } else if (dim.isGap()) {
+    printer << "G";
+  } else {
+    printer << dim.getDim();
+  }
+  printer << ":" << dim.getSize() << ":" << dim.getStride() << "]";
 }
 
+// Parses `[piece, ... | piece, ...]`: the `|` separates the ciphertext dims
+// from the slot dims (absent when there are no ciphertext dims). The written
+// boundary is returned in `writtenCtLen` for validation against the derived
+// split.
 static ParseResult parseLayoutDims(AsmParser& parser,
-                                   SmallVector<Attribute>& dims) {
+                                   SmallVector<Attribute>& dims,
+                                   std::optional<int64_t>& writtenCtLen) {
   if (parser.parseLSquare()) return failure();
   if (succeeded(parser.parseOptionalRSquare())) return success();
 
@@ -148,6 +214,15 @@ static ParseResult parseLayoutDims(AsmParser& parser,
     }
 
     if (succeeded(parser.parseOptionalComma())) continue;
+    if (succeeded(parser.parseOptionalVerticalBar())) {
+      if (writtenCtLen.has_value()) {
+        return parser.emitError(parser.getNameLoc())
+               << "at most one `|` may separate ciphertext dims from slot "
+                  "dims";
+      }
+      writtenCtLen = static_cast<int64_t>(dims.size());
+      continue;
+    }
     return parser.parseRSquare();
   }
 }
@@ -177,8 +252,6 @@ static ParseResult parseLayoutRolls(AsmParser& parser,
   }
 }
 
-}  // namespace
-
 static LogicalResult verifyLayoutRolls(
     ArrayAttr dims, DenseI64ArrayAttr rolls,
     function_ref<InFlightDiagnostic()> emitError) {
@@ -205,6 +278,9 @@ static LogicalResult verifyLayoutRolls(
     if (!di || !dj) {
       return emitError() << "roll indices must refer to #rotom.dim entries";
     }
+    // The two dims of a roll must have equal extents: roll(i, j) shifts dims[i]
+    // by dims[j]'s index modulo their shared extent, which is well-defined only
+    // when the extents match.
     if (di.getSize() != dj.getSize()) {
       return emitError() << "rolled dims must have the same extent (size)";
     }
@@ -265,10 +341,16 @@ void LayoutAttr::print(AsmPrinter& printer) const {
     printer << "]";
   }
 
+  SmallVector<DimAttr> dimVec;
+  dimVec.reserve(getDims().size());
+  for (Attribute attr : getDims()) dimVec.push_back(cast<DimAttr>(attr));
+  const size_t ctLen = inferCtPrefixLen(dimVec, getN());
+
   printer << ", dims = [";
-  llvm::interleaveComma(getDims(), printer, [&](Attribute attr) {
-    printDimTriple(printer, cast<DimAttr>(attr));
-  });
+  for (size_t i = 0; i < dimVec.size(); ++i) {
+    if (i > 0) printer << (i == ctLen ? " | " : ", ");
+    printDimTriple(printer, dimVec[i]);
+  }
   printer << "]>";
 }
 
@@ -276,6 +358,7 @@ Attribute LayoutAttr::parse(AsmParser& parser, Type type) {
   int64_t n;
   SmallVector<int64_t> rolls;
   SmallVector<Attribute> dims;
+  std::optional<int64_t> writtenCtLen;
 
   if (parser.parseLess()) return {};
 
@@ -292,11 +375,13 @@ Attribute LayoutAttr::parse(AsmParser& parser, Type type) {
     }
 
     if (parser.parseKeyword("dims") || parser.parseEqual() ||
-        failed(parseLayoutDims(parser, dims)) || parser.parseGreater()) {
+        failed(parseLayoutDims(parser, dims, writtenCtLen)) ||
+        parser.parseGreater()) {
       return {};
     }
   } else if (succeeded(parser.parseOptionalKeyword("dims"))) {
-    if (parser.parseEqual() || failed(parseLayoutDims(parser, dims)) ||
+    if (parser.parseEqual() ||
+        failed(parseLayoutDims(parser, dims, writtenCtLen)) ||
         parser.parseComma() || parser.parseKeyword("n") ||
         parser.parseEqual() || parser.parseInteger(n)) {
       return {};
@@ -316,6 +401,23 @@ Attribute LayoutAttr::parse(AsmParser& parser, Type type) {
     return {};
   }
 
+  // Verify the written `|` boundary.
+  SmallVector<DimAttr> dimVec;
+  dimVec.reserve(dims.size());
+  for (Attribute attr : dims) dimVec.push_back(cast<DimAttr>(attr));
+  const int64_t derivedCtLen =
+      static_cast<int64_t>(inferCtPrefixLen(dimVec, n));
+  if (n > 0 && writtenCtLen.value_or(0) != derivedCtLen) {
+    parser.emitError(parser.getNameLoc())
+        << "the written `|` ciphertext/slot split (" << writtenCtLen.value_or(0)
+        << " ciphertext dims) does not match the derived split ("
+        << derivedCtLen
+        << "): the slot side is the longest dims suffix whose extents fit "
+           "n = "
+        << n;
+    return {};
+  }
+
   MLIRContext* context = parser.getContext();
   return LayoutAttr::getChecked(
       [&]() { return parser.emitError(parser.getNameLoc()); }, context,
@@ -340,73 +442,13 @@ LogicalResult LayoutAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 
   if (failed(verifyLayoutRolls(dims, rolls, emitError))) return failure();
 
-  MLIRContext* ctx = dims.getContext();
-  std::vector<DimAttr> ctDims;
-  std::vector<DimAttr> slotDims;
-
-  int64_t nRem = n;
-  for (auto it = preprocessed->originalDims.rbegin();
-       it != preprocessed->originalDims.rend(); ++it) {
-    DimAttr d = *it;
-    const int64_t size = d.getSize();
-
-    if (nRem <= 1) {
-      ctDims.insert(ctDims.begin(), d);
-      continue;
-    }
-
-    // Size > nRem: split into ct and slot dims
-    if (size > nRem) {
-      if (size % nRem != 0) {
-        return emitError() << "dim size " << size
-                           << " must be divisible by remaining slot capacity "
-                           << nRem;
-      }
-
-      slotDims.insert(slotDims.begin(),
-                      DimAttr::get(ctx, d.getDim(), nRem, /*stride=*/1));
-      ctDims.insert(ctDims.begin(), DimAttr::get(ctx, d.getDim(), size / nRem,
-                                                 /*stride=*/nRem));
-      nRem /= size;
-      continue;
-    }
-
-    // Size == nRem: add to slot dims
-    if (size == nRem) {
-      slotDims.insert(slotDims.begin(), d);
-      nRem /= size;
-      continue;
-    }
-
-    // Size divides nRem: add to slot dims
-    if (nRem % size == 0) {
-      slotDims.insert(slotDims.begin(), d);
-      nRem /= size;
-      continue;
-    }
-
-    // Size does not divide nRem (e.g. odd input channels): keep in ctDims
-    // so slotDims can remain pow2 and nRem is unchanged.
-    ctDims.insert(ctDims.begin(), d);
-  }
-
-  // If there is remaining slot capacity, insert a gap dim at the front.
-  if (nRem > 1) {
-    slotDims.insert(slotDims.begin(),
-                    DimAttr::get(ctx, /*dim=*/-2, /*size=*/nRem,
-                                 /*stride=*/1));
-  }
-
-  // Remove gap dims from ctDims.
-  std::vector<DimAttr> ctDimsFiltered;
-  ctDimsFiltered.reserve(ctDims.size());
-  for (DimAttr d : ctDims) {
-    if (!d.isGap()) ctDimsFiltered.push_back(d);
-  }
-  ctDims.swap(ctDimsFiltered);
-
-  // Enforce the Rotom invariant: slot-dim sizes/strides must be powers of 2.
-  for (DimAttr d : slotDims) {
+  SmallVector<DimAttr> dimVec;
+  dimVec.reserve(dims.size());
+  for (Attribute attr : dims) dimVec.push_back(cast<DimAttr>(attr));
+  const size_t ctLen = inferCtPrefixLen(dimVec, n);
+  int64_t slotExtent = 1;
+  for (size_t p = ctLen; p < dimVec.size(); ++p) {
+    DimAttr d = dimVec[p];
     if (!llvm::isPowerOf2_64(static_cast<uint64_t>(d.getSize()))) {
       return emitError() << "slot dim size must be a power of two, got "
                          << d.getSize();
@@ -415,6 +457,16 @@ LogicalResult LayoutAttr::verify(function_ref<InFlightDiagnostic()> emitError,
       return emitError() << "slot dim stride must be a power of two, got "
                          << d.getStride();
     }
+    slotExtent *= d.getSize();
+  }
+
+  // The slot side must fill the ciphertext exactly. Unused slots must be
+  // represented as an explicit gap piece.
+  if (slotExtent != n) {
+    return emitError() << "slot dims must fill the ciphertext exactly (slot "
+                          "extent "
+                       << slotExtent << " vs n = " << n
+                       << "); spell unused capacity as an explicit gap piece";
   }
 
   return success();
@@ -434,9 +486,38 @@ LogicalResult SeedAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+void canonicalizeLayoutDims(MLIRContext* ctx, SmallVector<DimAttr>& dims,
+                            int64_t n, SmallVector<int64_t>& rolls) {
+  const size_t ctLen = inferCtPrefixLen(dims, n);
+  int64_t slotExtent = 1;
+  for (size_t p = ctLen; p < dims.size(); ++p) slotExtent *= dims[p].getSize();
+  if (slotExtent <= 0 || n % slotExtent != 0) return;
+  const int64_t fill = n / slotExtent;
+  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.
+  for (int64_t& encoded : rolls) {
+    if (encoded >= static_cast<int64_t>(ctLen)) ++encoded;
+  }
+}
+
+LayoutAttr LayoutAttr::getCanonical(MLIRContext* context,
+                                    ArrayRef<DimAttr> dims, int64_t n,
+                                    ArrayRef<int64_t> rolls) {
+  SmallVector<DimAttr> dimVec(dims.begin(), dims.end());
+  SmallVector<int64_t> rollVec(rolls.begin(), rolls.end());
+  canonicalizeLayoutDims(context, dimVec, n, rollVec);
+  SmallVector<Attribute> attrs(dimVec.begin(), dimVec.end());
+  return get(context, ArrayAttr::get(context, attrs), n,
+             DenseI64ArrayAttr::get(context, rollVec));
+}
+
 LayoutAttr LayoutAttr::get(MLIRContext* context, ArrayAttr dims, int64_t n) {
-  return get(context, dims, n,
-             DenseI64ArrayAttr::get(context, ArrayRef<int64_t>{}));
+  SmallVector<DimAttr> dimVec;
+  dimVec.reserve(dims.size());
+  for (Attribute attr : dims) dimVec.push_back(cast<DimAttr>(attr));
+  return getCanonical(context, dimVec, n);
 }
 
 }  // namespace rotom
diff --git a/lib/Dialect/Rotom/IR/RotomAttributes.h b/lib/Dialect/Rotom/IR/RotomAttributes.h
index a6ad4141b0..147db1ed2d 100644
--- a/lib/Dialect/Rotom/IR/RotomAttributes.h
+++ b/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_
diff --git a/lib/Dialect/Rotom/IR/RotomAttributes.td b/lib/Dialect/Rotom/IR/RotomAttributes.td
index b547d0ab07..0fab9c1a92 100644
--- a/lib/Dialect/Rotom/IR/RotomAttributes.td
+++ b/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);
   }];
diff --git a/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLowering.cpp b/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLowering.cpp
index bdc78a0528..32f2bbd668 100644
--- a/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLowering.cpp
+++ b/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLowering.cpp
@@ -20,12 +20,12 @@ 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) {
-  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()) {
+static FailureOr<int64_t> varIndexForRotomDim(const SmallVector<DimAttr>& axes,
+                                              DimAttr want) {
+  for (int64_t i = 0; i < static_cast<int64_t>(axes.size()); ++i) {
+    if (axes[i].getDim() == want.getDim() &&
+        axes[i].getSize() == want.getSize() &&
+        axes[i].getStride() == want.getStride()) {
       return i;
     }
   }
@@ -40,32 +40,36 @@ static std::string modExpr(llvm::StringRef expr, int64_t mod) {
   return out;
 }
 
-static LogicalResult emitSegmentAddress(
-    llvm::raw_ostream& os, bool& firstTerm, ArrayRef<LayoutPieceKind> pieces,
-    ArrayRef<int64_t> pieceIndex, const SmallVector<DimAttr>& traversalDims,
-    const SmallVector<DimAttr>& gapDims,
-    const SmallVector<DimAttr>& replicationDims,
-    int64_t numActiveTraversalComponents, size_t segStart, size_t segEnd,
-    bool foldGapVarsToZero, ArrayRef<int64_t> rolls, ArrayAttr rotomDims,
-    bool isSlotLine) {
+static std::string floorDivExpr(llvm::StringRef expr, int64_t d) {
+  std::string out;
+  llvm::raw_string_ostream os(out);
+  os << "floor((" << expr << ") / " << d << ")";
+  return out;
+}
+
+/// Ordinal of pieces[p] among the pieces of its kind.
+/// Used to name the existential variables for replication and gap pieces.
+static int64_t kindOrdinal(ArrayRef<LayoutPiece> pieces, size_t p) {
+  int64_t ordinal = 0;
+  for (size_t q = 0; q < p; ++q) {
+    if (pieces[q].kind == pieces[p].kind) ++ordinal;
+  }
+  return ordinal;
+}
+
+static LogicalResult emitSegmentAddress(llvm::raw_ostream& os, bool& firstTerm,
+                                        ArrayRef<LayoutPiece> pieces,
+                                        const SmallVector<DimAttr>& axes,
+                                        int64_t numAxes, size_t segStart,
+                                        size_t segEnd, bool foldGapVarsToZero,
+                                        ArrayRef<int64_t> rolls,
+                                        ArrayAttr rotomDims) {
   llvm::SmallVector<int64_t> suffixCoeff(pieces.size(), 0);
   int64_t suffix = 1;
   for (size_t p = segEnd; p > segStart;) {
     --p;
     suffixCoeff[p] = suffix;
-    DimAttr d;
-    switch (pieces[p]) {
-      case LayoutPieceKind::Traversal:
-        d = traversalDims[pieceIndex[p]];
-        break;
-      case LayoutPieceKind::Gap:
-        d = gapDims[pieceIndex[p]];
-        break;
-      case LayoutPieceKind::Replication:
-        d = replicationDims[pieceIndex[p]];
-        break;
-    }
-    suffix *= d.getSize();
+    suffix *= pieces[p].dim.getSize();
   }
 
   auto emitTerm = [&](int64_t coeff, llvm::StringRef expr) -> LogicalResult {
@@ -89,31 +93,48 @@ static LogicalResult emitSegmentAddress(
   llvm::DenseMap<int64_t, int64_t> replicationCoeff;
   for (size_t p = segStart; p < segEnd; ++p) {
     const int64_t coeff = suffixCoeff[p];
-    if (pieces[p] == LayoutPieceKind::Gap) {
+    if (pieces[p].kind == LayoutPieceKind::Gap) {
       if (foldGapVarsToZero) continue;
-      gapCoeff[pieceIndex[p]] = coeff;
-    } else if (pieces[p] == LayoutPieceKind::Replication) {
-      replicationCoeff[pieceIndex[p]] = coeff;
+      gapCoeff[kindOrdinal(pieces, p)] = coeff;
+    } else if (pieces[p].kind == LayoutPieceKind::Replication) {
+      replicationCoeff[kindOrdinal(pieces, p)] = coeff;
     }
   }
 
-  llvm::DenseMap<int64_t, int64_t> traversalCoeff;
+  // A tensor dim can contribute several pieces to one segment (a mixed-radix
+  // split places more than one digit of the same index on the same axis), so
+  // collect a list of (coeff, digit descriptor) per dim rather than one entry.
+  struct AxisDigit {
+    int64_t coeff;
+    int64_t divBy;
+    int64_t modBy;
+  };
+  llvm::DenseMap<int64_t, llvm::SmallVector<AxisDigit>> digitsByAxis;
   for (size_t p = segStart; p < segEnd; ++p) {
-    if (pieces[p] != LayoutPieceKind::Traversal) continue;
-    const int64_t ti = pieceIndex[p];
-    if (traversalDims[ti].getSize() == 1) continue;
-    traversalCoeff[ti] = suffixCoeff[p];
+    if (pieces[p].kind != LayoutPieceKind::Traversal) continue;
+    const int64_t ti = pieces[p].axisIndex;
+    if (axes[ti].getSize() == 1) continue;
+    digitsByAxis[ti].push_back(
+        {suffixCoeff[p], pieces[p].divBy, pieces[p].modBy});
   }
 
-  llvm::SmallVector<std::string> traversalExprs;
-  traversalExprs.reserve(traversalDims.size());
-  for (int64_t i = 0; i < static_cast<int64_t>(traversalDims.size()); ++i) {
-    traversalExprs.push_back("i" + std::to_string(i));
+  llvm::SmallVector<std::string> axisExprs;
+  axisExprs.reserve(axes.size());
+  for (int64_t i = 0; i < static_cast<int64_t>(axes.size()); ++i) {
+    axisExprs.push_back("i" + std::to_string(i));
   }
 
   // Apply roll(a,b) transforms left-to-right:
   // t_a <- (t_a - t_b) mod extent(a).
-  if (isSlotLine && !rolls.empty()) {
+  //
+  // 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, 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 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];
@@ -126,39 +147,41 @@ static LogicalResult emitSegmentAddress(
       auto fromDim = dyn_cast<DimAttr>(rotomDims[fromIdx]);
       auto toDim = dyn_cast<DimAttr>(rotomDims[toIdx]);
       if (!fromDim || !toDim) return failure();
-      FailureOr<int64_t> maybeFromTrav =
-          traversalIndexForRotomDim(traversalDims, fromDim);
-      FailureOr<int64_t> maybeToTrav =
-          traversalIndexForRotomDim(traversalDims, toDim);
-      if (failed(maybeFromTrav) || failed(maybeToTrav)) return failure();
-      const int64_t fromTrav = *maybeFromTrav;
-      const int64_t toTrav = *maybeToTrav;
+      FailureOr<int64_t> maybeFromVar = varIndexForRotomDim(axes, fromDim);
+      FailureOr<int64_t> maybeToVar = varIndexForRotomDim(axes, toDim);
+      if (failed(maybeFromVar) || failed(maybeToVar)) return failure();
+      const int64_t fromVar = *maybeFromVar;
+      const int64_t toVar = *maybeToVar;
       std::string diffExpr =
-          "(" + traversalExprs[fromTrav] + " - " + traversalExprs[toTrav] + ")";
-      traversalExprs[fromTrav] = modExpr(diffExpr, fromDim.getSize());
+          "(" + axisExprs[fromVar] + " - " + axisExprs[toVar] + ")";
+      axisExprs[fromVar] = modExpr(diffExpr, fromDim.getSize());
     }
   }
 
-  for (int64_t oldIdx = 0; oldIdx < static_cast<int64_t>(traversalDims.size());
+  for (int64_t oldIdx = 0; oldIdx < static_cast<int64_t>(axes.size());
        ++oldIdx) {
-    if (traversalDims[oldIdx].getSize() == 1) continue;
-    auto it = traversalCoeff.find(oldIdx);
-    if (it != traversalCoeff.end()) {
-      if (failed(emitTerm(it->second, traversalExprs[oldIdx])))
-        return failure();
+    if (axes[oldIdx].getSize() == 1) continue;
+    auto it = digitsByAxis.find(oldIdx);
+    if (it == digitsByAxis.end()) continue;
+    for (const AxisDigit& tp : it->second) {
+      // Mixed-radix digit: (i / divBy) mod modBy (modBy 0 => no modulus). A
+      // whole-dim piece (divBy 1, modBy 0) leaves the index untouched.
+      std::string expr = axisExprs[oldIdx];
+      if (tp.divBy > 1) expr = floorDivExpr(expr, tp.divBy);
+      if (tp.modBy > 0) expr = modExpr(expr, tp.modBy);
+      if (failed(emitTerm(tp.coeff, expr))) return failure();
     }
   }
-  for (int64_t g = 0; g < static_cast<int64_t>(gapDims.size()); ++g) {
+  for (int64_t g = 0; g < static_cast<int64_t>(pieces.size()); ++g) {
     auto it = gapCoeff.find(g);
     if (it != gapCoeff.end() &&
         failed(emitTerm(it->second, "g" + std::to_string(g))))
       return failure();
   }
-  for (int64_t e = 0; e < static_cast<int64_t>(replicationDims.size()); ++e) {
+  for (int64_t e = 0; e < static_cast<int64_t>(pieces.size()); ++e) {
     auto it = replicationCoeff.find(e);
     if (it != replicationCoeff.end()) {
-      const auto varName =
-          "d" + std::to_string(numActiveTraversalComponents + e);
+      const auto varName = "d" + std::to_string(numAxes + e);
       if (failed(emitTerm(it->second, varName))) return failure();
     }
   }
@@ -166,20 +189,24 @@ static LogicalResult emitSegmentAddress(
 }
 
 static FailureOr<std::string> emitSplitCtSlotIsl(
-    int64_t n, size_t prefix, ArrayRef<LayoutPieceKind> pieces,
-    ArrayRef<int64_t> pieceIndex, const SmallVector<DimAttr>& traversalDims,
-    const SmallVector<DimAttr>& replicationDims,
-    const SmallVector<DimAttr>& gapDims, int64_t numTraversalComponents,
-    int64_t numReplication, int64_t numGap, ArrayRef<int64_t> rolls,
+    int64_t n, size_t prefix, ArrayRef<LayoutPiece> pieces,
+    const SmallVector<DimAttr>& axes, ArrayRef<int64_t> rolls,
     ArrayAttr rotomDims) {
   if (prefix > pieces.size()) return failure();
 
+  const int64_t numAxes = static_cast<int64_t>(axes.size());
+  int64_t numReplication = 0;
+  int64_t numGap = 0;
+  for (const LayoutPiece& piece : pieces) {
+    if (piece.kind == LayoutPieceKind::Replication) ++numReplication;
+    if (piece.kind == LayoutPieceKind::Gap) ++numGap;
+  }
+
   int64_t numCt = 1;
   for (size_t p = 0; p < prefix; ++p) {
-    if (pieces[p] == LayoutPieceKind::Traversal) {
-      numCt *= traversalDims[pieceIndex[p]].getSize();
-    } else if (pieces[p] == LayoutPieceKind::Replication) {
-      numCt *= replicationDims[pieceIndex[p]].getSize();
+    if (pieces[p].kind == LayoutPieceKind::Traversal ||
+        pieces[p].kind == LayoutPieceKind::Replication) {
+      numCt *= pieces[p].dim.getSize();
     }
   }
   if (numCt < 1) numCt = 1;
@@ -188,7 +215,7 @@ static FailureOr<std::string> emitSplitCtSlotIsl(
   llvm::raw_string_ostream os(s);
 
   os << "{ [";
-  for (int64_t i = 0; i < numTraversalComponents; ++i) {
+  for (int64_t i = 0; i < numAxes; ++i) {
     if (i) os << ", ";
     os << "i" << i;
   }
@@ -200,8 +227,8 @@ static FailureOr<std::string> emitSplitCtSlotIsl(
     first = false;
   };
 
-  for (int64_t i = 0; i < numTraversalComponents; ++i) {
-    const DimAttr d = traversalDims[i];
+  for (int64_t i = 0; i < numAxes; ++i) {
+    const DimAttr d = axes[i];
     emitAnd();
     os << "0 <= i" << i << " <= " << (d.getSize() - 1);
   }
@@ -218,10 +245,10 @@ static FailureOr<std::string> emitSplitCtSlotIsl(
     os << " and exists ";
     bool firstVar = true;
     for (size_t p = 0; p < pieces.size(); ++p) {
-      if (pieces[p] == LayoutPieceKind::Replication) {
+      if (pieces[p].kind == LayoutPieceKind::Replication) {
         if (!firstVar) os << ", ";
         firstVar = false;
-        os << "d" << (numTraversalComponents + pieceIndex[p]);
+        os << "d" << (numAxes + kindOrdinal(pieces, p));
       }
     }
     os << " : ";
@@ -231,31 +258,26 @@ static FailureOr<std::string> emitSplitCtSlotIsl(
   emitAnd();
   bool firstTerm = true;
   os << "ct = ";
-  if (failed(emitSegmentAddress(os, firstTerm, pieces, pieceIndex,
-                                traversalDims, gapDims, replicationDims,
-                                numTraversalComponents, 0, prefix,
-                                foldGapVarsToZero, rolls, rotomDims,
-                                /*isSlotLine=*/false)))
+  if (failed(emitSegmentAddress(os, firstTerm, pieces, axes, numAxes, 0, prefix,
+                                foldGapVarsToZero, rolls, rotomDims)))
     return failure();
   if (firstTerm) os << "0";
 
   emitAnd();
   firstTerm = true;
   os << "slot = ";
-  if (failed(emitSegmentAddress(os, firstTerm, pieces, pieceIndex,
-                                traversalDims, gapDims, replicationDims,
-                                numTraversalComponents, prefix, pieces.size(),
-                                foldGapVarsToZero, rolls, rotomDims,
-                                /*isSlotLine=*/true)))
+  if (failed(emitSegmentAddress(os, firstTerm, pieces, axes, numAxes, prefix,
+                                pieces.size(), foldGapVarsToZero, rolls,
+                                rotomDims)))
     return failure();
   if (firstTerm) os << "0";
 
   if (numLocalVars > 0) {
-    for (int64_t k = 0; k < numReplication; ++k) {
-      const DimAttr d = replicationDims[k];
+    for (size_t p = 0; p < pieces.size(); ++p) {
+      if (pieces[p].kind != LayoutPieceKind::Replication) continue;
       emitAnd();
-      os << "0 <= d" << (numTraversalComponents + k)
-         << " <= " << (d.getSize() - 1);
+      os << "0 <= d" << (numAxes + kindOrdinal(pieces, p))
+         << " <= " << (pieces[p].dim.getSize() - 1);
     }
   }
 
@@ -270,24 +292,17 @@ static FailureOr<std::string> lowerToIslImpl(LayoutAttr layout) {
   const LayoutData& data = *maybeData;
 
   llvm::DenseMap<int64_t, llvm::DenseSet<int64_t>> seenDimStride;
-  for (DimAttr d : data.traversalDims) {
+  for (DimAttr d : data.axes) {
     auto& seenStridesForDim = seenDimStride[d.getDim()];
     if (seenStridesForDim.contains(d.getStride())) return failure();
     seenStridesForDim.insert(d.getStride());
   }
 
-  const int64_t numTraversalComponents =
-      static_cast<int64_t>(data.traversalDims.size());
-  const int64_t numReplication =
-      static_cast<int64_t>(data.replicationDims.size());
-  const int64_t numGap = static_cast<int64_t>(data.gapDims.size());
   DenseI64ArrayAttr rollsAttr = layout.getRolls();
   ArrayRef<int64_t> rolls =
       rollsAttr ? rollsAttr.asArrayRef() : ArrayRef<int64_t>{};
-  return emitSplitCtSlotIsl(
-      data.n, data.ctPrefixLen, data.pieces, data.pieceIndex,
-      data.traversalDims, data.replicationDims, data.gapDims,
-      numTraversalComponents, numReplication, numGap, rolls, layout.getDims());
+  return emitSplitCtSlotIsl(data.n, data.ctPrefixLen, data.pieces, data.axes,
+                            rolls, layout.getDims());
 }
 
 }  // namespace
diff --git a/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLoweringTest.cpp b/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLoweringTest.cpp
index 0b70e2f6ab..4065c241d3 100644
--- a/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLoweringTest.cpp
+++ b/lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLoweringTest.cpp
@@ -8,6 +8,8 @@
 #include "lib/Dialect/Rotom/Utils/RotomTensorExtLayoutLowering.h"
 #include "lib/Utils/Layout/Evaluate.h"
 #include "lib/Utils/Layout/IslConversion.h"
+#include "lib/Utils/Layout/Utils.h"
+#include "mlir/include/mlir/Analysis/Presburger/IntegerRelation.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"   // from @llvm-project
 #include "mlir/include/mlir/IR/MLIRContext.h"         // from @llvm-project
 #include "mlir/include/mlir/Support/LogicalResult.h"  // from @llvm-project
@@ -77,8 +79,8 @@ TEST(RotomTensorExtLayoutLoweringTest, ColumnMajor4x4Evaluate) {
   };
   std::vector<std::vector<int>> packed = evaluateLayout<int>(
       relation.value(), [&](const std::vector<int64_t>& domainPoint) -> int {
-        // Traversal dims are {dim1, dim0}, so relation vars are [col, row].
-        return matrix[domainPoint[1]][domainPoint[0]];
+        // Domain vars are tensor dims in order: [row, col].
+        return matrix[domainPoint[0]][domainPoint[1]];
       });
 
   std::vector<std::vector<int>> expected = {
@@ -110,12 +112,14 @@ TEST(RotomTensorExtLayoutLoweringTest, TiledRowMajor4x4Evaluate) {
       {9, 10, 11, 12},
       {13, 14, 15, 16},
   };
-  ASSERT_EQ(relation->getNumDomainVars(), 4);
+  // dim ids 0 and 1 are each split mixed-radix into two pieces, so the relation
+  // has one domain variable per tensor axis: i0 = row, i1 = col -- the pieces
+  // of an axis share its variable rather than each binding their own. The
+  // packing itself is plain 2x2-tiled row-major.
+  ASSERT_EQ(relation->getNumDomainVars(), 2);
   std::vector<std::vector<int>> packed = evaluateLayout<int>(
       relation.value(), [&](const std::vector<int64_t>& domainPoint) -> int {
-        const int64_t row = domainPoint[0] * 2 + domainPoint[2];
-        const int64_t col = domainPoint[1] * 2 + domainPoint[3];
-        return matrix[row][col];
+        return matrix[domainPoint[0]][domainPoint[1]];
       });
 
   std::vector<std::vector<int>> expected = {
@@ -182,8 +186,8 @@ TEST(RotomTensorExtLayoutLoweringTest, SplitColumnMajor4x4Evaluate) {
   };
   std::vector<std::vector<int>> packed = evaluateLayout<int>(
       relation.value(), [&](const std::vector<int64_t>& domainPoint) -> int {
-        // Traversal dims are {dim1, dim0}, so relation vars are [col, row].
-        return matrix[domainPoint[1]][domainPoint[0]];
+        // Domain vars are tensor dims in order: [row, col].
+        return matrix[domainPoint[0]][domainPoint[1]];
       });
 
   // Column-major packing, split into 4 ciphertexts of 4 slots: one column per
@@ -197,26 +201,34 @@ TEST(RotomTensorExtLayoutLoweringTest, SplitColumnMajor4x4Evaluate) {
   EXPECT_EQ(packed, expected);
 }
 
-TEST(RotomTensorExtLayoutLoweringTest, PreprocessAddsImplicitGap) {
+// The canonicalizing builder makes unused slot capacity explicit: a 4-vector
+// at n = 8 gains a front gap piece, so the stored dims show every slot.
+TEST(RotomTensorExtLayoutLoweringTest, BuilderInsertsExplicitGapFill) {
   MLIRContext context;
   context.loadDialect<RotomDialect>();
   DimAttr d0 = DimAttr::get(&context, /*dim=*/0, /*size=*/4, /*stride=*/1);
   ArrayAttr dims = ArrayAttr::get(&context, {d0});
   LayoutAttr layout = LayoutAttr::get(&context, dims, /*n=*/8);
 
+  ASSERT_EQ(layout.getDims().size(), 2u);
+  EXPECT_TRUE(cast<DimAttr>(layout.getDims()[0]).isGap());
+  EXPECT_EQ(cast<DimAttr>(layout.getDims()[0]).getSize(), 2);
+
   FailureOr<LayoutData> data = preprocessLayoutAttr(layout);
   ASSERT_TRUE(succeeded(data));
   EXPECT_EQ(data->n, 8);
   EXPECT_EQ(data->ctPrefixLen, 0);
-  ASSERT_EQ(data->gapDims.size(), 1);
-  EXPECT_EQ(data->gapDims[0].getDim(), -2);
-  EXPECT_EQ(data->gapDims[0].getSize(), 2);
   ASSERT_EQ(data->pieces.size(), 2);
-  EXPECT_EQ(data->pieces[0], LayoutPieceKind::Gap);
-  EXPECT_EQ(data->pieces[1], LayoutPieceKind::Traversal);
+  EXPECT_EQ(data->pieces[0].kind, LayoutPieceKind::Gap);
+  EXPECT_EQ(data->pieces[0].dim.getDim(), -2);
+  EXPECT_EQ(data->pieces[0].dim.getSize(), 2);
+  EXPECT_EQ(data->pieces[1].kind, LayoutPieceKind::Traversal);
 }
 
-TEST(RotomTensorExtLayoutLoweringTest, PreprocessPreservesTraversalDimsOrder) {
+TEST(RotomTensorExtLayoutLoweringTest, PreprocessSortsAxesByDimId) {
+  // The deduped axes are canonicalized to ascending dim id regardless of
+  // piece order, so the ISL lowering's domain variables always line up
+  // positionally with tensor dims.
   MLIRContext context;
   context.loadDialect<RotomDialect>();
   DimAttr d0 = DimAttr::get(&context, /*dim=*/0, /*size=*/4, /*stride=*/1);
@@ -226,9 +238,9 @@ TEST(RotomTensorExtLayoutLoweringTest, PreprocessPreservesTraversalDimsOrder) {
 
   FailureOr<LayoutData> data = preprocessLayoutAttr(layout);
   ASSERT_TRUE(succeeded(data));
-  ASSERT_EQ(data->traversalDims.size(), 2);
-  EXPECT_EQ(data->traversalDims[0].getDim(), 1);
-  EXPECT_EQ(data->traversalDims[1].getDim(), 0);
+  ASSERT_EQ(data->axes.size(), 2);
+  EXPECT_EQ(data->axes[0].getDim(), 0);
+  EXPECT_EQ(data->axes[1].getDim(), 1);
 }
 
 TEST(RotomTensorExtLayoutLoweringTest, RolledRowMajor2x2Evaluate) {
diff --git a/tests/Dialect/Rotom/IR/layout.mlir b/tests/Dialect/Rotom/IR/layout.mlir
index 6e925f3c6c..8ca92a1596 100644
--- a/tests/Dialect/Rotom/IR/layout.mlir
+++ b/tests/Dialect/Rotom/IR/layout.mlir
@@ -9,21 +9,30 @@ func.func private @ok(tensor<16xi32> {foo.bar = #layout_ok})
 
 // -----
 
-// This forces a non-pow2 slot dim size after splitting: size=3 divides 12 but not pow2.
-#bad = #rotom.layout<n = 12, dims = [[0:3:1]]> // expected-error {{slot dim size must be a power of two, got 3}}
+// A non-pow2 slot dim size is rejected.
+#bad = #rotom.layout<n = 12, dims = [[G:4:1], [0:3:1]]> // expected-error {{slot dim size must be a power of two, got 3}}
 func.func private @bad(tensor<16xi32> {foo.bar = #bad})
 
 // -----
 
-// Splitting case: size > n causes a ct/slot split, and slot-side size becomes n (must be pow2).
-#split_ok = #rotom.layout<n = 8, dims = [[0:16:1]]>
+// An oversized dim indexes ciphertexts (one element per ciphertext); the
+// unused slots are an explicit gap, so the slot side is all gap here.
+#split_ok = #rotom.layout<n = 8, dims = [[0:16:1] | [G:8:1]]>
 func.func private @split_ok(tensor<16xi32> {foo.bar = #split_ok})
 
 // -----
 
-// size > n but not divisible => verifier error (mirrors Python assert size % n == 0).
-#split_bad = #rotom.layout<n = 8, dims = [[0:10:1]]> // expected-error {{dim size 10 must be divisible by remaining slot capacity 8}}
-func.func private @split_bad(tensor<16xi32> {foo.bar = #split_bad})
+// The slot side must fill the ciphertext exactly; a written layout may not
+// leave capacity implicit.
+#underfilled = #rotom.layout<n = 8, dims = [[0:4:1]]> // expected-error {{slot dims must fill the ciphertext exactly (slot extent 4 vs n = 8); spell unused capacity as an explicit gap piece}}
+func.func private @underfilled(tensor<16xi32> {foo.bar = #underfilled})
+
+// -----
+
+// The written `|` must sit exactly at the derived boundary: both pieces here
+// fit the 8 slots, so neither may be claimed as a ciphertext dim.
+#bad_split = #rotom.layout<n = 8, dims = [[0:2:1] | [1:4:1]]> // expected-error {{the written `|` ciphertext/slot split (1 ciphertext dims) does not match the derived split (0): the slot side is the longest dims suffix whose extents fit n = 8}}
+func.func private @bad_split(tensor<16xi32> {foo.bar = #bad_split})
 
 // -----
 
diff --git a/tests/Dialect/Rotom/IR/syntax.mlir b/tests/Dialect/Rotom/IR/syntax.mlir
index 325c704b18..1bc657d1f0 100644
--- a/tests/Dialect/Rotom/IR/syntax.mlir
+++ b/tests/Dialect/Rotom/IR/syntax.mlir
@@ -2,20 +2,35 @@
 
 #d0 = #rotom.dim<[0:4:1]>
 #d1 = #rotom.dim<[1:4:1]>
-#plain = #rotom.layout<n = 8, dims = [[0:4:1]]>
+// Unused slot capacity is spelled as an explicit gap piece (the builders
+// insert it; written layouts must show it).
+#plain = #rotom.layout<n = 8, dims = [[G:2:1], [0:4:1]]>
 #rolled = #rotom.layout<n = 16, rolls = [(0, 1)], dims = [#d0, #d1]>
+// Replication and gap dims print as R and G; the numeric ids -1 and -2 are
+// accepted on input and round-trip to the letter forms.
+#repl_gap = #rotom.layout<n = 16, dims = [[R:2:1], [-2:2:1], [0:4:1]]>
+// The `|` separates ciphertext dims from slot dims (omitted when every dim
+// is a slot dim): here each of the 8 elements sits in its own ciphertext,
+// and the slot side is all gap.
+#split = #rotom.layout<n = 4, dims = [[0:8:1] | [G:4:1]]>
 
 // CHECK: #dim = #rotom.dim<[2:8:4]>
-// CHECK: #layout = #rotom.layout<n = 8, dims = {{\[\[0:4:1\]\]}}>
-// CHECK: #layout1 = #rotom.layout<n = 16, rolls = [(0, 1)], dims = {{\[\[0:4:1\], \[1:4:1\]\]}}>
+// CHECK: #layout = #rotom.layout<n = 8, dims = {{\[\[G:2:1\], \[0:4:1\]\]}}>
+// CHECK: #layout1 = #rotom.layout<n = 16, dims = {{\[\[R:2:1\], \[G:2:1\], \[0:4:1\]\]}}>
+// CHECK: #layout2 = #rotom.layout<n = 16, rolls = [(0, 1)], dims = {{\[\[0:4:1\], \[1:4:1\]\]}}>
+// CHECK: #layout3 = #rotom.layout<n = 4, dims = {{\[\[0:8:1\] \| \[G:4:1\]\]}}>
 // CHECK: module attributes
 // CHECK-SAME: rotom.dim_attr = #dim
 // CHECK-SAME: rotom.plain_layout = #layout
-// CHECK-SAME: rotom.rolled_layout = #layout1
+// CHECK-SAME: rotom.repl_gap_layout = #layout1
+// CHECK-SAME: rotom.rolled_layout = #layout2
+// CHECK-SAME: rotom.split_layout = #layout3
 module attributes {
   rotom.dim_attr = #rotom.dim<[2:8:4]>,
   rotom.plain_layout = #plain,
-  rotom.rolled_layout = #rolled
+  rotom.repl_gap_layout = #repl_gap,
+  rotom.rolled_layout = #rolled,
+  rotom.split_layout = #split
 } {
   func.func @f(%arg0: tensor<4x4xf32>) {
     return
diff --git a/tests/Dialect/Rotom/Transforms/doctest_seed_layout.mlir b/tests/Dialect/Rotom/Transforms/doctest_seed_layout.mlir
index 8ae956689b..735648a76e 100644
--- a/tests/Dialect/Rotom/Transforms/doctest_seed_layout.mlir
+++ b/tests/Dialect/Rotom/Transforms/doctest_seed_layout.mlir
@@ -3,15 +3,15 @@
 module {
   // CHECK: func.func @test_seeding(
   // CHECK-SAME: !secret.secret<tensor<4x4xf32>> {rotom.seed = #rotom.seed<layouts = [
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[0:2:1\], \[1:4:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[1:4:1\], \[0:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[0:4:1\], \[1:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[1:2:1\], \[0:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[0:2:1\], \[1:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[1:4:1\], \[0:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[0:4:1\], \[1:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[1:2:1\], \[0:4:1\]\]}}>
   // CHECK-SAME: ]>}, %{{.*}}: tensor<4x4xf32> {rotom.seed = #rotom.seed<layouts = [
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[0:2:1\], \[1:4:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[1:4:1\], \[0:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[0:4:1\], \[1:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[1:2:1\], \[0:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[0:2:1\], \[1:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[1:4:1\], \[0:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[0:4:1\], \[1:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[1:2:1\], \[0:4:1\]\]}}>
   // CHECK-SAME: ]>})
   func.func @test_seeding(%arg0: !secret.secret<tensor<4x4xf32>>, %arg1: tensor<4x4xf32>) -> !secret.secret<tensor<4x4xf32>> {
     // CHECK: secret.generic(%{{.*}}: !secret.secret<tensor<4x4xf32>>)
diff --git a/tests/Dialect/Rotom/Transforms/materialize_column_major.mlir b/tests/Dialect/Rotom/Transforms/materialize_column_major.mlir
index 55ad6fe353..6a19529d5a 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_column_major.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_column_major.mlir
@@ -9,7 +9,7 @@
 // CHECK-DAG: #tensor_ext.layout<
 // CHECK-DAG: [i0, i1] -> [ct, slot] :
 // CHECK-DAG: ct = 0
-// CHECK-DAG: slot = 4 * i0 + i1
+// CHECK-DAG: slot = i0 + 4 * i1
 module {
   func.func @f(%arg0: tensor<4x4xf32> {rotom.layout = #layout}) {
     return
diff --git a/tests/Dialect/Rotom/Transforms/materialize_ct_slot_split.mlir b/tests/Dialect/Rotom/Transforms/materialize_ct_slot_split.mlir
index fca9a6eb35..a867e8cbff 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_ct_slot_split.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_ct_slot_split.mlir
@@ -1,21 +1,23 @@
 // RUN: heir-opt %s --rotom-materialize-tensor-ext-layout | FileCheck %s
 
-// Rotom ``[0:2:2];[1:2:2][0:2:1][1:2:1]`` with ``n = 8``: first dim is
-// ciphertext traversal, remaining dims pack within each CT. This is a tiled
-// row-major layout.
+// Rotom ``[0:2:2][1:2:2][0:2:1][1:2:1]`` with ``n = 8`` on a 4x4 tensor: each
+// tensor axis is split mixed-radix into a stride-2 (high) and stride-1 (low)
+// piece. The cumulative product of extents crosses n after the slot pieces, so
+// the high part of axis 0 indexes ciphertexts. This is a 2x2-tiled row-major
+// layout (ct = tile-row, slot = 4*tile-col + 2*within-row + within-col).
 #d0 = #rotom.dim<[0:2:2]>
 #d1 = #rotom.dim<[1:2:2]>
 #d2 = #rotom.dim<[0:2:1]>
 #d3 = #rotom.dim<[1:2:1]>
-#layout = #rotom.layout<n = 8, dims = [#d0, #d1, #d2, #d3]>
+#layout = #rotom.layout<n = 8, dims = [#d0 | #d1, #d2, #d3]>
 
-// CHECK:   func.func @f(%arg0: tensor<2x2x2x2xf32> {tensor_ext.layout =
+// CHECK:   func.func @f(%arg0: tensor<4x4xf32> {tensor_ext.layout =
 // CHECK-DAG: #tensor_ext.layout<
-// CHECK-DAG: [i0, i1, i2, i3] -> [ct, slot] :
-// CHECK-DAG: ct = i0
-// CHECK-DAG: slot = 4 * i1 + 2 * i2 + i3
+// CHECK-DAG: [i0, i1] -> [ct, slot] :
+// CHECK-DAG: ct = floor((i0) / 2)
+// CHECK-DAG: slot = 2 * (i0 - 2 * floor((i0) / 2)) + 4 * floor((i1) / 2) + (i1 - 2 * floor((i1) / 2))
 module {
-  func.func @f(%arg0: tensor<2x2x2x2xf32> {rotom.layout = #layout}) {
+  func.func @f(%arg0: tensor<4x4xf32> {rotom.layout = #layout}) {
     return
   }
 }
diff --git a/tests/Dialect/Rotom/Transforms/materialize_explicit_gap_dim.mlir b/tests/Dialect/Rotom/Transforms/materialize_explicit_gap_dim.mlir
index 78c8ce2867..b055cf7f0b 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_explicit_gap_dim.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_explicit_gap_dim.mlir
@@ -2,7 +2,7 @@
 
 // Rotom ``[0:4:1][G:2:1]`` with ``n = 8``: in-slot row with explicit gap.
 #d0 = #rotom.dim<[0:4:1]>
-#g0 = #rotom.dim<[-2:2:1]>
+#g0 = #rotom.dim<[G:2:1]>
 #layout = #rotom.layout<n = 8, dims = [#d0, #g0]>
 
 // CHECK:   func.func @f(%arg0: tensor<4xf32> {tensor_ext.layout =
diff --git a/tests/Dialect/Rotom/Transforms/materialize_implicit_gap_dim.mlir b/tests/Dialect/Rotom/Transforms/materialize_implicit_gap_dim.mlir
index d0b1659c18..1dd0f61b87 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_implicit_gap_dim.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_implicit_gap_dim.mlir
@@ -1,10 +1,12 @@
 // RUN: heir-opt %s --rotom-materialize-tensor-ext-layout | FileCheck %s
 
 
-// Rotom ``[G:2:1][0:4:1]`` with ``n = 8``: Row-major, first 4. Implicit gap
-// dimension, ``[G:2:1]``, should be added in front.
+// Rotom ``[G:2:1][0:4:1]`` with ``n = 8``: row-major in the low 4 slots. The
+// unused capacity is the explicit front gap piece, which contributes no
+// address term (its blocks stay unclaimed).
+#gap = #rotom.dim<[G:2:1]>
 #d0 = #rotom.dim<[0:4:1]>
-#layout = #rotom.layout<n = 8, dims = [#d0]>
+#layout = #rotom.layout<n = 8, dims = [#gap, #d0]>
 
 // CHECK:   func.func @f(%arg0: tensor<4xf32> {tensor_ext.layout =
 // CHECK-DAG: #tensor_ext.layout<
diff --git a/tests/Dialect/Rotom/Transforms/materialize_repeated_column_major.mlir b/tests/Dialect/Rotom/Transforms/materialize_repeated_column_major.mlir
index c23a4fb92e..c570fefa74 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_repeated_column_major.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_repeated_column_major.mlir
@@ -2,18 +2,18 @@
 
 // Rotom ``[R:4:1];[1:4:1][0:4:1]`` with ``n = 16`` (replication + column-major
 // traversals in ``dims``). Replication is projected to ciphertext index ``ct``
-// via existential ``d2``; slots pack ``4 * i0 + i1``.
+// via existential ``d2``; slots pack column-major ``i0 + 4 * i1``.
 #d0 = #rotom.dim<[0:4:1]>
 #d1 = #rotom.dim<[1:4:1]>
-#d2 = #rotom.dim<[-1:4:1]>
-#layout = #rotom.layout<n = 16, dims = [#d2, #d1, #d0]>
+#d2 = #rotom.dim<[R:4:1]>
+#layout = #rotom.layout<n = 16, dims = [#d2 | #d1, #d0]>
 
 // CHECK:   func.func @f(%arg0: tensor<4x4xf32> {tensor_ext.layout =
 // CHECK-DAG: #tensor_ext.layout<
 // CHECK-DAG: [i0, i1] -> [ct, slot] :
 // CHECK-DAG: exists d2
 // CHECK-DAG: ct = d2
-// CHECK-DAG: slot = 4 * i0 + i1
+// CHECK-DAG: slot = i0 + 4 * i1
 module {
   func.func @f(%arg0: tensor<4x4xf32> {rotom.layout = #layout}) {
     return
diff --git a/tests/Dialect/Rotom/Transforms/materialize_replication_dim.mlir b/tests/Dialect/Rotom/Transforms/materialize_replication_dim.mlir
index a815c5de32..4bb72854cc 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_replication_dim.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_replication_dim.mlir
@@ -3,7 +3,7 @@
 // Rotom ``[0:4:1][R:2:1]`` with ``n = 8``: row-major where each value is
 // repeated twice.
 #d0 = #rotom.dim<[0:4:1]>
-#r0 = #rotom.dim<[-1:2:4]>
+#r0 = #rotom.dim<[R:2:4]>
 #layout = #rotom.layout<n = 8, dims = [#d0, #r0]>
 
 // CHECK:   func.func @f(%arg0: tensor<4xf32> {tensor_ext.layout =
diff --git a/tests/Dialect/Rotom/Transforms/materialize_slot_gap_dims.mlir b/tests/Dialect/Rotom/Transforms/materialize_slot_gap_dims.mlir
index 69f5ec07f7..e6f27ea40e 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_slot_gap_dims.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_slot_gap_dims.mlir
@@ -4,8 +4,8 @@
 // variables with g_k = 0 (payload only at gap index 0; other indices zero-fill).
 #d0 = #rotom.dim<[0:2:1]>
 #d1 = #rotom.dim<[1:2:2]>
-#g0 = #rotom.dim<[-2:2:1]>
-#g1 = #rotom.dim<[-2:2:4]>
+#g0 = #rotom.dim<[G:2:1]>
+#g1 = #rotom.dim<[G:2:4]>
 #layout = #rotom.layout<n = 16, dims = [#d0, #d1, #g0, #g1]>
 
 // CHECK:   func.func @f(%arg0: tensor<2x2xf32> {tensor_ext.layout =
diff --git a/tests/Dialect/Rotom/Transforms/materialize_tiled_duplicate_dim.mlir b/tests/Dialect/Rotom/Transforms/materialize_tiled_duplicate_dim.mlir
index fcbffd72cb..55fc63ae7a 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_tiled_duplicate_dim.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_tiled_duplicate_dim.mlir
@@ -1,21 +1,23 @@
 // RUN: heir-opt %s --rotom-materialize-tensor-ext-layout | FileCheck %s
 
-// Tiled row-major style: dim ids repeat across traversals, but (dim, stride)
-// pairs are distinct. With ``n = 8``, Rotom ``;`` split: first traversal in
-// ``dims`` is ciphertext index, the rest pack within each CT.
+// A tensor axis split mixed-radix: dim id 0 appears as two pieces ([0:2:2] and
+// [0:2:1]) that together form axis 0 (extent 4), and likewise dim id 1. The
+// strides are the within-axis cumulative products (1, then 2), so the two pieces
+// share one domain variable. The cumulative product of extents crosses n = 8
+// after the slot pieces, putting the high part of axis 0 on the ciphertext axis.
 #d0s2 = #rotom.dim<[0:2:2]>
 #d1s2 = #rotom.dim<[1:2:2]>
 #d0s1 = #rotom.dim<[0:2:1]>
 #d1s1 = #rotom.dim<[1:2:1]>
-#layout = #rotom.layout<n = 8, dims = [#d0s2, #d1s2, #d0s1, #d1s1]>
+#layout = #rotom.layout<n = 8, dims = [#d0s2 | #d1s2, #d0s1, #d1s1]>
 
-// CHECK:   func.func @f(%arg0: tensor<2x2x2x2xf32> {tensor_ext.layout =
+// CHECK:   func.func @f(%arg0: tensor<4x4xf32> {tensor_ext.layout =
 // CHECK-DAG: #tensor_ext.layout<
-// CHECK-DAG: [i0, i1, i2, i3] -> [ct, slot] :
-// CHECK-DAG: ct = i0
-// CHECK-DAG: slot = 4 * i1 + 2 * i2 + i3
+// CHECK-DAG: [i0, i1] -> [ct, slot] :
+// CHECK-DAG: ct = floor((i0) / 2)
+// CHECK-DAG: slot = 2 * (i0 - 2 * floor((i0) / 2)) + 4 * floor((i1) / 2) + (i1 - 2 * floor((i1) / 2))
 module {
-  func.func @f(%arg0: tensor<2x2x2x2xf32> {rotom.layout = #layout}) {
+  func.func @f(%arg0: tensor<4x4xf32> {rotom.layout = #layout}) {
     return
   }
 }
diff --git a/tests/Dialect/Rotom/Transforms/materialize_vector.mlir b/tests/Dialect/Rotom/Transforms/materialize_vector.mlir
index 9b079e13a3..d99e96183d 100644
--- a/tests/Dialect/Rotom/Transforms/materialize_vector.mlir
+++ b/tests/Dialect/Rotom/Transforms/materialize_vector.mlir
@@ -1,7 +1,8 @@
 // RUN: heir-opt %s --rotom-materialize-tensor-ext-layout | FileCheck %s
 
+#gap = #rotom.dim<[G:2:1]>
 #d0 = #rotom.dim<[0:4:1]>
-#layout = #rotom.layout<n = 8, dims = [#d0]>
+#layout = #rotom.layout<n = 8, dims = [#gap, #d0]>
 
 // CHECK:   func.func @f(%arg0: tensor<4xf32> {tensor_ext.layout =
 // CHECK-DAG: #tensor_ext.layout<
diff --git a/tests/Dialect/Rotom/Transforms/seed_layout.mlir b/tests/Dialect/Rotom/Transforms/seed_layout.mlir
index b4c295ad85..1fc9e964b6 100644
--- a/tests/Dialect/Rotom/Transforms/seed_layout.mlir
+++ b/tests/Dialect/Rotom/Transforms/seed_layout.mlir
@@ -3,15 +3,15 @@
 module {
   // CHECK: func.func @test_seeding(
   // CHECK-SAME: !secret.secret<tensor<4x4xf32>> {rotom.seed = #rotom.seed<layouts = [
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[0:2:1\], \[1:4:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[1:4:1\], \[0:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[0:4:1\], \[1:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[1:2:1\], \[0:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[0:2:1\], \[1:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[1:4:1\], \[0:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[0:4:1\], \[1:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[1:2:1\], \[0:4:1\]\]}}>
   // CHECK-SAME: ]>}, %{{.*}}: tensor<4x4xf32> {rotom.seed = #rotom.seed<layouts = [
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[0:2:1\], \[1:4:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[1:4:1\], \[0:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[0:4:1\], \[1:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[1:2:1\], \[0:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[0:2:1\], \[1:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[1:4:1\], \[0:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[0:4:1\], \[1:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[1:2:1\], \[0:4:1\]\]}}>
   // CHECK-SAME: ]>})
   func.func @test_seeding(%arg0: !secret.secret<tensor<4x4xf32>>, %arg1: tensor<4x4xf32>) -> !secret.secret<tensor<4x4xf32>> {
     // CHECK: secret.generic(%{{.*}}: !secret.secret<tensor<4x4xf32>>)
@@ -51,15 +51,15 @@ module {
 module {
   // CHECK: func.func @test_seeding_non_pow2(
   // CHECK-SAME: !secret.secret<tensor<3x3xf32>> {rotom.seed = #rotom.seed<layouts = [
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[0:2:1\], \[1:4:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[1:4:1\], \[0:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[0:4:1\], \[1:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[1:2:1\], \[0:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[0:2:1\], \[1:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[1:4:1\], \[0:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[0:4:1\], \[1:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[1:2:1\], \[0:4:1\]\]}}>
   // CHECK-SAME: ]>}, %{{.*}}: tensor<3x3xf32> {rotom.seed = #rotom.seed<layouts = [
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[0:2:1\], \[1:4:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\], \[1:4:1\], \[0:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[0:4:1\], \[1:2:1\]\]}}>
-  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\], \[1:2:1\], \[0:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[0:2:1\], \[1:4:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[0:2:2\] \| \[1:4:1\], \[0:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[0:4:1\], \[1:2:1\]\]}}>
+  // CHECK-SAME: #rotom.layout<n = 8, dims = {{\[\[1:2:2\] \| \[1:2:1\], \[0:4:1\]\]}}>
   // CHECK-SAME: ]>})
   func.func @test_seeding_non_pow2(%arg0: !secret.secret<tensor<3x3xf32>>, %arg1: tensor<3x3xf32>) -> !secret.secret<tensor<3x3xf32>> {
     // CHECK: secret.generic(%{{.*}}: !secret.secret<tensor<3x3xf32>>)