[mlir] [linalg] Specify alignment during promotion.

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.

Differential Revision: https://reviews.llvm.org/D78556

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

@@ -115,8 +115,9 @@ def PreconditionPromoteSubviewsLinalgOp : CPred<
 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

@@ -123,12 +123,14 @@ SmallVector<Value, 0> promoteSubviewsLinalgOp(PatternRewriter &rewriter,
 
 /// 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
 

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

@@ -176,7 +176,8 @@ struct PromotionInfo {
 /// 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
@@ -204,6 +205,7 @@ void applyPermutationToVector(SmallVector<T, N> &inVec,
 LinalgOp promoteSubViewOperands(OpBuilder &b, LinalgOp op,
                                 llvm::SetVector<Value> subViews,
                                 bool dynamicBuffers = false,
+                                int64_t alignment = 0,
                                 OperationFolder *folder = nullptr);
 
 } // namespace linalg

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

@@ -349,7 +349,8 @@ mlir::linalg::promoteSubviewsLinalgOp(PatternRewriter &rewriter,
 
 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");
 
@@ -372,7 +373,8 @@ SmallVector<Value, 0> mlir::linalg::promoteSelectedSubviewsLinalgOpAndSetMarker(
       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));

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

@@ -65,16 +65,21 @@ static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc,
 }
 
 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
@@ -97,6 +102,7 @@ static Value allocBuffer(Type elementType, Value size, bool dynamicBuffers,
 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);
@@ -117,8 +123,8 @@ static PromotionInfo promoteFullTileBuffer(OpBuilder &b, Location loc,
     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);
@@ -132,7 +138,7 @@ static PromotionInfo promoteFullTileBuffer(OpBuilder &b, Location loc,
 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 {};
 
@@ -142,8 +148,8 @@ mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
   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);
   }
@@ -178,6 +184,7 @@ mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
 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");
 
@@ -189,8 +196,9 @@ LinalgOp mlir::linalg::promoteSubViewOperands(OpBuilder &b, LinalgOp op,
 
   // 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;
@@ -248,7 +256,7 @@ static void promoteSubViews(FuncOp f, bool dynamicBuffers) {
         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);
     }
   });

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

@@ -444,3 +444,25 @@ func @promote_first_subview_matmul(%arg0: memref<?x?xf32, offset: ?, strides: [?
 // 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_fill(%arg0: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
+  %c2000 = constant 2000 : index
+  %c4000 = constant 4000 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %cf = constant 1.0 : f32
+  %3 = std.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] :
+ 	 memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+  linalg.fill(%3, %cf) { __internal_linalg_transform__ = "_promote_views_aligned_"}
+  	:  memref<?x?xf32, offset: ?, strides: [?, ?]>, f32
+  return
+}
+// CHECK-LABEL: func @aligned_promote_fill
+// CHECK:	  %[[cf:.*]] = constant {{.*}} : f32
+// CHECK:         %[[s0:.*]] = 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:         linalg.fill(%[[v0]], {{%.*}}) : memref<?x?xf32>, f32
+// CHECK:         linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
+// CHECK:         linalg.fill(%[[v0]], %[[cf]]) : memref<?x?xf32>, f32

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

@@ -157,4 +157,12 @@ def : Pat<(MatmulOp:$op $_, $_, $_),
               HasLinalgTransformMarker<"_promote_first_view_">]>>
            )]>;
 
+def : Pat<(FillOp:$op $_, $_),
+          (PromoteSelectedSubviewsLinalgOp<[0], "aligned_promotion", 32>),
+          [(Constraint<And<[
+              PreconditionPromoteSubviewsLinalgOp,
+              HasOperandsOfType<"SubViewOp">,
+              HasLinalgTransformMarker<"_promote_views_aligned_">]>>
+           )]>;
+
 #endif // TEST_LINALG_TRANSFORMS_PATTERNS