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.

Sign up for GitHub

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] Handle 0-rank case in fold instead of RewriterPattern #130168

Merged
merged 1 commit into from
Mar 24, 2025
Merged

[mlir][Vector] Handle 0-rank case in fold instead of RewriterPattern #130168

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
12 changes: 5 additions & 7 deletions mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1678,7 +1678,7 @@ static Value foldExtractFromBroadcast(ExtractOp extractOp) {
return source;

unsigned extractResultRank = getRank(extractOp.getType());
if (extractResultRank >= broadcastSrcRank)
if (extractResultRank > broadcastSrcRank)
return Value();
// Check that the dimension of the result haven't been broadcasted.
auto extractVecType = llvm::dyn_cast<VectorType>(extractOp.getType());
Expand Down Expand Up @@ -2159,13 +2159,11 @@ class ExtractOpFromBroadcast final : public OpRewritePattern<ExtractOp> {
// folding patterns.
if (extractResultRank < broadcastSrcRank)
return failure();
// For scalar result, the input can only be a rank-0 vector, which will
// be handled by the folder.
if (extractResultRank == 0)
return failure();

// Special case if broadcast src is a 0D vector.
if (extractResultRank == 0) {
assert(broadcastSrcRank == 0 && llvm::isa<VectorType>(source.getType()));
rewriter.replaceOpWithNewOp<vector::ExtractElementOp>(extractOp, source);
return success();
}
rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
extractOp, extractOp.getType(), source);
return success();
Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Dialect/Vector/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -736,7 +736,7 @@ func.func @fold_extract_broadcast_same_input_output_vec(%a : vector<4xf32>,

// CHECK-LABEL: fold_extract_broadcast_0dvec_input_scalar_output
// CHECK-SAME: %[[A:.*]]: vector<f32>
// CHECK: %[[B:.+]] = vector.extractelement %[[A]][] : vector<f32>
// CHECK: %[[B:.+]] = vector.extract %[[A]][] : f32 from vector<f32>
// CHECK: return %[[B]] : f32
func.func @fold_extract_broadcast_0dvec_input_scalar_output(%a : vector<f32>,
%idx0 : index, %idx1 : index, %idx2: index) -> f32 {
Expand Down Expand Up @@ -2834,7 +2834,7 @@ func.func @extract_from_0d_splat_broadcast_regression(%a: f32, %b: vector<f32>,
%3 = vector.extract %2[] : f32 from vector<f32>

// Broadcast 0D to 3D and extract scalar.
// CHECK: %[[extract1:.*]] = vector.extractelement %[[b]][] : vector<f32>
// CHECK: %[[extract1:.*]] = vector.extract %[[b]][] : f32 from vector<f32>
%4 = vector.broadcast %b : vector<f32> to vector<1x2x4xf32>
%5 = vector.extract %4[0, 0, 1] : f32 from vector<1x2x4xf32>

Expand Down
Loading