Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir][vector] Add Vector-dialect interleave-to-shuffle pattern, enable in VectorToSPIRV #91800

Merged
merged 3 commits into from
May 13, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,20 @@ def ApplyLowerInterleavePatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}

def ApplyInterleaveToShufflePatternsOp : Op<Transform_Dialect,
"apply_patterns.vector.interleave_to_shuffle",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
let description = [{
Indicates that 1D vector interleave operations should be rewritten as
vector shuffle operations.

This is motivated by some current codegen backends not handling vector
interleave operations.
}];

let assemblyFormat = "attr-dict";
}

def ApplyRewriteNarrowTypePatternsOp : Op<Transform_Dialect,
"apply_patterns.vector.rewrite_narrow_types",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,9 @@ void populateVectorInterleaveLoweringPatterns(RewritePatternSet &patterns,
int64_t targetRank = 1,
PatternBenefit benefit = 1);

void populateVectorInterleaveToShufflePatterns(RewritePatternSet &patterns,
PatternBenefit benefit = 1);

} // namespace vector
} // namespace mlir
#endif // MLIR_DIALECT_VECTOR_TRANSFORMS_LOWERINGPATTERNS_H
4 changes: 4 additions & 0 deletions mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
Expand Down Expand Up @@ -828,6 +829,9 @@ void mlir::populateVectorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
// than the generic one that extracts all elements.
patterns.add<VectorReductionToFPDotProd>(typeConverter, patterns.getContext(),
PatternBenefit(2));

// Need this until vector.interleave is handled.
vector::populateVectorInterleaveToShufflePatterns(patterns);
}

void mlir::populateVectorReductionToSPIRVDotProductPatterns(
Expand Down
5 changes: 5 additions & 0 deletions mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,11 @@ void transform::ApplyLowerInterleavePatternsOp::populatePatterns(
vector::populateVectorInterleaveLoweringPatterns(patterns);
}

void transform::ApplyInterleaveToShufflePatternsOp::populatePatterns(
RewritePatternSet &patterns) {
vector::populateVectorInterleaveToShufflePatterns(patterns);
}

void transform::ApplyRewriteNarrowTypePatternsOp::populatePatterns(
RewritePatternSet &patterns) {
populateVectorNarrowTypeRewritePatterns(patterns);
Expand Down
41 changes: 41 additions & 0 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorInterleave.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LogicalResult.h"

#define DEBUG_TYPE "vector-interleave-lowering"

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

/// Rewrite vector.interleave op into an equivalent vector.shuffle op, when
/// applicable: `sourceType` must be 1D and non-scalable.
///
/// Example:
///
/// ```mlir
/// vector.interleave %a, %b : vector<7xi16>
/// ```
///
/// Is rewritten into:
///
/// ```mlir
/// vector.shuffle %arg0, %arg1 [0, 7, 1, 8, 2, 9, 3, 10, 4, 11, 5, 12, 6, 13]
/// : vector<7xi16>, vector<7xi16>
/// ```
class InterleaveToShuffle : public OpRewritePattern<vector::InterleaveOp> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note, vector.shuffle doesn't support scalable vectors:

You can just disable this pattern when sourceType.isScalable() returns true. It would be nice to have a test to verity that as well.

Also, would you mind adding a comment with IR BEFORE and AFTER (similar to what's been added to UnrollInterleaveOp)? I know that right now this is a small file and the pattern is self-explanatory, but these files tend to grow a lot at which point comments can be really helpful.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note, vector.shuffle doesn't support scalable vectors:

* [[vector][mlir] Restrict vector.shuffle to fixed-width vectors #88733](https://github.com/llvm/llvm-project/pull/88733)

You can just disable this pattern when sourceType.isScalable() returns true. It would be nice to have a test to verity that as well.

Thanks, done!

Also, would you mind adding a comment with IR BEFORE and AFTER (similar to what's been added to UnrollInterleaveOp)? I know that right now this is a small file and the pattern is self-explanatory, but these files tend to grow a lot at which point comments can be really helpful.

Good idea. Done!

public:
InterleaveToShuffle(MLIRContext *context, PatternBenefit benefit = 1)
: OpRewritePattern(context, benefit){};

LogicalResult matchAndRewrite(vector::InterleaveOp op,
PatternRewriter &rewriter) const override {
VectorType sourceType = op.getSourceVectorType();
if (sourceType.getRank() != 1 || sourceType.isScalable()) {
return failure();
}
int64_t n = sourceType.getNumElements();
auto seq = llvm::seq<int64_t>(2 * n);
auto zip = llvm::to_vector(llvm::map_range(
seq, [n](int64_t i) { return (i % 2 ? n : 0) + i / 2; }));
rewriter.replaceOpWithNewOp<ShuffleOp>(op, op.getLhs(), op.getRhs(), zip);
return success();
}
};

} // namespace

void mlir::vector::populateVectorInterleaveLoweringPatterns(
RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) {
patterns.add<UnrollInterleaveOp>(targetRank, patterns.getContext(), benefit);
}

void mlir::vector::populateVectorInterleaveToShufflePatterns(
RewritePatternSet &patterns, PatternBenefit benefit) {
patterns.add<InterleaveToShuffle>(patterns.getContext(), benefit);
}
21 changes: 21 additions & 0 deletions mlir/test/Dialect/Vector/vector-interleave-to-shuffle.mlir
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
// RUN: mlir-opt %s --transform-interpreter | FileCheck %s

// CHECK-LABEL: @vector_interleave_to_shuffle
func.func @vector_interleave_to_shuffle(%a: vector<7xi16>, %b: vector<7xi16>) -> vector<14xi16>
{
%0 = vector.interleave %a, %b : vector<7xi16>
return %0 : vector<14xi16>
}
// CHECK: vector.shuffle %arg0, %arg1 [0, 7, 1, 8, 2, 9, 3, 10, 4, 11, 5, 12, 6, 13] : vector<7xi16>, vector<7xi16>

module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%f = transform.structured.match ops{["func.func"]} in %module_op
: (!transform.any_op) -> !transform.any_op

transform.apply_patterns to %f {
transform.apply_patterns.vector.interleave_to_shuffle
} : !transform.any_op
transform.yield
}
}
Loading