Add domain schedule attribute to tensor_ext.assign_layout and convert_layout

* utilize this for a domain schedule for the 2d conv nchw fchw
* update implementAssignLayout to take the domain schedule and pass that into the loop generation string
* add syntax and folder tests

PiperOrigin-RevId: 906428225

Author: asraa

lib/Dialect/TensorExt/IR/TensorExtOps.cpp

@@ -198,7 +198,29 @@ LogicalResult ConvertLayoutOp::verify() {
 }
 
 LogicalResult AssignLayoutOp::verify() {
-  return verifyLayoutMatchesType(getLayout(), getValue().getType(), *this);
+  LogicalResult layoutVerification =
+      verifyLayoutMatchesType(getLayout(), getValue().getType(), *this);
+  if (failed(layoutVerification)) {
+    return layoutVerification;
+  }
+
+  if (!getDomainSchedule().empty()) {
+    auto layout = dyn_cast<LayoutAttr>(getLayout());
+    if (!layout) {
+      return emitOpError()
+             << "requires LayoutAttr when domainSchedule is provided";
+    }
+    presburger::IntegerRelation rel = layout.getIntegerRelation();
+    for (int64_t idx : getDomainSchedule()) {
+      if (idx < 0 || idx >= rel.getNumDomainVars()) {
+        return emitOpError()
+               << "domainSchedule index " << idx << " is out of bounds [0, "
+               << rel.getNumDomainVars() << ")";
+      }
+    }
+  }
+
+  return success();
 }
 
 LogicalResult UnpackOp::verify() {

lib/Dialect/TensorExt/IR/TensorExtOps.td

@@ -110,10 +110,20 @@ def TensorExt_ConvertLayoutOp : TensorExt_Op<"convert_layout", [Pure, AllTypesMa
     plaintext masks (ciphertext-plaintext multiplications), and additions.
 
     This op is inserted by layout selection passes.
+
+    The optional `domainSchedule` attribute specifies which extra domain
+    dimensions should be part of the loop nest schedule when lowering. Generated
+    loops nests are only used when the layout conversion is used for a plaintext
+    layout assignment. See `tensor_ext.assign_layout` for more details.
   }];
 
   let assemblyFormat = "operands attr-dict `:` type($output)";
-  let arguments = (ins AnyType:$value, LayoutLike:$from_layout, LayoutLike:$to_layout);
+  let arguments = (ins
+      AnyType:$value,
+      LayoutLike:$from_layout,
+      LayoutLike:$to_layout,
+      DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$domainSchedule
+  );
   let results = (outs AnyType:$output);
   let hasVerifier = 1;
   let hasFolder = 1;
@@ -130,10 +140,20 @@ def TensorExt_AssignLayoutOp : TensorExt_Op<"assign_layout", [Pure, AllTypesMatc
     lowered as ciphertext-plaintext ops.
 
     This op is inserted by layout selection passes.
+
+    The optional `domainSchedule` attribute specifies which domain dimensions
+    should be part of the loop nest schedule when lowering. The range dimensions
+    are always included. Adding extra domain indices to the loop schedule will
+    help ISL generate a loop nest more efficiently, but comes at the tradeoff of
+    a slower generated loop nest.
   }];
 
   let assemblyFormat = "operands attr-dict `:` type($output)";
-  let arguments = (ins AnyType:$value, LayoutLike:$layout);
+  let arguments = (ins
+      AnyType:$value,
+      LayoutLike:$layout,
+      DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$domainSchedule
+  );
   let results = (outs AnyType:$output);
   let hasVerifier = 1;
 }

lib/Dialect/TensorExt/Transforms/Patterns.cpp

@@ -5,9 +5,10 @@
 #include "lib/Dialect/TensorExt/IR/TensorExtDialect.h"
 #include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "lib/Utils/AttributeUtils.h"
