[Mlir-commits] [mlir] 128d727 - [mlir] [linalg] Specify alignment during promotion.
Alex Zinenko
llvmlistbot at llvm.org
Tue Apr 21 09:26:24 PDT 2020
Author: Pierre Oechsel
Date: 2020-04-21T18:13:40+02:00
New Revision: 128d72751fd33b3602447e639c6ee09643a4589c
URL: https://github.com/llvm/llvm-project/commit/128d72751fd33b3602447e639c6ee09643a4589c
DIFF: https://github.com/llvm/llvm-project/commit/128d72751fd33b3602447e639c6ee09643a4589c.diff
LOG: [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
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td
mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp
mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
mlir/test/Dialect/Linalg/transform-patterns.mlir
mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td
index 2eaed14e8377..a51352cd4d0e 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td
+++ b/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
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h
index e7a8925f746b..78d588aaf00b 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h
+++ b/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
diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index 9f83707d3ef5..7a1e18398cad 100644
--- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ b/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
diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp
index e96ee2780252..15454a4d6bd4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp
+++ b/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));
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index 09e1bebb354a..5bde420df076 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/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);
}
});
diff --git a/mlir/test/Dialect/Linalg/transform-patterns.mlir b/mlir/test/Dialect/Linalg/transform-patterns.mlir
index e4320d803fcb..d230aa993611 100644
--- a/mlir/test/Dialect/Linalg/transform-patterns.mlir
+++ b/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
diff --git a/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td b/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td
index 8444f4cc3dc4..313e2f8171a8 100644
--- a/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td
+++ b/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
More information about the Mlir-commits
mailing list