[mlir] [linalg] Specify alignment during promotion.

[poechsel]

Summary

The buffer allocated by a promotion can be subject to other transformations afterward. For example it could be vectorized, in which case it is needed to ensure that this buffer is memory-aligned.

[ftynse]

Please provide a longer commit description (https://mlir.llvm.org/getting_started/Contributing/#commit-messages)

mlir/test/Dialect/Linalg/transform-patterns.mlir

448

I wonder if we could have a significantly simpler test here that only looks at promotion rather than replicating matmul every time. There seems to be a lot of things that are tested elsewhere and will lead to additional churn when the transformation is modified.

Could you maybe try linalg.fill instead?

[poechsel]

Simplified test.

[nicolasvasilache]

Thanks Pierre!

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td

 //===----------------------------------------------------------------------===//
 // Linalg promote subview operands precondition and DRR.
 //===----------------------------------------------------------------------===//
 def PreconditionPromoteSubviewsLinalgOp : CPred<
   "succeeded(promoteSubviewsLinalgOpPrecondition(op))">;
 def PromoteSubviewsLinalgOp : NativeCodeCall<
   "promoteSubviewsLinalgOp($_builder, op)">;
 
-class PromoteSelectedSubviewsLinalgOp<list<int> operands, string marker=""> :
+class PromoteSelectedSubviewsLinalgOp<list<int> operands, string marker="",
+                                      int alignment=0> :
   NativeCodeCall<"promoteSelectedSubviewsLinalgOpAndSetMarker($_builder, op, {" #
-    StrJoinInt<operands>.result # "}, \"" # marker # "\")">;
+    StrJoinInt<operands>.result # "}, \"" # marker # "\", " # alignment # ")">;
 
 #endif // LINALG_TRANSFORMS

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h

 
 /// Promote std.subviews feeding linalg operations.
 LogicalResult promoteSubviewsLinalgOpPrecondition(Operation *op);
 SmallVector<Value, 0> promoteSubviewsLinalgOp(PatternRewriter &rewriter,
                                               Operation *op);
 
 /// Similar to `promoteSubviewsLinalgOp` but only tries to promote
 /// the views corresponding to the operands specified in
-/// `operandIndicesToPromote`.
+/// `operandIndicesToPromote`. Generated allocations are memory-aligned
+/// according to the `alignment` parameter.
 /// If linalgMarker is specified and the transformation is successfull
 /// sets the attribute `kLinalgTransformMarker` to `linalgMarker`.
 SmallVector<Value, 0> promoteSelectedSubviewsLinalgOpAndSetMarker(
     PatternRewriter &rewriter, Operation *op,
-    ArrayRef<int64_t> operandIndicesToPromote, StringRef linalgMarker = "");
+    ArrayRef<int64_t> operandIndicesToPromote, StringRef linalgMarker = "",
+    int64_t alignment = 0);
 } // namespace linalg
 } // namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_LINALGTRANSFORMS_H_

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 ///   3. Take a partial slice of the full view in step 2. and copy into it.
 /// Infers statically sized buffers from subViews unless `dynamicBuffers` is
 /// true.
 ///
 /// Returns a list of PromotionInfo which hold the promoted buffer and the
 /// full and partial views indexing into the buffer.
 SmallVector<PromotionInfo, 8>
 promoteSubViews(OpBuilder &b, Location loc, ArrayRef<Value> subViews,
-                bool dynamicBuffers = false, OperationFolder *folder = nullptr);
+                bool dynamicBuffers = false, int64_t alignment = 0,
+                OperationFolder *folder = nullptr);
 
 /// Returns all the operands of `linalgOp` that are not views.
 /// Asserts that these operands are value types to allow transformations like
 /// tiling to just use the values when cloning `linalgOp`.
 SmallVector<Value, 4> getAssumedNonViewOperands(LinalgOp linalgOp);
 
 /// Apply the permutation defined by `permutation` to `inVec`.
 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.

 /// Prepares the SubView promotion later performed by `promoteSubViews`
 /// (where most of the transformation happens). It arranges the new
 /// operands for `LinalgOp op` and deallocates the new buffer(s)
 /// It is the entry point for declarative transformation
 /// Returns the cloned `LinalgOp` with the new operands
 LinalgOp promoteSubViewOperands(OpBuilder &b, LinalgOp op,
                                 llvm::SetVector<Value> subViews,
                                 bool dynamicBuffers = false,
+                                int64_t alignment = 0,
                                 OperationFolder *folder = nullptr);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_UTILS_H_

mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp

mlir::linalg::promoteSubviewsLinalgOp(PatternRewriter &rewriter,
   toPromote.reserve(nBuffers);
   for (int64_t i = 0; i < nBuffers; ++i)
     toPromote.push_back(i);
   return promoteSelectedSubviewsLinalgOpAndSetMarker(rewriter, op, toPromote);
 }
 
 SmallVector<Value, 0> mlir::linalg::promoteSelectedSubviewsLinalgOpAndSetMarker(
     PatternRewriter &rewriter, Operation *op,
-    ArrayRef<int64_t> operandIndicesToPromote, StringRef linalgMarker) {
+    ArrayRef<int64_t> operandIndicesToPromote, StringRef linalgMarker,
+    int64_t alignment) {
   LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: Promote subviews for linalg op: "
                     << *op << ":\n");
 
   assert(succeeded(promoteSubviewsLinalgOpPrecondition(op)) &&
          "DRR failure case must be a precondition");
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op)) {
     // TODO(ntv): add a level of indirection to linalg.generic.
     if (convOp.padding())
       llvm_unreachable("Unexpected conv with padding");
   }
 
   LinalgOp linOp = cast<LinalgOp>(op);
   assert(linOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   SetVector<Value> subViews;
   for (int64_t index : operandIndicesToPromote)
     if (auto sv =
             dyn_cast_or_null<SubViewOp>(linOp.getBuffer(index).getDefiningOp()))
       subViews.insert(sv);
 
   if (!subViews.empty()) {
-    auto newOp = promoteSubViewOperands(rewriter, linOp, subViews);
+    auto newOp =
+        promoteSubViewOperands(rewriter, linOp, subViews, false, alignment);
     if (!linalgMarker.empty())
       newOp.setAttr(LinalgTransforms::kLinalgTransformMarker,
                     rewriter.getStringAttr(linalgMarker));
     return {};
   }
   llvm_unreachable("DRR failure case must be a precondition");
 }

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc,
   for (auto e : affineMinOp.getAffineMap().getResults())
     if (auto cst = e.dyn_cast<AffineConstantExpr>())
       minConst = std::min(minConst, cst.getValue());
   assert(minConst != std::numeric_limits<int64_t>::max());
   return b.create<ConstantIndexOp>(loc, minConst);
 }
 
 static Value allocBuffer(Type elementType, Value size, bool dynamicBuffers,
-                         OperationFolder *folder) {
+                         OperationFolder *folder, int64_t alignment = 0) {
   auto *ctx = size.getContext();
   auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
+  IntegerAttr alignment_attr;
+  if (alignment)
+    alignment_attr = IntegerAttr::get(IntegerType::get(64, ctx), alignment);
   if (!dynamicBuffers)
     if (auto cst = dyn_cast_or_null<ConstantIndexOp>(size.getDefiningOp()))
       return std_alloc(
-          MemRefType::get(width * cst.getValue(), IntegerType::get(8, ctx)));
+          MemRefType::get(width * cst.getValue(), IntegerType::get(8, ctx)), {},
+          alignment_attr);
   Value mul =
       folded_std_muli(folder, folded_std_constant_index(folder, width), size);
