[Mlir-commits] [mlir] 5661ed6 - [mlir][vector] Fix crashes in MaskOp::fold and CanonializeEmptyMaskOp (#183781)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-02-27T17:50:22Z
New Revision: 5661ed60e37d63d5deff6f4c943b928b6b957049

URL: https://github.com/llvm/llvm-project/commit/5661ed60e37d63d5deff6f4c943b928b6b957049
DIFF: https://github.com/llvm/llvm-project/commit/5661ed60e37d63d5deff6f4c943b928b6b957049.diff

LOG: [mlir][vector] Fix crashes in MaskOp::fold and CanonializeEmptyMaskOp (#183781)

Two related crashes were fixed in vector.mask handling:

1. MaskOp::fold() crashes with a null pointer dereference when the mask
is all-true and the mask body has no maskable operation (only a
vector.yield). getMaskableOp() returns nullptr in this case, and the
fold was calling nullptr->dropAllUses(). Fixed by returning failure()
when there is no maskable op, deferring to the canonicalizer.

2. CanonializeEmptyMaskOp creates an invalid arith.select when the mask
type is a vector (e.g., vector<1xi1>) but the result type is a scalar
(e.g., i32). arith.select with a vector condition requires the value
types to be vectors of the same shape. Fixed by bailing out when any
result type doesn't match the mask shape.

Regression tests are added for both cases.

Fixes #177833

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index aa6734049bbd7..b935ad77c1c14 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7630,7 +7630,11 @@ LogicalResult MaskOp::fold(FoldAdaptor adaptor,
     return failure();
 
   // Move maskable operation outside of the `vector.mask` region.
+  // If there is no maskable op (empty body), the fold cannot proceed; the
+  // canonicalizer handles this case instead.
   Operation *maskableOp = getMaskableOp();
+  if (!maskableOp)
+    return failure();
   maskableOp->dropAllUses();
   maskableOp->moveBefore(getOperation());
 
@@ -7661,6 +7665,17 @@ class CanonializeEmptyMaskOp : public OpRewritePattern<MaskOp> {
     if (!maskOp.hasPassthru())
       return failure();
 
+    // arith.select with a vector condition requires the value types to be
+    // vectors of the same shape. Since vector.mask always has a vector mask
+    // type, bail out when any result type doesn't match the mask shape to
+    // avoid creating invalid IR.
+    VectorType maskType = maskOp.getMask().getType();
+    for (Type resultType : maskOp.getResultTypes()) {
+      auto vecResultType = dyn_cast<VectorType>(resultType);
+      if (!vecResultType || vecResultType.getShape() != maskType.getShape())
+        return failure();
+    }
+
     Block *block = maskOp.getMaskBlock();
     auto terminator = cast<vector::YieldOp>(block->front());
     assert(terminator.getNumOperands() == 1 &&

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 82b2cb633d1c9..3980c179b5d0a 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4082,3 +4082,36 @@ func.func @extract_strided_slice_outer_shorter(%arg0: vector<4x8x16xf32>) -> vec
   %1 = vector.extract_strided_slice %0 {offsets = [0], sizes = [2], strides = [1]} : vector<3x4x16xf32> to vector<2x4x16xf32>
   return %1 : vector<2x4x16xf32>
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/177833:
+// CanonializeEmptyMaskOp must not create arith.select with a vector condition
+// and scalar value types (incompatible with arith.select semantics). The
+// vector.mask op should remain unchanged.
+//
+// CHECK-LABEL: func @no_fold_empty_mask_scalar_result_with_passthru
+//  CHECK-SAME:     %[[MASK:.*]]: vector<1xi1>, %[[PASSTHRU:.*]]: i32, %[[VAL:.*]]: i32
+//       CHECK:   vector.mask %[[MASK]], %[[PASSTHRU]] { vector.yield %[[VAL]] : i32 } : vector<1xi1> -> i32
+//   CHECK-NOT:   arith.select
+func.func @no_fold_empty_mask_scalar_result_with_passthru(
+    %mask : vector<1xi1>, %passthru : i32, %val : i32) -> i32 {
+  %result = vector.mask %mask, %passthru { vector.yield %val : i32 } : vector<1xi1> -> i32
+  return %result : i32
+}
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/177833:
+// MaskOp::fold must not crash with a null pointer dereference when the mask
+// is all-true and the mask body has no maskable op (only a yield).
+//
+// CHECK-LABEL: func @no_fold_alltrue_mask_empty_body_scalar_result
+//  CHECK-SAME:     %[[PASSTHRU:.*]]: i32, %[[VAL:.*]]: i32
+//       CHECK:   vector.mask
+func.func @no_fold_alltrue_mask_empty_body_scalar_result(
+    %passthru : i32, %val : i32) -> i32 {
+  %all_true = vector.constant_mask [1] : vector<1xi1>
+  %result = vector.mask %all_true, %passthru { vector.yield %val : i32 } : vector<1xi1> -> i32
+  return %result : i32
+}