Rotom Layout Dialect by edwjchen · Pull Request #2829 · google/heir

edwjchen · 2026-03-30T15:33:57Z

This is first iteration port of Rotom's layout representation as a MLIR Dialect into HEIR.

The overall goal of this dialect is to act as a starting point for layout assignment. This key difference about this layout representation compared to the one used in HEIR is that it retains tensor metadata information, such as the specific iteration order along a tensor dimension. This allows Rotom to create alignment rules, which check for two given layouts and a tensor operator, if the two layouts can implement that tensor operator without any layout conversions. These rules ensure that the traversal patterns of input tensors—as encoded in their layouts—are compatible with the iteration structure of the target tensor operation.

Rotom's layout representation primarily contains two main ingredients, a list of traversal dimensions, and a list of rolls. Traversal dimensions represent a nested loop-nest iteration over the plaintext tensor, mapping each element into a ciphertext slot. Rolls represent a transformation along this iteration space, by modifying the index traversal of one traversal dimension by another. This PR includes support for transforming a Rotom layout composed of only traversal dimensions into a Presburger relation used in HEIR's layout representation. Rolls will be implemented in a follow-up PR.

Rotom's Dialect:

RotomAttributes: Defines the traversal dimensions and layout structure.
RotomOp: An entry function to parse an ssa value into a Rotom traversal dimension.
RotomEnum: Defines the traversal dimension type (whether its filled or empty)

MaterializeTensorExtLayout:

A transformation pass to convert Rotom's layout representation into HEIR's layout representation.

tools/heir-opt.cpp

asraa · 2026-04-02T13:31:09Z

lib/Dialect/Rotom/IR/RotomAttributes.td

+    "int64_t":$dim,
+    "int64_t":$extent,
+    DefaultValuedParameter<"int64_t", "1">:$stride,
+    "::mlir::heir::rotom::DimType":$dim_type


if i'm understanding, dim_type is only relevant when dim is -1 (gapped or replicate), in which case dim_type differentiates the two (gap = empty, replicate = fill)? If so, maybe it's better to consolidate the dim into a single attr that defines that (and not use an enum attr)?