-  return std_alloc(MemRefType::get(-1, IntegerType::get(8, ctx)), mul);
+  return std_alloc(MemRefType::get(-1, IntegerType::get(8, ctx)), mul,
+                   alignment_attr);
 }
 
 // Performs promotion of a `subView` into a local buffer of the size of the
 // *ranges* of the `subView`. This produces a buffer whose size may be bigger
 // than the actual size of the `subView` at the boundaries.
 // This is related to the full/partial tile problem.
 // Returns a PromotionInfo containing a `buffer`, `fullLocalView` and
 // `partialLocalView` such that:
 //   * `buffer` is always the size of the full tile.
 //   * `fullLocalView` is a dense contiguous view into that buffer.
 //   * `partialLocalView` is a dense non-contiguous slice of `fullLocalView`
 //     that corresponds to the size of `subView` and accounting for boundary
 //     effects.
 // The point of the full tile buffer is that constant static tile sizes are
 // folded and result in a buffer type with statically known size and alignment
 // properties.
 // To account for general boundary effects, padding must be performed on the
 // boundary tiles. For now this is done with an unconditional `fill` op followed
 // by a partial `copy` op.
 static PromotionInfo promoteFullTileBuffer(OpBuilder &b, Location loc,
                                            SubViewOp subView,
                                            bool dynamicBuffers,
+                                           int64_t alignment,
                                            OperationFolder *folder) {
   auto zero = folded_std_constant_index(folder, 0);
   auto one = folded_std_constant_index(folder, 1);
 
   auto viewType = subView.getType();
   auto rank = viewType.getRank();
   Value allocSize = one;
   SmallVector<Value, 8> fullSizes, partialSizes;
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
   for (auto en : llvm::enumerate(subView.getRanges())) {
     auto rank = en.index();
     auto rangeValue = en.value();
     // Try to extract a tight constant
     Value size = extractSmallestConstantBoundingSize(b, loc, rangeValue.size);
     allocSize = folded_std_muli(folder, allocSize, size).getValue();
     fullSizes.push_back(size);
     partialSizes.push_back(folded_std_dim(folder, subView, rank));
   }
   SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1);
-  auto buffer =
-      allocBuffer(viewType.getElementType(), allocSize, dynamicBuffers, folder);
+  auto buffer = allocBuffer(viewType.getElementType(), allocSize,
+                            dynamicBuffers, folder, alignment);
   auto fullLocalView = folded_std_view(
       folder, MemRefType::get(dynSizes, viewType.getElementType()), buffer,
       fullSizes);
   SmallVector<Value, 4> zeros(fullSizes.size(), zero);
   SmallVector<Value, 4> ones(fullSizes.size(), one);
   auto partialLocalView =
       folded_std_subview(folder, fullLocalView, zeros, partialSizes, ones);
   return PromotionInfo{buffer, fullLocalView, partialLocalView};
 }
 
 SmallVector<PromotionInfo, 8>
 mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
                               ArrayRef<Value> subViews, bool dynamicBuffers,
-                              OperationFolder *folder) {
+                              int64_t alignment, OperationFolder *folder) {
   if (subViews.empty())
     return {};
 
   ScopedContext scope(b, loc);
   SmallVector<PromotionInfo, 8> res;
   res.reserve(subViews.size());
   DenseMap<Value, PromotionInfo> promotionInfoMap;
   for (auto v : subViews) {
     SubViewOp subView = cast<SubViewOp>(v.getDefiningOp());
-    auto promotionInfo =
-        promoteFullTileBuffer(b, loc, subView, dynamicBuffers, folder);
+    auto promotionInfo = promoteFullTileBuffer(b, loc, subView, dynamicBuffers,
+                                               alignment, folder);
     promotionInfoMap.insert(std::make_pair(subView.getResult(), promotionInfo));
     res.push_back(promotionInfo);
   }
 
   for (auto v : subViews) {
     SubViewOp subView = cast<SubViewOp>(v.getDefiningOp());
     auto info = promotionInfoMap.find(v);
     if (info == promotionInfoMap.end())
linalg_copy(cast<SubViewOp>(v.getDefiningOp()),
                 info->second.partialLocalView);
   }
   return res;
 }
 
 LinalgOp mlir::linalg::promoteSubViewOperands(OpBuilder &b, LinalgOp op,
                                               SetVector<Value> subViews,
                                               bool dynamicBuffers,
+                                              int64_t alignment,
                                               OperationFolder *folder) {
   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
     // TODO(ntv): add a level of indirection to linalg.generic.
     if (convOp.padding())
       llvm_unreachable("Unexpected conv with padding");
   }
 
   // 1. Promote the specified views and use them in the new op.
   ScopedContext scope(b, op.getLoc());
