Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a1d43c1

Browse files
bjacobbjacob
authored
[mlir][vector] Add Vector-dialect interleave-to-shuffle pattern, enable in VectorToSPIRV (#92012)
This is the second attempt at merging #91800, which bounced due to a linker error apparently caused by an undeclared dependency. `MLIRVectorToSPIRV` needed to depend on `MLIRVectorTransforms`. In fact that was a preexisting issue already flagged by the tool in https://discourse.llvm.org/t/ninja-can-now-check-for-missing-cmake-dependencies-on-generated-files/74344. Context: iree-org/iree#17346. Test IREE integrate showing it's fixing the problem it's intended to fix, i.e. it allows IREE to drop its local revert of #89131: iree-org/iree#17359 This is added to VectorToSPIRV because SPIRV doesn't currently handle `vector.interleave` (see motivating context above). This is limited to 1D, non-scalable vectors.
1 parent fd4b5f4 commit a1d43c1

8 files changed

Lines changed: 90 additions & 0 deletions

File tree

‎mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,20 @@ def ApplyLowerInterleavePatternsOp : Op<Transform_Dialect,
306306
let assemblyFormat = "attr-dict";
307307
}
308308

309+
def ApplyInterleaveToShufflePatternsOp : Op<Transform_Dialect,
310+
"apply_patterns.vector.interleave_to_shuffle",
311+
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
312+
let description = [{
313+
Indicates that 1D vector interleave operations should be rewritten as
314+
vector shuffle operations.
315+
316+
This is motivated by some current codegen backends not handling vector
317+
interleave operations.
318+
}];
319+
320+
let assemblyFormat = "attr-dict";
321+
}
322+
309323
def ApplyRewriteNarrowTypePatternsOp : Op<Transform_Dialect,
310324
"apply_patterns.vector.rewrite_narrow_types",
311325
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {

‎mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,9 @@ void populateVectorInterleaveLoweringPatterns(RewritePatternSet &patterns,
273273
int64_t targetRank = 1,
274274
PatternBenefit benefit = 1);
275275

276+
void populateVectorInterleaveToShufflePatterns(RewritePatternSet &patterns,
277+
PatternBenefit benefit = 1);
278+
276279
} // namespace vector
277280
} // namespace mlir
278281
#endif // MLIR_DIALECT_VECTOR_TRANSFORMS_LOWERINGPATTERNS_H

‎mlir/lib/Conversion/VectorToSPIRV/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,6 @@ add_mlir_conversion_library(MLIRVectorToSPIRV
1414
MLIRSPIRVDialect
1515
MLIRSPIRVConversion
1616
MLIRVectorDialect
17+
MLIRVectorTransforms
1718
MLIRTransforms
1819
)

‎mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
1919
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
2020
#include "mlir/Dialect/Vector/IR/VectorOps.h"
21+
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
2122
#include "mlir/IR/Attributes.h"
2223
#include "mlir/IR/BuiltinAttributes.h"
2324
#include "mlir/IR/BuiltinTypes.h"
@@ -828,6 +829,9 @@ void mlir::populateVectorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
828829
// than the generic one that extracts all elements.
829830
patterns.add<VectorReductionToFPDotProd>(typeConverter, patterns.getContext(),
830831
PatternBenefit(2));
832+
833+
// Need this until vector.interleave is handled.
834+
vector::populateVectorInterleaveToShufflePatterns(patterns);
831835
}
832836

