[Mlir-commits] [mlir] [mlir][vector] Extend `CreateMaskFolder` (PR #75842)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Dec 18 10:46:12 PST 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector
Author: Andrzej Warzyński (banach-space)
<details>
<summary>Changes</summary>
Extends `CreateMaskFolder` pattern so that the following:
```mlir
%c8 = arith.constant 8 : index
%c16 = arith.constant 16 : index
%0 = vector.vscale
%1 = arith.muli %0, %c16 : index
%10 = vector.create_mask %c8, %1 : vector<8x[16]xi1>
```
is folded as:
```mlir
%0 = vector.constant_mask [8, 16] : vector<8x[16]xi1>
```
---
Full diff: https://github.com/llvm/llvm-project/pull/75842.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+76-20)
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+13)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 540959b486db9c..3619c1c00f1664 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5657,30 +5657,79 @@ LogicalResult CreateMaskOp::verify() {
namespace {
-// Pattern to rewrite a CreateMaskOp with a ConstantMaskOp.
+/// Pattern to rewrite a CreateMaskOp with a ConstantMaskOp.
+///
+/// Ex 1:
+/// %c2 = arith.constant 2 : index
+/// %c3 = arith.constant 3 : index
+/// %0 = vector.create_mask %c3, %c2 : vector<4x3xi1>
+/// Becomes:
+/// vector.constant_mask [3, 2] : vector<4x3xi1>
+///
+/// Ex 2:
+/// %c_neg_1 = arith.constant -1 : index
+/// %0 = vector.create_mask %c_neg_1 : vector<[8]xi1>
+/// becomes:
+/// vector.constant_mask [0] : vector<[8]xi1>
+///
+/// Ex 3:
+/// %c8 = arith.constant 8 : index
+/// %c16 = arith.constant 16 : index
+/// %0 = vector.vscale
+/// %1 = arith.muli %0, %c16 : index
+/// %10 = vector.create_mask %c8, %1 : vector<8x[16]xi1>
+/// becomes:
+/// %0 = vector.constant_mask [8, 16] : vector<8x[16]xi1>
class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
public:
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(CreateMaskOp createMaskOp,
PatternRewriter &rewriter) const override {
- // Return if any of 'createMaskOp' operands are not defined by a constant.
- auto isNotDefByConstant = [](Value operand) {
- return !getConstantIntValue(operand).has_value();
- };
- if (llvm::any_of(createMaskOp.getOperands(), isNotDefByConstant))
- return failure();
+ VectorType retTy = createMaskOp.getResult().getType();
+ bool isScalable = retTy.isScalable();
+
+ // Check every mask operand
+ for (auto [opIdx, operand] : llvm::enumerate(createMaskOp.getOperands())) {
+ // Most basic case - this operand is a constant value. Note that for
+ // scalable dimensions, CreateMaskOp can be folded only if the
+ // corresponding operand is negative or zero.
+ if (auto op = getConstantIntValue(operand)) {
+ APInt intVal;
+ if (isScalable && !(matchPattern(operand, m_ConstantInt(&intVal)) ||
+ intVal.isStrictlyPositive()))
+ return failure();
- // CreateMaskOp for scalable vectors can be folded only if all dimensions
- // are negative or zero.
- if (auto vType = llvm::dyn_cast<VectorType>(createMaskOp.getType())) {
- if (vType.isScalable())
- for (auto opDim : createMaskOp.getOperands()) {
- APInt intVal;
- if (matchPattern(opDim, m_ConstantInt(&intVal)) &&
- intVal.isStrictlyPositive())
- return failure();
- }
+ continue;
+ }
+
+ // Non-constant operands are not allowed for non-scalable vectors.
+ if (!isScalable)
+ return failure();
+
+ // For scalable vectors, "arith.muli %vscale, %dimSize" means an "all
+ // true" mask, so can also be treated as constant.
+ auto mul = llvm::dyn_cast_or_null<arith::MulIOp>(operand.getDefiningOp());
+ if (!mul)
+ return failure();
+ auto mulLHS = mul.getOperands()[0];
+ auto mulRHS = mul.getOperands()[1];
+ bool isOneOpVscale =
+ (isa<vector::VectorScaleOp>(mulLHS.getDefiningOp()) ||
+ isa<vector::VectorScaleOp>(mulRHS.getDefiningOp()));
+
+ auto isConstantValMatchingDim =
+ [=, dim = createMaskOp.getResult().getType().getShape()[opIdx]](
+ Value operand) {
+ auto constantVal = getConstantIntValue(operand);
+ return (constantVal.has_value() && constantVal.value() == dim);
+ };
+
+ bool isOneOpConstantMatchingDim =
+ isConstantValMatchingDim(mulLHS) || isConstantValMatchingDim(mulRHS);
+
+ if (!isOneOpVscale || !isOneOpConstantMatchingDim)
+ return failure();
}
// Gather constant mask dimension sizes.
@@ -5688,15 +5737,22 @@ class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
maskDimSizes.reserve(createMaskOp->getNumOperands());
for (auto [operand, maxDimSize] : llvm::zip_equal(
createMaskOp.getOperands(), createMaskOp.getType().getShape())) {
- int64_t dimSize = getConstantIntValue(operand).value();
- dimSize = std::min(dimSize, maxDimSize);
+ auto dimSize = getConstantIntValue(operand);
+ if (not dimSize) {
+ // Although not a constant, it is safe to assume that `operand` is
+ // "vscale * maxDimSize".
+ maskDimSizes.push_back(maxDimSize);
+ continue;
+ }
+ int64_t dimSizeVal = std::min(dimSize.value(), maxDimSize);
// If one of dim sizes is zero, set all dims to zero.
if (dimSize <= 0) {
maskDimSizes.assign(createMaskOp.getType().getRank(), 0);
break;
}
- maskDimSizes.push_back(dimSize);
+ maskDimSizes.push_back(dimSizeVal);
}
+
// Replace 'createMaskOp' with ConstantMaskOp.
rewriter.replaceOpWithNewOp<ConstantMaskOp>(
createMaskOp, createMaskOp.getResult().getType(),
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 1021c73cc57d34..a30016ea857d97 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -53,6 +53,19 @@ func.func @create_vector_mask_to_constant_mask_truncation_zero() -> (vector<4x3x
// -----
+// CHECK-LABEL: create_vector_mask_to_constant_mask_scalable_all_true
+func.func @create_vector_mask_to_constant_mask_scalable_all_true() -> (vector<8x[16]xi1>) {
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %0 = vector.vscale
+ %1 = arith.muli %0, %c16 : index
+ // CHECK: vector.constant_mask [8, 16] : vector<8x[16]xi1>
+ %10 = vector.create_mask %c8, %1 : vector<8x[16]xi1>
+ return %10 : vector<8x[16]xi1>
+}
+
+// -----
+
// CHECK-LABEL: create_mask_transpose_to_transposed_create_mask
// CHECK-SAME: %[[DIM0:.*]]: index, %[[DIM1:.*]]: index, %[[DIM2:.*]]: index
func.func @create_mask_transpose_to_transposed_create_mask(
``````````
</details>
https://github.com/llvm/llvm-project/pull/75842
More information about the Mlir-commits
mailing list