-#include "mlir/include/mlir/IR/PatternMatch.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/Value.h"         // from @llvm-project
-#include "mlir/include/mlir/Support/LLVM.h"     // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/PatternMatch.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 {
 namespace heir {
@@ -33,8 +34,15 @@ LogicalResult FoldConvertLayoutIntoAssignLayoutPattern::matchAndRewrite(
         continue;
       }
 
-      auto newOp = rewriter.replaceOpWithNewOp<AssignLayoutOp>(
-          convertLayoutOp, op.getValue(), convertLayoutOp.getToLayout());
+      DenseI64ArrayAttr domainSchedule =
+          !convertLayoutOp.getDomainSchedule().empty()
+              ? convertLayoutOp.getDomainScheduleAttr()
+              : op.getDomainScheduleAttr();
+
+      auto newOp = AssignLayoutOp::create(
+          rewriter, convertLayoutOp.getLoc(), op.getValue(),
+          convertLayoutOp.getToLayout(), domainSchedule);
+      rewriter.replaceOp(convertLayoutOp, newOp.getResult());
       // Ensure the newOp has its layout attribute properly set
       setAttributeAssociatedWith(newOp.getResult(),
                                  TensorExtDialect::kLayoutAttrName,

lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp

@@ -10,6 +10,7 @@
 #include "lib/Utils/Layout/Utils.h"
 #include "llvm/include/llvm/ADT/DynamicAPInt.h"     // from @llvm-project
 #include "llvm/include/llvm/ADT/STLExtras.h"        // from @llvm-project
+#include "llvm/include/llvm/ADT/SmallVector.h"      // from @llvm-project
 #include "llvm/include/llvm/Support/Debug.h"        // from @llvm-project
 #include "llvm/include/llvm/Support/raw_ostream.h"  // from @llvm-project
 #include "mlir/include/mlir/Analysis/FlatLinearValueConstraints.h"  // from @llvm-project
@@ -48,7 +49,8 @@ static std::string printRelation(const IntegerRelation& rel) {
 static FailureOr<Value> implementAssignLayoutNew(
     Value input, LayoutAttr layout, int64_t ciphertextSize,
     ImplicitLocOpBuilder& builder,
-    const std::function<void(Operation*)>& createdOpCallback) {
+    const std::function<void(Operation*)>& createdOpCallback,
+    ArrayRef<int64_t> domainSchedule = {}) {
   IntegerRelation rel = layout.getIntegerRelation();
 
   RankedTensorType dataSemanticType =
@@ -106,7 +108,10 @@ static FailureOr<Value> implementAssignLayoutNew(
   TypedValue<RankedTensorType> ciphertextTensor =
       cast<TypedValue<RankedTensorType>>(zeroOp.getResult());
   MLIRLoopNestGenerator generator(builder, createdOpCallback);
-  auto loopNestCstr = generateLoopNestAsCStr(rel);
+
+  SmallVector<int> domainIndices = llvm::to_vector(llvm::map_range(
+      domainSchedule, [](int64_t idx) { return static_cast<int>(idx); }));
+  auto loopNestCstr = generateLoopNestAsCStr(rel, domainIndices);
   if (failed(loopNestCstr)) {
     return builder.emitError() << "Failed to generate loop nest for relation "
                                << printRelation(rel);
@@ -141,7 +146,8 @@ static FailureOr<Value> implementAssignLayoutNew(
         auto inserted = tensor::InsertOp::create(builder, loc, extracted,
                                                  iterArgs[0], insertIndices);
         return scf::ValueVector({inserted});
-      });
+      },
+      domainIndices);
   if (failed(loop)) {
     return builder.emitError() << "Failed to generate loop nest for relation "
                                << printRelation(rel);
@@ -293,11 +299,12 @@ static FailureOr<Value> implementAssignLayoutPermutation(
 FailureOr<Value> implementAssignLayout(
     Value input, Attribute layout, int64_t ciphertextSize,
     ImplicitLocOpBuilder& builder,
-    const std::function<void(Operation*)>& createdOpCallback) {
+    const std::function<void(Operation*)>& createdOpCallback,
+    ArrayRef<int64_t> domainSchedule) {
   OpBuilder::InsertionGuard guard(builder);
   if (LayoutAttr layoutAttr = dyn_cast<LayoutAttr>(layout)) {
     return implementAssignLayoutNew(input, layoutAttr, ciphertextSize, builder,
-                                    createdOpCallback);
+                                    createdOpCallback, domainSchedule);
   } else if (DenseIntElementsAttr elementAttr =
                  dyn_cast<DenseIntElementsAttr>(layout)) {
     return implementAssignLayoutPermutation(input, elementAttr, ciphertextSize,

lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.h

@@ -23,7 +23,8 @@ namespace heir {
 FailureOr<Value> implementAssignLayout(
     Value input, Attribute layout, int64_t ciphertextSize,
     ImplicitLocOpBuilder& builder,
-    const std::function<void(Operation*)>& createdOpCallback);
+    const std::function<void(Operation*)>& createdOpCallback,
+    ArrayRef<int64_t> domainSchedule = {});
 
 // Lower tensor_ext.unpack. Returns the final value produced by the unpacking
 // implementation. Applies createdOpCallback to each created operation.

lib/Transforms/LayoutOptimization/Patterns.cpp

@@ -31,7 +31,8 @@ LogicalResult tryFoldLayoutConversionIntoPrevious(
     auto newFrom = priorConversion.getFromLayoutAttr();
     auto newTo = op.getToLayoutAttr();
     auto newConversion = ConvertLayoutOp::create(
-        rewriter, op.getLoc(), priorConversion.getValue(), newFrom, newTo);
+        rewriter, op.getLoc(), priorConversion.getValue(), newFrom, newTo,
+        op.getDomainScheduleAttr());
     newConversion->setAttr(kLayoutAttrName, newTo);
 
     rewriter.replaceAllUsesWith(op, newConversion);
@@ -45,8 +46,12 @@ LogicalResult tryFoldLayoutConversionIntoPrevious(
   if (auto priorAssignment = op.getValue().getDefiningOp<AssignLayoutOp>()) {
     // merge the conversion into the assignment return success();
     auto newLayout = op.getToLayoutAttr();
-    auto newAssign = AssignLayoutOp::create(
-        rewriter, op.getLoc(), priorAssignment.getValue(), newLayout);
+    auto domainSchedule = !op.getDomainSchedule().empty()
+                              ? op.getDomainScheduleAttr()
+                              : priorAssignment.getDomainScheduleAttr();
+    auto newAssign = AssignLayoutOp::create(rewriter, op.getLoc(),
+                                            priorAssignment.getValue(),
+                                            newLayout, domainSchedule);
     newAssign->setAttr(kLayoutAttrName, newLayout);
 
     rewriter.replaceAllUsesWith(op, newAssign);

lib/Transforms/LayoutPropagation/LayoutPropagation.cpp

@@ -111,11 +111,13 @@ void debugAssignLayout(Value value, LayoutAttr layout) {
 
 std::pair<Value, LayoutAttr> convertToLayout(
     MLIRContext* ctx, mlir::IRRewriter& builder, Operation* op, Value value,
-    LayoutAttr oldLayout, const IntegerRelation& newRelation) {
+    LayoutAttr oldLayout, const IntegerRelation& newRelation,
+    ArrayRef<int64_t> domainSchedule = {}) {
   builder.setInsertionPoint(op);
   LayoutAttr layoutAttr = LayoutAttr::getFromIntegerRelation(ctx, newRelation);
   ConvertLayoutOp convertLayoutOp = ConvertLayoutOp::create(
-      builder, op->getLoc(), value, oldLayout, layoutAttr);
+      builder, op->getLoc(), value, oldLayout, layoutAttr,
+      builder.getDenseI64ArrayAttr(domainSchedule));
   convertLayoutOp->setAttr(tensor_ext::TensorExtDialect::kLayoutAttrName,
                            layoutAttr);
   Value toReplace = convertLayoutOp.getResult();
@@ -722,9 +724,13 @@ LogicalResult LayoutPropagation::visitOperation(Conv2DNchwFchwOp op) {
                                "inserting layout conversion.\n");
 
     // Insert a layout conversion op to make the matrix layout expanded and
-    // squat diagonal
+    // squat diagonal. The added domain schedule ensures ISL can efficiently
+    // generate a loop nest implementing the layout. However, the choice of
+    // which domain indices to include is arbitrary (so long as ISL remains
+    // fast).
     auto [toReplace, newFilterLayoutAttr] = convertToLayout(
-        ctx, builder, op, filter, filterLayout, convRelation.value());
+        ctx, builder, op, filter, filterLayout, convRelation.value(),
+        /*domainSchedule=*/{0, 1});
     debugAssignLayout(toReplace, newFilterLayoutAttr);
     assignedLayouts.insert({toReplace, newFilterLayoutAttr});
   }

tests/Dialect/TensorExt/IR/ops.mlir

@@ -32,3 +32,13 @@ func.func @test_halevi_shoup_reduction(%0: tensor<16xi32>, %1: tensor<16x16xi32>
   %2 = tensor_ext.rotate_and_reduce %0, %1 {period = 1 : index, steps = 16 : index} : (tensor<16xi32>, tensor<16x16xi32>) -> tensor<16xi32>
   return %2 : tensor<16xi32>
 }
+
+func.func @test_convert_layout_with_schedule(%0: tensor<16x16xi32>) -> tensor<16x16xi32> {
+  %1 = tensor_ext.convert_layout %0 {from_layout = #layout1, to_layout = #layout2, domainSchedule = array<i64: 0, 1>} : tensor<16x16xi32>
+  return %1 : tensor<16x16xi32>
+}
+
+func.func @test_assign_layout_with_schedule(%0: tensor<16x16xi32>) -> tensor<16x16xi32> {
+  %1 = tensor_ext.assign_layout %0 {layout = #layout1, domainSchedule = array<i64: 1, 0>} : tensor<16x16xi32>
+  return %1 : tensor<16x16xi32>
+}

tests/Transforms/fold_convert_layout_to_assign_layout/fold.mlir

@@ -24,3 +24,27 @@ func.func @fold_multiple(%arg0 : tensor<16x16xi16>) -> (tensor<16x16xi16>, tenso
   %2 = tensor_ext.convert_layout %0 {from_layout = #row_major_matrix, to_layout = #col_major_matrix2} : tensor<16x16xi16>
   return %1, %2 : tensor<16x16xi16>, tensor<16x16xi16>
 }
+
+// CHECK: @assign_layout_with_schedule
+func.func @assign_layout_with_schedule(%arg0 : tensor<16x16xi16>) -> tensor<16x16xi16> {
+  // CHECK: tensor_ext.assign_layout {{.*}}domainSchedule = array<i64: 1, 0>
+  %0 = tensor_ext.assign_layout %arg0 {layout = #row_major_matrix, domainSchedule = array<i64: 1, 0>} : tensor<16x16xi16>
+  %1 = tensor_ext.convert_layout %0 {from_layout = #row_major_matrix, to_layout = #col_major_matrix} : tensor<16x16xi16>
+  return %1 : tensor<16x16xi16>
+}
+
+// CHECK: @convert_layout_with_schedule
+func.func @convert_layout_with_schedule(%arg0 : tensor<16x16xi16>) -> tensor<16x16xi16> {
+  // CHECK: tensor_ext.assign_layout {{.*}}domainSchedule = array<i64: 0, 1>
+  %0 = tensor_ext.assign_layout %arg0 {layout = #row_major_matrix} : tensor<16x16xi16>
+  %1 = tensor_ext.convert_layout %0 {from_layout = #row_major_matrix, to_layout = #col_major_matrix, domainSchedule = array<i64: 0, 1>} : tensor<16x16xi16>
+  return %1 : tensor<16x16xi16>
+}
+
+// CHECK: @both_with_schedule
+func.func @both_with_schedule(%arg0 : tensor<16x16xi16>) -> tensor<16x16xi16> {
+  // CHECK: tensor_ext.assign_layout {{.*}}domainSchedule = array<i64: 0, 1>
+  %0 = tensor_ext.assign_layout %arg0 {layout = #row_major_matrix, domainSchedule = array<i64: 1, 0>} : tensor<16x16xi16>
+  %1 = tensor_ext.convert_layout %0 {from_layout = #row_major_matrix, to_layout = #col_major_matrix, domainSchedule = array<i64: 0, 1>} : tensor<16x16xi16>
+  return %1 : tensor<16x16xi16>
+}