833837
void mlir::populateVectorReductionToSPIRVDotProductPatterns(

‎mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,11 @@ void transform::ApplyLowerInterleavePatternsOp::populatePatterns(
164164
vector::populateVectorInterleaveLoweringPatterns(patterns);
165165
}
166166

167+
void transform::ApplyInterleaveToShufflePatternsOp::populatePatterns(
168+
RewritePatternSet &patterns) {
169+
vector::populateVectorInterleaveToShufflePatterns(patterns);
170+
}
171+
167172
void transform::ApplyRewriteNarrowTypePatternsOp::populatePatterns(
168173
RewritePatternSet &patterns) {
169174
populateVectorNarrowTypeRewritePatterns(patterns);

‎mlir/lib/Dialect/Vector/Transforms/LowerVectorInterleave.cpp‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
1717
#include "mlir/IR/BuiltinTypes.h"
1818
#include "mlir/IR/PatternMatch.h"
19+
#include "mlir/Support/LogicalResult.h"
1920

2021
#define DEBUG_TYPE "vector-interleave-lowering"
2122

@@ -77,9 +78,49 @@ class UnrollInterleaveOp : public OpRewritePattern<vector::InterleaveOp> {
7778
int64_t targetRank = 1;
7879
};
7980

81+
/// Rewrite vector.interleave op into an equivalent vector.shuffle op, when
82+
/// applicable: `sourceType` must be 1D and non-scalable.
83+
///
84+
/// Example:
85+
///
86+
/// ```mlir
87+
/// vector.interleave %a, %b : vector<7xi16>
88+
/// ```
89+
///
90+
/// Is rewritten into:
91+
///
92+
/// ```mlir
93+
/// vector.shuffle %arg0, %arg1 [0, 7, 1, 8, 2, 9, 3, 10, 4, 11, 5, 12, 6, 13]
94+
/// : vector<7xi16>, vector<7xi16>
95+
/// ```
96+
class InterleaveToShuffle : public OpRewritePattern<vector::InterleaveOp> {
97+
public:
98+
InterleaveToShuffle(MLIRContext *context, PatternBenefit benefit = 1)
99+
: OpRewritePattern(context, benefit) {};
100+
101+
LogicalResult matchAndRewrite(vector::InterleaveOp op,
102+
PatternRewriter &rewriter) const override {
103+
VectorType sourceType = op.getSourceVectorType();
104+
if (sourceType.getRank() != 1 || sourceType.isScalable()) {
105+
return failure();
106+
}
107+
int64_t n = sourceType.getNumElements();
108+
auto seq = llvm::seq<int64_t>(2 * n);
109+
auto zip = llvm::to_vector(llvm::map_range(
110+
seq, [n](int64_t i) { return (i % 2 ? n : 0) + i / 2; }));
111+
rewriter.replaceOpWithNewOp<ShuffleOp>(op, op.getLhs(), op.getRhs(), zip);
112+
return success();
113+
}
114+
};
115+
80116
} // namespace
81117

82118
void mlir::vector::populateVectorInterleaveLoweringPatterns(
83119
RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) {
84120
patterns.add<UnrollInterleaveOp>(targetRank, patterns.getContext(), benefit);
85121
}
122+
123+
void mlir::vector::populateVectorInterleaveToShufflePatterns(
124+
RewritePatternSet &patterns, PatternBenefit benefit) {
125+
patterns.add<InterleaveToShuffle>(patterns.getContext(), benefit);
126+
}

‎mlir/test/Dialect/Vector/vector-interleave-to-shuffle.mlir‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
// RUN: mlir-opt %s --transform-interpreter | FileCheck %s
2+
3+
// CHECK-LABEL: @vector_interleave_to_shuffle
4+
func.func @vector_interleave_to_shuffle(%a: vector<7xi16>, %b: vector<7xi16>) -> vector<14xi16>
5+
{
6+
%0 = vector.interleave %a, %b : vector<7xi16>
7+
return %0 : vector<14xi16>
8+
}
9+
// CHECK: vector.shuffle %arg0, %arg1 [0, 7, 1, 8, 2, 9, 3, 10, 4, 11, 5, 12, 6, 13] : vector<7xi16>, vector<7xi16>
10+
11+
module attributes {transform.with_named_sequence} {
12+
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
13+
%f = transform.structured.match ops{["func.func"]} in %module_op
14+
: (!transform.any_op) -> !transform.any_op
15+
16+
transform.apply_patterns to %f {
17+
transform.apply_patterns.vector.interleave_to_shuffle
18+
} : !transform.any_op
19+
transform.yield
20+
}
21+
}

‎utils/bazel/llvm-project-overlay/mlir/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5877,6 +5877,7 @@ cc_library(
58775877
":Support",
58785878
":TransformUtils",
58795879
":VectorDialect",
5880+
":VectorTransforms",
58805881
"//llvm:Support",
58815882
],
58825883
)

0 commit comments

Comments
 (0)