wip: Rotom layout alignment, assignment, and lowering

.gitignore

@@ -1,11 +1,9 @@
 # macOS
 .DS_Store
 
-# bazel output symlinks
-bazel-bin
-bazel-heir
-bazel-out
-bazel-testlogs
+# bazel output symlinks (bazel-bin, bazel-out, bazel-testlogs, and the
+# per-workspace/worktree bazel-<dir> convenience symlink)
+/bazel-*
 
 # default cmake output directory
 build/

lib/Dialect/Rotom/IR/RotomAttributes.h

@@ -26,11 +26,26 @@ struct LayoutData {
   llvm::SmallVector<DimAttr> gapDims;
   llvm::SmallVector<LayoutPieceKind> pieces;
   llvm::SmallVector<int64_t> pieceIndex;
+  // pieceModBy and pieceDivBy are a piece's Rotom extent and stride (from
+  // [dim:extent:stride]), named for how the emitter consumes them: the mixed-
+  // radix digit of tensor index i is (i / pieceDivBy) mod pieceModBy. pieceModBy
+  // is 0 on the most-significant digit of an axis -- there i / pieceDivBy is
+  // already below the extent, so the modulus is redundant and dropped. A whole
+  // dim packed as one piece uses pieceModBy 0 and pieceDivBy 1 => digit == i;
+  // the same holds for non-traversal pieces.
+  llvm::SmallVector<int64_t> pieceModBy;
+  llvm::SmallVector<int64_t> pieceDivBy;
 };
 
 /// Preprocess a Rotom layout.
 FailureOr<LayoutData> preprocessLayoutAttr(LayoutAttr attr);
 
+/// 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
@@ -57,14 +59,27 @@ def Rotom_LayoutAttr : Rotom_Attr<"Layout", "layout"> {
 
     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
-    uses modular addition to modify the indices of  `dims[i]` by the indices of
-    `dims[j]`.
+    rewrites `dims[i]`'s index to `(idx_i - s * idx_j) mod extent(dims[i])`,
+    where the scale `s` defaults to 1 and may be spelled as a third tuple
+    element `(i, j, s)`. 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.
+
+    When a traversal dim is packed as several mixed-radix pieces, a roll FROM
+    any of its pieces rewrites the whole dim's index (each piece then takes its
+    digit of the rolled index), and a roll BY one of its pieces shifts by that
+    piece's digit -- together with a scale this expresses baby-step/giant-step
+    diagonal packings.
+
+    Optional **rollScales** stores one scale per roll pair; it must be omitted
+    when every scale is 1 (the canonical form the parser and builders produce).
   }];
 
   let parameters = (ins
     "::mlir::ArrayAttr":$dims,
     "int64_t":$n,
-    OptionalParameter<"::mlir::DenseI64ArrayAttr">:$rolls
+    OptionalParameter<"::mlir::DenseI64ArrayAttr">:$rolls,
+    OptionalParameter<"::mlir::DenseI64ArrayAttr">:$rollScales
   );
 
   let genVerifyDecl = 1;
@@ -74,6 +89,10 @@ def Rotom_LayoutAttr : Rotom_Attr<"Layout", "layout"> {
     /// Layout with no `roll(i,j)` metadata (empty rolls storage).
     static ::mlir::heir::rotom::LayoutAttr get(::mlir::MLIRContext *context,
                                                ::mlir::ArrayAttr dims, int64_t n);
+    /// Layout with unit-scale rolls (empty scale storage).
+    static ::mlir::heir::rotom::LayoutAttr get(::mlir::MLIRContext *context,
+                                               ::mlir::ArrayAttr dims, int64_t n,
+                                               ::mlir::DenseI64ArrayAttr rolls);
   }];
 }
 

lib/Dialect/Rotom/Transforms/BUILD

@@ -8,6 +8,62 @@ package(
     default_visibility = ["//visibility:public"],
 )
 
