[mlir] [vector] Add linearization pattern for vector.create_mask #138214

nbpatel · 2025-05-01T23:06:15Z

This PR is a breakdown [3 / 4] of the PR #136193
The PR adds linearization patterns for vector.create_mask

llvmbot · 2025-05-01T23:06:50Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-vector

Author: Nishant Patel (nbpatel)

Changes

This PR is a breakdown [3 / 4] of the PR #136193
The PR adds linearization patterns for vector.create_mask

Full diff: https://github.com/llvm/llvm-project/pull/138214.diff

3 Files Affected:

(modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+62-3)
(modified) mlir/test/Dialect/Vector/linearize.mlir (+38)
(modified) mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp (+2-1)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index b9cef003fa365..cdd937eed6569 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -445,6 +445,64 @@ struct LinearizeVectorSplat final
   }
 };
 
+/// This pattern converts the CreateMaskOp to work on a
+/// linearized vector. The pattern currently
+/// supports only 2D masks with a unit outer dimension.
+/// Following,
+///   vector.create_mask %dims : vector<1x4xi1>
+/// is converted to:
+///   %out_1d = vector.create_mask %dims : vector<4xi1>
+///   %out_nd = vector.shape_cast %out_1d : vector<4xi1> to vector<1x4xi1>
+struct LinearizeVectorCreateMask final
+    : OpConversionPattern<vector::CreateMaskOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LinearizeVectorCreateMask(const TypeConverter &typeConverter,
+                            MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit) {}
+
+  LogicalResult
+  matchAndRewrite(vector::CreateMaskOp createMaskOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcTy = createMaskOp.getType();
+    auto srcShape = srcTy.getShape();
+    if (srcShape.size() != 2)
+      return rewriter.notifyMatchFailure(createMaskOp,
+                                         "only 2D mask is supported.");
+
+    if (srcShape[0] != 1)
+      return rewriter.notifyMatchFailure(
+          createMaskOp, "only unit outer dimension is supported.");
+
+    auto dstTy = getTypeConverter()->convertType(srcTy);
+    if (!dstTy)
+      return rewriter.notifyMatchFailure(createMaskOp, "cannot convert type.");
+
+    // Compare the first operand with 0. If it's less than or equal to 0,
+    // create a zero mask, else strip the first operand and create a mask
+    // using the second operand.
+    auto firstOperand = adaptor.getOperands().front();
+    auto zero =
+        rewriter.create<mlir::arith::ConstantIndexOp>(createMaskOp.getLoc(), 0);
+    auto isZeroOrNegative = rewriter.create<mlir::arith::CmpIOp>(
+        createMaskOp.getLoc(), mlir::arith::CmpIPredicate::sle, firstOperand,
+        zero);
+    auto isZeroOrNegativeSplat = rewriter.create<mlir::vector::SplatOp>(
+        createMaskOp.getLoc(), dstTy, isZeroOrNegative);
+
+    // Use a select operation to choose between the masks.
+    auto zeroMask = rewriter.create<mlir::arith::ConstantOp>(
+        createMaskOp.getLoc(), dstTy, rewriter.getZeroAttr(dstTy));
+    auto newMask = rewriter.create<mlir::vector::CreateMaskOp>(
+        createMaskOp.getLoc(), dstTy, adaptor.getOperands().back());
+    auto result = rewriter.create<mlir::arith::SelectOp>(
+        createMaskOp.getLoc(), isZeroOrNegativeSplat, zeroMask, newMask);
+
+    rewriter.replaceOp(createMaskOp, result.getResult());
+    return success();
+  }
+};
+
 } // namespace
 
 /// Return true if the operation `op` does not support scalable vectors and
