Skip to content

Commit 316e627

Browse files
author
gysit
committed
[mlir][linalg] Support the empty anchor op string when padding.
Add support for an empty anchor op string in vectorization. An empty anchor op string is useful after fusion when there are multiple different operations to vectorize. Depends On D114689 Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D114690
1 parent 3e32f82 commit 316e627

File tree

3 files changed

+25
-8
lines changed

3 files changed

+25
-8
lines changed

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -268,7 +268,6 @@ struct CodegenStrategy {
268268
vectorize(StringRef opName,
269269
LinalgTransformationFilter::FilterFunction f = nullptr,
270270
bool vectorizePadding = false) {
271-
assert(!opName.empty() && "expected an op name");
272271
transformationSequence.emplace_back(std::make_unique<Vectorize>(
273272
opName, linalg::LinalgVectorizationOptions(), f, vectorizePadding));
274273
return *this;

mlir/test/Dialect/Linalg/codegen-strategy.mlir

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" -split-input-file | FileCheck %s --check-prefix=CHECK-OUTER
33
// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 tile-interchange=1,2,0 generalize iterator-interchange=0,2,1" -split-input-file | FileCheck %s --check-prefix=CHECK-INTERCHANGE
44
// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 pad pack-paddings=1,1,0 hoist-paddings=3,3,0" -split-input-file | FileCheck %s --check-prefix=CHECK-PAD
5+
// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 fuse pad vectorize" -split-input-file | FileCheck %s --check-prefix=CHECK-FUSE
56

67
// CHECK-INTRINSIC: func @matmul(
78
// CHECK-OUTER: func @matmul(
@@ -56,3 +57,20 @@ func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<7
5657
%0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32>
5758
return %0 : tensor<72x72xf32>
5859
}
60+
61+
// -----
62+
63+
// CHECK-FUSE: func @matmul(
64+
func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
65+
66+
// Check the padding and vectorization applies to the fill operation due to the empty anchor op string.
67+
// CHECK-FUSE: %[[CST:.*]] = arith.constant dense<0.000000e+00>
68+
// CHECK-FUSE: vector.transfer_write %[[CST]]
69+
%cst = arith.constant 0.0 : f32
70+
%0 = linalg.fill(%cst, %arg0) : f32, tensor<72x72xf32> -> tensor<72x72xf32>
71+
72+
// Check the matmul is padded and vectorized despite the empty anchor op string.
73+
// CHECK-FUSE: vector.outerproduct
74+
%1 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%0: tensor<72x72xf32>) -> tensor<72x72xf32>
75+
return %1 : tensor<72x72xf32>
76+
}

mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -165,24 +165,24 @@ void TestLinalgCodegenStrategy::runStrategy(
165165
vector::VectorTransferSplit vectorTransferSplit) {
166166
assert(!anchorOpName.empty());
167167
CodegenStrategy strategy;
168-
StringRef genericOpName = GenericOp::getOperationName();
169168
strategy
170169
.tileAndFuseIf(fuse && !tileSizes.empty(), anchorOpName,
171170
tilingAndFusionOptions)
172171
.tileIf(!fuse && !tileSizes.empty(), anchorOpName, tilingOptions)
173-
.promoteIf(promote, anchorOpName,
172+
.promoteIf(!fuse && promote, anchorOpName,
174173
LinalgPromotionOptions()
175174
.setAlignment(16)
176175
.setUseFullTileBuffersByDefault(promoteFullTile))
177-
.tileIf(!registerTileSizes.empty(), anchorOpName, registerTilingOptions)
178-
.promoteIf(registerPromote, anchorOpName,
176+
.tileIf(!fuse && !registerTileSizes.empty(), anchorOpName,
177+
registerTilingOptions)
178+
.promoteIf(!fuse && registerPromote, anchorOpName,
179179
LinalgPromotionOptions()
180180
.setAlignment(16)
181181
.setUseFullTileBuffersByDefault(registerPromoteFullTile))
182-
.padIf(pad, anchorOpName, paddingOptions)
183-
.generalizeIf(generalize, anchorOpName)
182+
.padIf(pad, "", paddingOptions)
183+
.generalizeIf(generalize, "")
184184
.interchangeIf(!iteratorInterchange.empty(), iteratorInterchange)
185-
.vectorizeIf(vectorize, generalize ? genericOpName : anchorOpName)
185+
.vectorizeIf(vectorize, "")
186186
.vectorLowering(
187187
LinalgVectorLoweringOptions()
188188
.setVectorTransformsOptions(

0 commit comments

Comments
 (0)