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<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
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<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
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<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 @@
LinalgOp promoteSubViewOperands(OpBuilder &b, LinalgOp op,
llvm::SetVector<Value> 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<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 @@
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<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 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<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 @@
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 @@
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 @@
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 @@
// 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 @@
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<?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
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<And<[
+ PreconditionPromoteSubviewsLinalgOp,
+ HasOperandsOfType<"SubViewOp">,
+ HasLinalgTransformMarker<"_promote_views_aligned_">]>>
+ )]>;
+
#endif // TEST_LINALG_TRANSFORMS_PATTERNS