+add_heir_transforms(
+    generated_target_name = "layout_assignment_pass_inc_gen",
+    header_filename = "LayoutAssignment/LayoutAssignment.h.inc",
+    pass_name = "LayoutAssignment",
+    registration_name = "RotomLayoutAssignment",
+    td_file = "LayoutAssignment/LayoutAssignment.td",
+)
+
+cc_library(
+    name = "LayoutAssignment",
+    srcs = [
+        "LayoutAssignment/Candidate.cpp",
+        "LayoutAssignment/CostModel.cpp",
+        "LayoutAssignment/DimMaps.cpp",
+        "LayoutAssignment/LayoutAssignment.cpp",
+        "LayoutAssignment/ValueUtils.cpp",
+        "LayoutAssignment/gen/Contraction.cpp",
+        "LayoutAssignment/gen/Elementwise.cpp",
+        "LayoutAssignment/gen/ReduceTranspose.cpp",
+        "LayoutAssignment/gen/Reshape.cpp",
+        "LayoutAssignment/gen/Structural.cpp",
+    ],
+    hdrs = [
+        "LayoutAssignment/AssignmentContext.h",
+        "LayoutAssignment/Candidate.h",
+        "LayoutAssignment/CostModel.h",
+        "LayoutAssignment/DimMaps.h",
+        "LayoutAssignment/Generators.h",
+        "LayoutAssignment/LayoutAssignment.h",
+        "LayoutAssignment/ValueUtils.h",
+    ],
+    deps = [
+        ":layout_assignment_pass_inc_gen",
+        "@heir//lib/Dialect/Rotom/IR:Dialect",
+        "@heir//lib/Dialect/Rotom/Utils:ContractionAlignment",
+        "@heir//lib/Dialect/Rotom/Utils:LayoutAlignment",
+        "@heir//lib/Dialect/Rotom/Utils:RotomTensorExtLayoutLowering",
+        "@heir//lib/Dialect/Secret/IR:SecretPatterns",
+        "@heir//lib/Dialect/TensorExt/IR:Dialect",
+        "@heir//lib/Kernel",
+        "@heir//lib/Utils:AttributeUtils",
+        "@heir//lib/Utils/Layout:IslConversion",
+        "@heir//lib/Utils/Layout:Utils",
+        "@heir//lib/Utils:MathUtils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+    ],
+)
+
 add_heir_transforms(
     generated_target_name = "materialize_tensor_ext_layout_pass_inc_gen",
     header_filename = "MaterializeTensorExtLayout/MaterializeTensorExtLayout.h.inc",
@@ -24,16 +80,47 @@ cc_library(
         ":materialize_tensor_ext_layout_pass_inc_gen",
         "@heir//lib/Dialect/Rotom/IR:Dialect",
         "@heir//lib/Dialect/Rotom/Utils:RotomTensorExtLayoutLowering",
+        "@heir//lib/Dialect/Secret/IR:Dialect",
         "@heir//lib/Dialect/TensorExt/IR:Dialect",
         "@heir//lib/Utils",
         "@heir//lib/Utils:AttributeUtils",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
     ],
 )
 
+add_heir_transforms(
+    generated_target_name = "outline_kernels_pass_inc_gen",
+    header_filename = "OutlineKernels/OutlineKernels.h.inc",
+    pass_name = "OutlineKernels",
+    registration_name = "RotomOutlineKernels",
+    td_file = "OutlineKernels/OutlineKernels.td",
+)
+
+cc_library(
+    name = "OutlineKernels",
+    srcs = ["OutlineKernels/OutlineKernels.cpp"],
+    hdrs = ["OutlineKernels/OutlineKernels.h"],
+    deps = [
+        ":outline_kernels_pass_inc_gen",
+        "@heir//lib/Dialect/Rotom/IR:Dialect",
+        "@heir//lib/Dialect/Rotom/Utils:ContractionAlignment",
+        "@heir//lib/Dialect/Secret/IR:Dialect",
+        "@heir//lib/Utils:AttributeUtils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+    ],
+)
+
 add_heir_transforms(
     generated_target_name = "seed_layout_pass_inc_gen",
     header_filename = "SeedLayout/SeedLayout.h.inc",
@@ -63,17 +150,20 @@ cc_library(
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
-        "@llvm-project//mlir:TensorDialect",
     ],
 )
 
 cc_library(
     name = "Transforms",
     hdrs = ["Passes.h"],
     deps = [
+        ":LayoutAssignment",
         ":MaterializeTensorExtLayout",
+        ":OutlineKernels",
         ":SeedLayout",
+        ":layout_assignment_pass_inc_gen",
         ":materialize_tensor_ext_layout_pass_inc_gen",
+        ":outline_kernels_pass_inc_gen",
         ":seed_layout_pass_inc_gen",
         "@heir//lib/Dialect/Rotom/IR:Dialect",
         "@llvm-project//mlir:Pass",

lib/Dialect/Rotom/Transforms/LayoutAssignment/AssignmentContext.h

@@ -0,0 +1,55 @@
+#ifndef LIB_DIALECT_ROTOM_TRANSFORMS_LAYOUTASSIGNMENT_ASSIGNMENTCONTEXT_H_
+#define LIB_DIALECT_ROTOM_TRANSFORMS_LAYOUTASSIGNMENT_ASSIGNMENTCONTEXT_H_
+
+#include <cstdint>
+#include <optional>
+
+#include "lib/Dialect/Rotom/IR/RotomAttributes.h"
+#include "lib/Dialect/Rotom/Transforms/LayoutAssignment/Candidate.h"
+#include "lib/Kernel/KernelName.h"
+#include "mlir/include/mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"      // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::heir::rotom {
+
+// The narrow interface the per-op candidate generators use to read and write
+// the candidate map and query costs, so each generator depends only on this
+// rather than on the LayoutAssignment pass struct. The pass implements it.
+class AssignmentContext {
+ public:
+  virtual ~AssignmentContext() = default;
+
+  // Seed `value` from its rotom.seed attribute (no-op if already processed).
+  virtual void seedValue(Value value) = 0;
+
+  // The candidate layouts assigned to `value` (seeding it first if needed).
+  virtual SmallVector<Candidate> candidatesForValue(Value value) = 0;
+
+  // Record `newCandidates` as the candidates for `value`, after validating and
+  // folding each into its assignment.
+  virtual void setCandidates(Value value,
+                             ArrayRef<Candidate> newCandidates) = 0;
+
+  // Set the chosen candidate set on every tensor result of `op`.
+  virtual void assignResultsFromCandidates(Operation* op,
+                                           ArrayRef<Candidate> chosen) = 0;
+
+  // Bring all of `op`'s operands onto a common result layout (the generic
+  // N-operand combiner).
+  virtual SmallVector<Candidate> chooseCommonOperandCandidates(
+      Operation* op, KernelKind kind) = 0;
+
+  // The specialized binary convert-then-compute combiner (add/sub/mul).
+  virtual SmallVector<Candidate> chooseElementwiseKernels(
+      ArrayRef<Value> operands, KernelKind kind,
+      function_ref<int64_t(LayoutAttr)> computeCostFn,
+      std::optional<KernelName> rotomKernel = std::nullopt) = 0;
+
+  // Real (cached) rotation cost of converting `from` onto `to`.
+  virtual int64_t cachedConversionCost(LayoutAttr from, LayoutAttr to) = 0;
+};
+
+}  // namespace mlir::heir::rotom
+
+#endif  // LIB_DIALECT_ROTOM_TRANSFORMS_LAYOUTASSIGNMENT_ASSIGNMENTCONTEXT_H_