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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From 4501974c06e964093094fe87f4a86ae1d3f92adc Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Feb 2026 09:24:35 -0800
Subject: [PATCH] [mlir][vector] Fix crashes in MaskOp::fold and
 CanonializeEmptyMaskOp

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
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   | 15 ++++++++++
 mlir/test/Dialect/Vector/canonicalize.mlir | 33 ++++++++++++++++++++++
 2 files changed, 48 insertions(+)

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
+}