diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td --- a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td @@ -115,8 +115,9 @@ def PromoteSubviewsLinalgOp : NativeCodeCall< "promoteSubviewsLinalgOp($_builder, op)">; -class PromoteSelectedSubviewsLinalgOp operands, string marker=""> : +class PromoteSelectedSubviewsLinalgOp operands, string marker="", + int alignment=0> : NativeCodeCall<"promoteSelectedSubviewsLinalgOpAndSetMarker($_builder, op, {" # - StrJoinInt.result # "}, \"" # marker # "\")">; + StrJoinInt.result # "}, \"" # marker # "\", " # alignment # ")">; #endif // LINALG_TRANSFORMS diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h @@ -123,12 +123,14 @@ /// 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 promoteSelectedSubviewsLinalgOpAndSetMarker( PatternRewriter &rewriter, Operation *op, - ArrayRef operandIndicesToPromote, StringRef linalgMarker = ""); + ArrayRef operandIndicesToPromote, StringRef linalgMarker = "", + int64_t alignment = 0); } // namespace linalg } // namespace mlir diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h --- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h @@ -176,7 +176,8 @@ /// full and partial views indexing into the buffer. SmallVector promoteSubViews(OpBuilder &b, Location loc, ArrayRef 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 @@ LinalgOp promoteSubViewOperands(OpBuilder &b, LinalgOp op, llvm::SetVector subViews, bool dynamicBuffers = false, + int64_t alignment = 0, OperationFolder *folder = nullptr); } // namespace linalg diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp @@ -349,7 +349,8 @@ SmallVector mlir::linalg::promoteSelectedSubviewsLinalgOpAndSetMarker( PatternRewriter &rewriter, Operation *op, - ArrayRef operandIndicesToPromote, StringRef linalgMarker) { + ArrayRef operandIndicesToPromote, StringRef linalgMarker, + int64_t alignment) { LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: Promote subviews for linalg op: " << *op << ":\n"); @@ -372,7 +373,8 @@ 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)); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp @@ -65,16 +65,21 @@ } 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(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 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 @@ partialSizes.push_back(folded_std_dim(folder, subView, rank)); } SmallVector 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 @@ SmallVector mlir::linalg::promoteSubViews(OpBuilder &b, Location loc, ArrayRef subViews, bool dynamicBuffers, - OperationFolder *folder) { + int64_t alignment, OperationFolder *folder) { if (subViews.empty()) return {}; @@ -142,8 +148,8 @@ DenseMap promotionInfoMap; for (auto v : subViews) { SubViewOp subView = cast(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 @@ LinalgOp mlir::linalg::promoteSubViewOperands(OpBuilder &b, LinalgOp op, SetVector subViews, bool dynamicBuffers, + int64_t alignment, OperationFolder *folder) { assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); @@ -189,8 +196,9 @@ // 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 opViews; opViews.reserve(op.getNumInputsAndOutputs()); SmallVector, 8> writebackViews; @@ -248,7 +256,7 @@ 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); } }); diff --git a/mlir/test/Dialect/Linalg/transform-patterns.mlir b/mlir/test/Dialect/Linalg/transform-patterns.mlir --- a/mlir/test/Dialect/Linalg/transform-patterns.mlir +++ b/mlir/test/Dialect/Linalg/transform-patterns.mlir @@ -444,3 +444,25 @@ // CHECK-NOT: linalg.copy(%[[s1]], %[[l1]]) : memref, memref // CHECK-NOT: linalg.copy(%[[s2]], %[[l2]]) : memref, memref^ // CHECK: linalg.matmul(%[[v0]], %[[s1]], %[[s2]]) : memref, memref, memref + +func @aligned_promote_fill(%arg0: memref) { + %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 to memref + linalg.fill(%3, %cf) { __internal_linalg_transform__ = "_promote_views_aligned_"} + : memref, f32 + return +} +// CHECK-LABEL: func @aligned_promote_fill +// CHECK: %[[cf:.*]] = constant {{.*}} : f32 +// CHECK: %[[s0:.*]] = subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref to memref +// CHECK: %[[a0:.*]] = alloc({{%.*}}) {alignment = 32 : i64} : memref +// CHECK: %[[v0:.*]] = std.view %[[a0]][][{{%.*}}, {{%.*}}] : memref to memref +// CHECK: %[[l0:.*]] = subview %[[v0]][{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref to memref +// CHECK: linalg.fill(%[[v0]], {{%.*}}) : memref, f32 +// CHECK: linalg.copy(%[[s0]], %[[l0]]) : memref, memref +// CHECK: linalg.fill(%[[v0]], %[[cf]]) : memref, f32 diff --git a/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td b/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td --- a/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td +++ b/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td @@ -157,4 +157,12 @@ HasLinalgTransformMarker<"_promote_first_view_">]>> )]>; +def : Pat<(FillOp:$op $_, $_), + (PromoteSelectedSubviewsLinalgOp<[0], "aligned_promotion", 32>), + [(Constraint, + HasLinalgTransformMarker<"_promote_views_aligned_">]>> + )]>; + #endif // TEST_LINALG_TRANSFORMS_PATTERNS