Skip to content

[mlir][Vector] Fix mask unpacking in transfer op unrolling #144889

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][Vector] Fix mask unpacking in transfer op unrolling #144889

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
15 changes: 13 additions & 2 deletions mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1208,11 +1208,22 @@ static void maybeAssignMask(OpBuilder &b, OpTy xferOp, OpTy newXferOp,
if (xferOp.getMaskType().getRank() > 1) {
// Unpack one dimension of the mask.
OpBuilder::InsertionGuard guard(b);
Location loc = xferOp.getLoc();
b.setInsertionPoint(newXferOp); // Insert load before newXfer.

auto expr = dyn_cast<AffineDimExpr>(
compressUnusedDims(xferOp.getPermutationMap()).getResult(0));
assert(expr && "cannot extract from dimension");
// Transpose dim to be the outer most dimension, so we can use
// vector.extract on it.
Copy link
Contributor

Choose a reason for hiding this comment

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

I'd rephrase a bit:

vector.extract can only extract the most minor dimensions of an multi-dimensional vector.
Transpose `d0` to the most most minor dimension so we can extract the (n-1)-D submask.

TypedValue<VectorType> mask = xferOp.getMask();
SmallVector<int64_t> perm =
llvm::to_vector(llvm::seq<int64_t>(mask.getType().getRank()));
std::swap(perm[0], perm[expr.getPosition()]);
mask = b.create<vector::TransposeOp>(loc, mask, perm);
// Extract from the transposed mask.
llvm::SmallVector<int64_t, 1> indices({i});
Location loc = xferOp.getLoc();
auto newMask = b.create<vector::ExtractOp>(loc, xferOp.getMask(), indices);
auto newMask = b.create<vector::ExtractOp>(loc, mask, indices);
newXferOp.getMaskMutable().assign(newMask);
}

Expand Down
30 changes: 30 additions & 0 deletions mlir/test/Conversion/VectorToSCF/unrolled-vector-to-loops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,3 +84,33 @@ func.func @transfer_read_mask(%A : memref<?x?x?xf32>, %mask : vector<2x3x4xi1>)
%vec = vector.transfer_read %A[%c0, %c0, %c0], %f0, %mask {in_bounds = [true, true, true]}: memref<?x?x?xf32>, vector<2x3x4xf32>
return %vec : vector<2x3x4xf32>
}

// -----

func.func @transfer_read_perm_mask(%A : memref<?x?x?x?xf32>, %mask : vector<3x2x4xi1>) -> (vector<2x3x4xf32>) {
Copy link
Contributor

Choose a reason for hiding this comment

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

IMO, this and other tests in this file are missing LIT variables that would demonstrate that e.g. %MASK_1 is used for %XFER_READ_1.

%f0 = arith.constant 0.0: f32
%c0 = arith.constant 0: index

// CHECK: vector.extract %{{.*}}[0, 0] : vector<4xi1> from vector<3x2x4xi1>
Copy link
Contributor

Choose a reason for hiding this comment

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

can we see a little more in the test (at least the transpose on the mask) ?

Copy link
Member Author

Choose a reason for hiding this comment

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

I'm not sure if that's possible if unrolling completely, because the transpose will just fold with the vector.extract on further unrolling. The vector.extract indices do show that the mask is being read in a transposed fashion. The other solution is to have a test that doesn't unroll fully. Any ideas what would be prefered?

Copy link
Contributor

Choose a reason for hiding this comment

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

ok, I missed that the test was also doing that.
There should be a max-transfer-rank or similar parameter where one could stop the unrolling at 2-D transfer reads (e.g. for HW that support > 1-D loads) but it is likely not worth the trouble at this point.

// CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?x?xf32>, vector<4xf32>
// CHECK-NEXT: vector.insert {{.*}} [0, 0] : vector<4xf32> into vector<2x3x4xf32>
// CHECK-NEXT: vector.extract %{{.*}}[1, 0] : vector<4xi1> from vector<3x2x4xi1>
// CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?x?xf32>, vector<4xf32>
// CHECK-NEXT: vector.insert {{.*}} [0, 1] : vector<4xf32> into vector<2x3x4xf32>
// CHECK-NEXT: vector.extract %{{.*}}[2, 0] : vector<4xi1> from vector<3x2x4xi1>
// CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?x?xf32>, vector<4xf32>
// CHECK-NEXT: vector.insert {{.*}} [0, 2] : vector<4xf32> into vector<2x3x4xf32>
// CHECK-NEXT: vector.extract %{{.*}}[0, 1] : vector<4xi1> from vector<3x2x4xi1>
// CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?x?xf32>, vector<4xf32>
// CHECK-NEXT: vector.insert {{.*}} [1, 0] : vector<4xf32> into vector<2x3x4xf32>
// CHECK-NEXT: vector.extract %{{.*}}[1, 1] : vector<4xi1> from vector<3x2x4xi1>
// CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?x?xf32>, vector<4xf32>
// CHECK-NEXT: vector.insert {{.*}} [1, 1] : vector<4xf32> into vector<2x3x4xf32>
// CHECK-NEXT: vector.extract %{{.*}}[2, 1] : vector<4xi1> from vector<3x2x4xi1>
// CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?x?xf32>, vector<4xf32>
// CHECK-NEXT: vector.insert {{.*}} [1, 2] : vector<4xf32> into vector<2x3x4xf32>
// CHECK-NOT: scf.if
// CHECK-NOT: scf.for
%vec = vector.transfer_read %A[%c0, %c0, %c0, %c0], %f0, %mask {permutation_map = affine_map<(d0, d1, d2, d4) -> (d2, d0, d4)>, in_bounds = [true, true, true]}: memref<?x?x?x?xf32>, vector<2x3x4xf32>
Copy link
Contributor

@nicolasvasilache nicolasvasilache Jun 19, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

hmm it is surprising to not have the mask type as part of the op given that the mapping is not trivial between vector<2x3x4xf32> and vector<3x2x4xi1>.
@dcaballe should the parser/printer be improved? (in a future PR)

Copy link
Contributor

Choose a reason for hiding this comment

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

Printing mask type would make sense to me. We discussed something similar recently:

However, there's a broader question. Do we need to support both forms:

%vec = vector.transfer_read %A[%c0, %c0, %c0, %c0], %f0, %mask

vs

%vec = vector.mask %mask { vector.transfer_read %A[%c0, %c0, %c0, %c0], %f0 }

?

Also, @Groverkss , this is a very nice example that demonstrates a case where the shape of the mask and the output vectors are different. We miss such examples in ops.mlir and I'd be tempted to add it there. Just as a nice-to-have.

return %vec : vector<2x3x4xf32>
}
Loading