[Mlir-commits] [mlir] [mlir][Vector] Improve `vector.mask` verifier (PR #139823)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 13 18:32:31 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Diego Caballero (dcaballe)
<details>
<summary>Changes</summary>
This PR improves the `vector.mask` verifier to make sure it's not applying masking semantics to operations defined outside of the `vector.mask` region. Documentation is updated to emphasize that and make it clearer, even though it already stated that.
As part of this change, the logic that ensures that a terminator is present in the region mask has been simplified to make it less surprising to the user when a `vector.yield` is explicitly provided in the IR.
---
Full diff: https://github.com/llvm/llvm-project/pull/139823.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+7-2)
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+24-17)
- (modified) mlir/test/Dialect/Vector/invalid.mlir (+28)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 3aefcea8de994..2820759687293 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2482,8 +2482,13 @@ def Vector_MaskOp : Vector_Op<"mask", [
masked. Values used within the region are captured from above. Only one
*maskable* operation can be masked with a `vector.mask` operation at a time.
An operation is *maskable* if it implements the `MaskableOpInterface`. The
- terminator yields all results of the maskable operation to the result of
- this operation.
+ terminator yields all results from the maskable operation to the result of
+ this operation. No other values are allowed to be yielded.
+
+ An empty `vector.mask` operation is considered ill-formed but legal to
+ facilitate optimizations across the `vector.mask` operation. It is considered
+ a no-op regardless of its returned values and will be removed by the
+ canonicalizer.
The vector mask argument holds a bit for each vector lane and determines
which vector lanes should execute the maskable operation and which ones
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f6c3c6a61afb6..395680b5e814b 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6543,29 +6543,31 @@ void mlir::vector::MaskOp::print(OpAsmPrinter &p) {
}
void MaskOp::ensureTerminator(Region ®ion, Builder &builder, Location loc) {
- OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
- MaskOp>::ensureTerminator(region, builder, loc);
- // Keep the default yield terminator if the number of masked operations is not
- // the expected. This case will trigger a verification failure.
- Block &block = region.front();
- if (block.getOperations().size() != 2)
+ // Create default terminator if there are no ops to mask.
+ if (region.empty() || region.front().empty()) {
+ OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
+ MaskOp>::ensureTerminator(region, builder, loc);
return;
+ }
- // Replace default yield terminator with a new one that returns the results
- // from the masked operation.
- OpBuilder opBuilder(builder.getContext());
- Operation *maskedOp = &block.front();
- Operation *oldYieldOp = &block.back();
- assert(isa<vector::YieldOp>(oldYieldOp) && "Expected vector::YieldOp");
+ // If region has an explicit terminator, we don't modify it.
+ Block &block = region.front();
+ if (isa<vector::YieldOp>(block.back()))
+ return;
- // Empty vector.mask op.
- if (maskedOp == oldYieldOp)
+ // Create default terminator if the number of masked operations is not
+ // one. This case will trigger a verification failure.
+ if (block.getOperations().size() != 1) {
+ OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
+ MaskOp>::ensureTerminator(region, builder, loc);
return;
+ }
- opBuilder.setInsertionPoint(oldYieldOp);
+ // Create a terminator that yields the results from the masked operation.
+ OpBuilder opBuilder(builder.getContext());
+ Operation *maskedOp = &block.front();
+ opBuilder.setInsertionPointToEnd(&block);
opBuilder.create<vector::YieldOp>(loc, maskedOp->getResults());
- oldYieldOp->dropAllReferences();
- oldYieldOp->erase();
}
LogicalResult MaskOp::verify() {
@@ -6600,6 +6602,11 @@ LogicalResult MaskOp::verify() {
return emitOpError("expects number of results to match maskable operation "
"number of results");
+ if (!llvm::equal(maskableOp->getResults(), terminator.getOperands()))
+ return emitOpError(
+ "expects all the results from the MaskableOpInterface to "
+ "be returned by the terminator");
+
if (!llvm::equal(maskableOp->getResultTypes(), getResultTypes()))
return emitOpError(
"expects result type to match maskable operation result type");
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index be65d4c2eef58..de6333ca8dc7b 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1747,6 +1747,34 @@ func.func @vector_mask_0d_mask(%arg0: tensor<2x4xi32>,
// -----
+func.func @vector_mask_non_empty_external_return(%t0: tensor<?xf32>, %idx: index,
+ %m0: vector<16xi1>, %ext: vector<16xf32>) -> vector<16xf32> {
+ %ft0 = arith.constant 0.0 : f32
+ // expected-error at +1 {{'vector.mask' op expects all the results from the MaskableOpInterface to be returned by the terminator}}
+ %0 = vector.mask %m0 {
+ %1 =vector.transfer_read %t0[%idx], %ft0 : tensor<?xf32>, vector<16xf32>
+ vector.yield %ext : vector<16xf32>
+ } : vector<16xi1> -> vector<16xf32>
+
+ return %0 : vector<16xf32>
+}
+
+// -----
+
+func.func @vector_mask_non_empty_mixed_return(%t0: tensor<?xf32>, %idx: index,
+ %m0: vector<16xi1>, %ext: vector<16xf32>) -> (vector<16xf32>, vector<16xf32>) {
+ %ft0 = arith.constant 0.0 : f32
+ // expected-error at +1 {{'vector.mask' op expects number of results to match maskable operation number of results}}
+ %0:2 = vector.mask %m0 {
+ %1 =vector.transfer_read %t0[%idx], %ft0 : tensor<?xf32>, vector<16xf32>
+ vector.yield %1, %ext : vector<16xf32>, vector<16xf32>
+ } : vector<16xi1> -> (vector<16xf32>, vector<16xf32>)
+
+ return %0#0, %0#1 : vector<16xf32>, vector<16xf32>
+}
+
+// -----
+
func.func @vector_scalable_insert_unaligned(%subv: vector<4xi32>, %vec: vector<[16]xi32>) {
// expected-error at +1 {{op failed to verify that position is a multiple of the source length.}}
%0 = vector.scalable.insert %subv, %vec[2] : vector<4xi32> into vector<[16]xi32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/139823
More information about the Mlir-commits
mailing list