-  auto promotedBufferAndViews = promoteSubViews(
-      b, op.getLoc(), subViews.getArrayRef(), dynamicBuffers, folder);
+  auto promotedBufferAndViews =
+      promoteSubViews(b, op.getLoc(), subViews.getArrayRef(), dynamicBuffers,
+                      alignment, folder);
   SmallVector<Value, 8> opViews;
   opViews.reserve(op.getNumInputsAndOutputs());
   SmallVector<std::pair<Value, Value>, 8> writebackViews;
   writebackViews.reserve(subViews.size());
   unsigned promotedIdx = 0;
   for (auto view : op.getInputsAndOutputBuffers()) {
     if (subViews.count(view) != 0) {
       opViews.push_back(promotedBufferAndViews[promotedIdx].fullLocalView);
f.walk([dynamicBuffers, &folder, &toErase](LinalgOp op) {
     // and integer buffers can be promoted.
     SetVector<Value> subViews;
     OpBuilder b(op);
     for (auto it : op.getInputsAndOutputBuffers())
       if (auto sv = dyn_cast_or_null<SubViewOp>(it.getDefiningOp()))
         if (sv.getType().getElementType().isSignlessIntOrFloat())
           subViews.insert(sv);
     if (!subViews.empty()) {
-      promoteSubViewOperands(b, op, subViews, dynamicBuffers, &folder);
+      promoteSubViewOperands(b, op, subViews, dynamicBuffers, 0, &folder);
       toErase.push_back(op);
     }
   });
   for (auto op : toErase)
     op.erase();
 }
 
 namespace {

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // CHECK-NOT:     %[[l0:.*]] = subview %[[v1]][{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32> to memref<?x?xf32, #[[STRIDED_2D]]>
 // CHECK-NOT:     %[[a2:.*]] = alloc({{%.*}}) : memref<?xi8>
 // CHECK-NOT:     %[[v2:.*]] = std.view %[[a2]][][{{%.*}}, {{%.*}}] : memref<?xi8> to memref<?x?xf32>
 // CHECK-NOT:     %[[l0:.*]] = subview %[[v2]][{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32> to memref<?x?xf32, #[[STRIDED_2D]]>
 // CHECK:         linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
 // CHECK-NOT:     linalg.copy(%[[s1]], %[[l1]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
 // CHECK-NOT:     linalg.copy(%[[s2]], %[[l2]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>^
 // CHECK:         linalg.matmul(%[[v0]], %[[s1]], %[[s2]]) : memref<?x?xf32>, memref<?x?xf32, #[[STRIDED_2D]]>, memref<?x?xf32, #[[STRIDED_2D]]>
+
+func @aligned_promote_matmul(%arg0: memref<?x?xf32, offset: ?, strides: [?, 1]>,
+                             %arg1: memref<?x?xf32, offset: ?, strides: [?, 1]>,
+                             %arg2: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
+  %c2000 = constant 2000 : index
+  %c3000 = constant 3000 : index
+  %c4000 = constant 4000 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = dim %arg0, 0 : memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %1 = dim %arg0, 1 : memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %2 = dim %arg1, 1 : memref<?x?xf32, offset: ?, strides: [?, 1]>
+  loop.for %arg3 = %c0 to %0 step %c2000 {
+    loop.for %arg4 = %c0 to %2 step %c3000 {
+      loop.for %arg5 = %c0 to %1 step %c4000 {
+        %3 = std.subview %arg0[%arg3, %arg5][%c2000, %c4000][%c1, %c1] :
+             memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %4 = std.subview %arg1[%arg5, %arg4][%c4000, %c3000][%c1, %c1] :
+             memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %arg2[%arg3, %arg4][%c2000, %c3000][%c1, %c1] :
+             memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%3, %4, %5) {__internal_linalg_transform__ = "_promote_views_aligned_"} :
+                      memref<?x?xf32, offset: ?, strides: [?, ?]>,
+                      memref<?x?xf32, offset: ?, strides: [?, ?]>,
+                      memref<?x?xf32, offset: ?, strides: [?, ?]>
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL: func @aligned_promote_matmul
+// CHECK:         loop.for {{.*}} = %[[c0]] to {{.*}} step %[[c2000]] {
+// CHECK:           loop.for {{.*}} = %[[c0]] to {{.*}} step %[[c3000]] {
+// CHECK:             loop.for {{.*}} = %[[c0]] to {{.*}} step %[[c4000]] {
+// CHECK:               %[[s0:.*]] = subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
+// CHECK:               %[[s1:.*]] = subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
+// CHECK:               %[[s2:.*]] = subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
+// CHECK:               %[[a0:.*]] = alloc({{%.*}}) {alignment = 32 : i64} : memref<?xi8>
+// CHECK:               %[[v0:.*]] = std.view %[[a0]][][{{%.*}}, {{%.*}}] : memref<?xi8> to memref<?x?xf32>
+// CHECK:               %[[l0:.*]] = subview %[[v0]][{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32> to memref<?x?xf32, #[[STRIDED_2D]]>
+// CHECK:               %[[a1:.*]] = alloc({{%.*}}) {alignment = 32 : i64} : memref<?xi8>
+// CHECK:               %[[v1:.*]] = std.view %[[a1]][][{{%.*}}, {{%.*}}] : memref<?xi8> to memref<?x?xf32>
+// CHECK:               %[[l1:.*]] = subview %[[v1]][{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32> to memref<?x?xf32, #[[STRIDED_2D]]>
+// CHECK:               %[[a2:.*]] = alloc({{%.*}}) {alignment = 32 : i64} : memref<?xi8>
+// CHECK:               %[[v2:.*]] = std.view %[[a2]][][{{%.*}}, {{%.*}}] : memref<?xi8> to memref<?x?xf32>
+// CHECK:               %[[l2:.*]] = subview %[[v2]][{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32> to memref<?x?xf32, #[[STRIDED_2D]]>
+// CHECK:               linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
+// CHECK:               linalg.copy(%[[s1]], %[[l1]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
+// CHECK:               linalg.copy(%[[s2]], %[[l2]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
+// CHECK:               linalg.matmul(%[[v0]], %[[v1]], %[[v2]]) : memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

 def : Pat<(MatmulOp:$op $_, $_, $_),
           (PromoteSelectedSubviewsLinalgOp<[0], "first_view_promotion">),
           [(Constraint<And<[
               PreconditionPromoteSubviewsLinalgOp,
               HasOperandsOfType<"SubViewOp">,
               HasLinalgTransformMarker<"_promote_first_view_">]>>
            )]>;
 
+def : Pat<(MatmulOp:$op $_, $_, $_),
+          (PromoteSelectedSubviewsLinalgOp<[0, 1, 2], "aligned_promotion", 32>),
+          [(Constraint<And<[
+              PreconditionPromoteSubviewsLinalgOp,
+              HasOperandsOfType<"SubViewOp">,
+              HasLinalgTransformMarker<"_promote_views_aligned_">]>>
+           )]>;
+
 #endif // TEST_LINALG_TRANSFORMS_PATTERNS