@@ -530,9 +588,10 @@ void mlir::vector::populateForVectorLinearize(TypeConverter &typeConverter,
 void mlir::vector::populateVectorLinearizeBasePatterns(
     const TypeConverter &typeConverter, const ConversionTarget &target,
     RewritePatternSet &patterns) {
-  patterns.add<LinearizeConstantLike, LinearizeVectorizable,
-               LinearizeVectorBitCast, LinearizeVectorSplat>(
-      typeConverter, patterns.getContext());
+  patterns
+      .add<LinearizeConstantLike, LinearizeVectorizable, LinearizeVectorBitCast,
+           LinearizeVectorSplat, LinearizeVectorCreateMask>(
+          typeConverter, patterns.getContext());
 }
 
 void mlir::vector::populateVectorLinearizeShuffleLikeOpsPatterns(
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 01ad1ac48b012..01872426c77bb 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -345,3 +345,41 @@ func.func @linearize_scalable_vector_splat(%arg0: i32) -> vector<4x[2]xi32> {
   %0 = vector.splat %arg0 : vector<4x[2]xi32>
   return %0 : vector<4x[2]xi32>
 }
+
+// -----
+// ALL-LABEL: linearize_create_mask
+func.func @linearize_create_mask() -> vector<1x16xi1> {
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[C20:.*]] = arith.constant 20 : index
+  // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
+  // CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[C0]], %[[C0_0]] : index
+  // CHECK: %[[SPLAT:.*]] = vector.splat %[[CMP]] : vector<16xi1>
+  // CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<16xi1>
+  // CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[C20]] : vector<16xi1>
+  // CHECK: %[[SELECT:.*]] = arith.select %[[SPLAT]], %[[CST]], %[[MASK_1D]] : vector<16xi1>, vector<16xi1>
+  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[SELECT]] : vector<16xi1> to vector<1x16xi1>
+  // CHECK: return %[[CAST]] : vector<1x16xi1>
+  %c0 = arith.constant 0 : index
+  %c20 = arith.constant 20 : index
+  %0 = vector.create_mask %c0, %c20 : vector<1x16xi1>
+  return %0 : vector<1x16xi1>
+}
+
+// -----
+// ALL-LABEL: linearize_scalable_create_mask
+func.func @linearize_scalable_create_mask() -> vector<1x[16]xi1> {
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[C20:.*]] = arith.constant 20 : index
+  // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
+  // CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[C0]], %[[C0_0]] : index
+  // CHECK: %[[SPLAT:.*]] = vector.splat %[[CMP]] : vector<[16]xi1>
+  // CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<[16]xi1>
+  // CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[C20]] : vector<[16]xi1>
+  // CHECK: %[[SELECT:.*]] = arith.select %[[SPLAT]], %[[CST]], %[[MASK_1D]] : vector<[16]xi1>, vector<[16]xi1>
+  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[SELECT]] : vector<[16]xi1> to vector<1x[16]xi1>
+  // CHECK: return %[[CAST]] : vector<1x[16]xi1>
+  %c0 = arith.constant 0 : index
+  %c20 = arith.constant 20 : index
+  %0 = vector.create_mask %c0, %c20 : vector<1x[16]xi1>
+  return %0 : vector<1x[16]xi1>
+}
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index eda2594fbc7c7..2d5e90908d4d0 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -973,7 +973,8 @@ struct TestVectorLinearize final
     return "Linearizes ND vectors for N >= 2 into 1D vectors";
   }
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<vector::VectorDialect>();
+    registry.insert<vector::VectorDialect, memref::MemRefDialect,
+                    arith::ArithDialect>();
   }
 
   void runOnOperation() override {

nbpatel · 2025-05-01T23:09:04Z

@Hardcode84

nbpatel · 2025-05-06T18:08:49Z

Ping for review

dcaballe

LGTM, thanks! Mostly some folding needed.

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

dcaballe · 2025-05-06T23:28:03Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+        createMaskOp.getLoc(), dstTy, rewriter.getZeroAttr(dstTy));
+    auto newMask = rewriter.create<mlir::vector::CreateMaskOp>(
+        createMaskOp.getLoc(), dstTy, adaptor.getOperands().back());
+    auto result = rewriter.create<mlir::arith::SelectOp>(


We may want to use some createAndFold here to get rid of redundant IR...

Thanks for updating the test @nbpatel to use non-constant operands.

nit: my guess is that most of the time, the unit dimension will have extent which is the constant 1, and so createOrFold will still be advisable.

dcaballe · 2025-05-06T23:29:00Z

mlir/test/Dialect/Vector/linearize.mlir

+  // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
+  // CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[C0]], %[[C0_0]] : index
+  // CHECK: %[[SPLAT:.*]] = vector.splat %[[CMP]] : vector<[16]xi1>
+  // CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<[16]xi1>


most of these checks should be folded with createAndFold

Alternatively, use an %arg0: index instead of %c0

newling

Thanks for the clear PR!

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

newling · 2025-05-07T17:04:12Z

mlir/test/Dialect/Vector/linearize.mlir

+  // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
+  // CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[C0]], %[[C0_0]] : index
+  // CHECK: %[[SPLAT:.*]] = vector.splat %[[CMP]] : vector<[16]xi1>
+  // CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<[16]xi1>


Alternatively, use an %arg0: index instead of %c0

mlir/test/Dialect/Vector/linearize.mlir

dcaballe

Thanks for addressing the comments. A couple more:

dcaballe · 2025-05-13T01:36:05Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+///   %index = arith.index_cast %cmpi : i1 to index
+///   %mul = arith.muli %index, %arg1 : index


this mul looks like and bitwise 'and' operation to me?

oh yes, good catch. Thanks

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

nbpatel · 2025-05-13T02:53:22Z

@newling @dcaballe addressed the comments. Thanks for the feedback.

github-actions · 2025-05-13T02:53:43Z

✅ With the latest revision this PR passed the C/C++ code formatter.

newling

Thanks, LGTM. I've added 2 minor suggestions but I'm happy for this to land as is.

It looks like all review comments have been addressed, but I'll give other reviewers some more time to comment, otherwise I'll land this at the end of week if that's ok.

newling · 2025-05-14T16:57:18Z

mlir/test/Dialect/Vector/linearize.mlir

+
+// -----
+// CHECK-LABEL: linearize_scalable_create_mask
+// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) -> vector<1x[16]xi1>


nit: I think it is better to not capture arguments (ARG0 and and ARG1) if they are not subsequently used. Thanks for reducing the size of this test! In fact I think it could reduced even further to

// This test is the same as linearize_create_mask but with a scalable mask. Confirms that it linearizes.
// CHECK-LABEL: linearize_scalable_create_mask
// CHECK: vector.create_mask {{.*}} vector<[16]xi1>

newling · 2025-05-14T17:01:13Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    Location loc = createMaskOp.getLoc();
+    VectorType srcTy = createMaskOp.getType();
+    auto srcShape = srcTy.getShape();
+    if (srcShape.size() != 2)


Consider adding
// FIXME: add support for any vector with at most 1 non-unit dimension (like vector<1x4x1xi1>)
here

dcaballe

Thanks!

nbpatel · 2025-05-14T19:07:51Z

@dcaballe or @newling can one of you help me land this?

newling · 2025-05-14T22:49:03Z

@dcaballe or @newling can one of you help me land this?

Will do 👍

…m#138214) This PR is a breakdown [3 / 4] of the PR llvm#136193 The PR adds linearization patterns for vector.create_mask

nbpatel added 4 commits May 1, 2025 22:10

Add linearization pattern for vector.create_mask

3a83e2d

Merge branch 'main' into vector_linearize_create_mask

203ec82

Use CHECKS

24b0739

Add test case for scalable vector

2b8a653

nbpatel requested review from banach-space, dcaballe, hanhanW and nicolasvasilache as code owners May 1, 2025 23:06

llvmbot added mlir:vectorops mlir mlir:vector labels May 1, 2025

Clean up

8e8de7a

newling mentioned this pull request May 6, 2025

[mlir][vector] Linearize vector.create_mask (flatten to 1D) #138760

Closed

dcaballe reviewed May 6, 2025

View reviewed changes

newling reviewed May 7, 2025

View reviewed changes

nbpatel added 3 commits May 8, 2025 23:39

Address Feedback

528f913

Replace select with mul

c2c1a22

Fix typo

c5b2e81

dcaballe reviewed May 13, 2025

View reviewed changes

Address comments

8fca9c1

nbpatel added 2 commits May 13, 2025 02:54

Fix doc

2b7e06e

Clang-format

db484c9

newling approved these changes May 14, 2025

View reviewed changes

dcaballe approved these changes May 14, 2025

View reviewed changes

Merge branch 'main' into vector_linearize_create_mask

7b714ce

newling merged commit 1778d3b into llvm:main May 14, 2025
11 checks passed

		/// %index = arith.index_cast %cmpi : i1 to index
		/// %mul = arith.muli %index, %arg1 : index

[mlir] [vector] Add linearization pattern for vector.create_mask #138214

[mlir] [vector] Add linearization pattern for vector.create_mask #138214

Uh oh!

Conversation

nbpatel commented May 1, 2025

Uh oh!

llvmbot commented May 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nbpatel commented May 1, 2025

Uh oh!

nbpatel commented May 6, 2025

Uh oh!

dcaballe left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

newling May 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

newling left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dcaballe left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nbpatel commented May 13, 2025

Uh oh!

github-actions bot commented May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

newling left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcaballe left a comment

Choose a reason for hiding this comment

Uh oh!

nbpatel commented May 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

newling commented May 14, 2025

Uh oh!

Uh oh!

Uh oh!

llvmbot commented May 1, 2025 •

edited

Loading

newling May 12, 2025 •

edited

Loading

github-actions bot commented May 13, 2025 •

edited

Loading

nbpatel commented May 14, 2025 •

edited

Loading