def DimAttr : MyDialect_Attr<"Dim", "dim"> { let parameters = (ins "int64_t":$value); // Print "gap" for -1, and "replicate" for -2, or value otherwise let hasCustomAssemblyFormat = 1; let extraClassDefinition = [{ bool isGapped() { return getValue() == -1; } bool isReplicate() { return getValue() == -2; } }]; // Write verifier to ensure value >= -2 }

That's correct! I removed the dim_type and changed the DimAttr to be -1 for replicated dimensions and -2 for gap dimensions.

asraa · 2026-04-02T13:32:52Z

lib/Dialect/Rotom/IR/RotomAttributes.td

+  let parameters = (ins
+    "::mlir::ArrayAttr":$dims,
+    "int64_t":$n,
+    DefaultValuedParameter<"bool", "false">:$secret


i'm not totally sure if secret is required - we have tools in our output to determine if a value is secret, so I don't think you need this (it doesn't make a different to the layout specification I think)

+1. I think the existing secretness analysis should suffice for static analysis of what SSA values are secret. (Moreover, a layout itself is not semantically something that can be considered secret or not)

asraa · 2026-04-02T13:39:30Z

lib/Dialect/Rotom/IR/RotomOps.td

+def Rotom_ApplyDimOp : Rotom_Op<"apply_dim", [Pure]> {
+  let summary = "Annotate a value with a single Rotom dim.";
+  let description = [{
+    Attaches a verified `rotom.dim` attribute to an SSA value.


is this meant to apply a rotom dim to an SSA value that already has a rotom.layout? i'm not totally sure the reason this is here, but i would expect that this takes an SSA value with a layout, applies the dim, and then the result is a new layout with the extra dim

Maybe it would help to wait to add this op until there's a pass that uses it, because I also agree I don't quite understand what this op is meant to do. It seems kind of like mgmt.init or tensor_ext.assign_layout, but I don't understand why there's only a single dim involved.

Removed this function!

asraa · 2026-04-02T13:42:58Z

lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout.cpp

+}  // namespace
+}  // namespace mlir::heir::rotom
+
+namespace {


Interesting! I would suggest defining the transform pass with tablegen - that way it auto-populates the docs. This means creating a Transforms/MaterializeTensorExtLayout.td file and some boilerplate (here's a simple pass: https://github.com/google/heir/tree/main/lib/Transforms/AnnotateSecretness, or for one in the dialect folder: https://github.com/google/heir/tree/main/lib/Dialect/CKKS/Transforms). Then instead of this code, you would add the imported generated code:

#define GEN_PASS_DEF_MATERIALIZETENSOREXTLAYOUT
#include "lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout/MaterializeTensorExtLayout.h.inc"

+1 Try https://github.com/google/heir/blob/main/scripts/templates/templates.py with new_dialect_transform and it should generate all the boilerplate with a FIXME for the body of the pass.

I've refactored the transform pass to use tablegen. I've also moved the helper functions into Util/ so we can better unit test this lowering process.

asraa · 2026-04-02T13:47:12Z

lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout.cpp

+    registry.insert<tensor_ext::TensorExtDialect>();
+  }
+
+  /// Finds `rotom.layout` on the module, lowers it to an ISL string, sets


is this just testing conversion of a rotom layout to an ISL layout, or in general do you expect that modules might have any number of attributes with rotom layouts on it?

If it's just for testing, you could expect that it will be a single rotom.layout and use something like module->getAttrOfType<LayoutAttr>("rotom.layout")

This was for testing purposes, see the changes below with Util/ for the new unit tests!

asraa · 2026-04-02T13:49:36Z

lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout.cpp

+                            numReplication, numGap);
+}
+
+struct RotomMaterializeTensorExtLayoutPass


see below, but this is all autogenerated if you define a tablegen transform pass

j2kun · 2026-03-31T16:13:09Z

lib/Dialect/Rotom/IR/RotomAttributes.td

+
+  let parameters = (ins
+    "int64_t":$dim,
+    "int64_t":$extent,


Any reason not to call this "size"? IIUC you use it as a size, whereas "extent" to me means "upper bound"

Or, if not, the precise definition of the terms should be in the description string.

Refactored to use size over extent

j2kun · 2026-03-31T16:16:33Z

lib/Dialect/Rotom/IR/RotomAttributes.td

+    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


If you want to connect this to the Rotom notation, either link to the paper and cite a particular section, or else provide examples here. Otherwise there is no reference to understand what the ; means in that context.

Also instead of double backticks maybe use single backticks?

Added the citation! Rotom Section 4.2

j2kun · 2026-03-31T16:19:35Z

lib/Dialect/Rotom/IR/RotomAttributes.cpp

+    return emitError() << "`n` must be > 0, got " << n;
+  }
+
+  struct TraversalDim {


DimAttr is equivalent to this struct, no? Why have a separate one?

j2kun · 2026-03-31T16:30:12Z

lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout.cpp

+  /// `tensor_ext.layout`, or fails the pass if lowering is unsupported.
+  void runOnOperation() override {
+    ModuleOp m = getOperation();
+    auto attrs = m->getAttrDictionary();


I suppose this pass is fine for initial testing purposes, but ultimately these attributes are not going to be stored on the module, but attached to particular function arguments or SSA values.

If I were you, I'd start from the module, use walkValues (

heir/lib/Utils/Utils.h

Line 97 in a379494

void walkValues(Operation* op, ValueProcessor valueProcessor);

) and for each value use findAttributeAssociatedWith (

heir/lib/Utils/AttributeUtils.h

Line 26 in a379494

FailureOr<Attribute> findAttributeAssociatedWith(Value value,

) and setAttributeAssociatedWith for the converted output.

Then you can attach a rotom layout to any SSA value (say, a function argument or operation result) similar to tests like https://github.com/google/heir/blob/main/tests/Transforms/convert_to_ciphertext_semantics/doctest.mlir that will suffice for testing this conversion and be future-proof for the later time when you want to convert all layout attributes anywhere in the IR to HEIR's presburger relations.

j2kun · 2026-03-31T16:32:19Z

lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout.cpp

+/// non-replication fill dims; infer ciphertext prefix length and implicit front
+/// gap; optionally reorder traversal `i*` by ascending logical dim id when all
+/// those ids are distinct; then emit the CT/slot map.
+static FailureOr<std::string> toTensorExtIslString(LayoutAttr layout) {


I think we mentioned this in office hours, but it might be more ergonomic for testing to move these to a utility class that can be unit tested, and then the output layout can be materialized on an actual tensor value and compared to the expected behavior of a rotom layout. Something like here:

heir/lib/Utils/Layout/ConvolutionTest.cpp

Line 131 in a379494

TEST(ConvolutionTest, ConvFilterRelationEvaluateStrided) {

Added this change and two unit tests for row-major and column-major layouts.

j2kun · 2026-04-02T15:35:15Z

lib/Dialect/Rotom/IR/RotomAttributes.td

+  let parameters = (ins
+    "::mlir::ArrayAttr":$dims,
+    "int64_t":$n,
+    DefaultValuedParameter<"bool", "false">:$secret


+1. I think the existing secretness analysis should suffice for static analysis of what SSA values are secret. (Moreover, a layout itself is not semantically something that can be considered secret or not)

j2kun · 2026-04-02T15:36:02Z

lib/Dialect/Rotom/Transforms/MaterializeTensorExtLayout.cpp

+}  // namespace
+}  // namespace mlir::heir::rotom
+
+namespace {


+1 Try https://github.com/google/heir/blob/main/scripts/templates/templates.py with new_dialect_transform and it should generate all the boilerplate with a FIXME for the body of the pass.

j2kun · 2026-04-02T15:38:44Z

lib/Dialect/Rotom/IR/RotomDialect.td

+  }];
+  let cppNamespace = "::mlir::heir::rotom";
+
+  let extraClassDeclaration = [{


I think you can set let useDefaultAttributePrinterParser = 1; to avoid this boilerplate.

- removed dim_type on the DimAttr, and used dim to denote replicated and gap dimensions - removed ApplyDimOp, unnecessary function - refactored MaterializeTensorExtLayout to use the automated tablegen pass - refactored a dimension into - refactored the MaterializeTensorExtLayout pass to use walkValues, to attach a rotom layout to any SSA value - refactored the MaterializeTensorExtLayout pass into Util for better unit testing

Rotom layout dialect and materialization pass

88e168e

asraa reviewed Apr 2, 2026

View reviewed changes

j2kun reviewed Apr 2, 2026

View reviewed changes

edwjchen added 2 commits April 6, 2026 05:01

added tiled and multi-ciphertext tests

3c4bb96

Conversation

edwjchen commented Mar 30, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants