[Mlir-commits] [mlir] [mlir][vector] Fix crash on untraceable masks in getCompressedMaskOp (PR #207299)

[llvmlistbot at llvm.org]

https://github.com/marquisburg updated https://github.com/llvm/llvm-project/pull/207299

>From bb179a3054918920bf0f77f7d8b141d176e630de Mon Sep 17 00:00:00 2001
From: GhostDragonWasTook <marquis at marquiseco.co.za>
Date: Fri, 3 Jul 2026 02:40:25 +0200
Subject: [PATCH] [mlir][vector] Fix crash on untraceable masks in
 getCompressedMaskOp

The loop tracing a mask back to its creation op assumed the chain
always ends at a create_mask/constant_mask/arith.constant, a
block-argument mask made it fall through to isa<> on a null op
(segfault), and any other unhandled defining op spun it forever.

Fixes #206928
---
 .../Transforms/VectorEmulateNarrowType.cpp    |  8 +++++-
 .../emulate-narrow-type-unsupported.mlir      | 28 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 9faaebdcf8f35..e37296e50fb29 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -98,10 +98,16 @@ static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
     if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
       maskOp = extractOp.getSource().getDefiningOp();
       extractOps.push_back(extractOp);
+    } else {
+      // Unsupported mask-defining op (e.g. a block argument, which has no
+      // defining op, or an op we cannot trace through). Bail out rather than
+      // looping forever or dereferencing a null op below.
+      break;
     }
   }
 
-  if (!isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
+  if (!maskOp ||
+      !isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
           maskOp))
     return failure();
 
diff --git a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
index a5a6fc4acfe10..74f72ed5a5d9c 100644
--- a/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
+++ b/mlir/test/Dialect/Vector/emulate-narrow-type-unsupported.mlir
@@ -109,3 +109,31 @@ func.func @vector_maskedstore_2d_i8_negative(%arg0: index, %arg1: index, %arg2:
 //  CHECK-LABEL: func @vector_maskedstore_2d_i8_negative
 //        CHECK: memref.alloc() : memref<3x8xi8>
 //    CHECK-NOT: i32
+
+// -----
+
+// The mask is a block argument, so its defining op cannot be traced back to a
+// supported mask-creation op. This must gracefully bail out instead of crashing
+// (see llvm/llvm-project#206928).
+
+func.func @vector_maskedstore_i8_block_arg_mask_negative(%arg0: memref<?xi8>, %mask: vector<8xi1>, %value: vector<8xi8>) {
+  %c0 = arith.constant 0 : index
+  vector.maskedstore %arg0[%c0], %mask, %value : memref<?xi8>, vector<8xi1>, vector<8xi8>
+  return
+}
+
+//  CHECK-LABEL: func @vector_maskedstore_i8_block_arg_mask_negative
+//        CHECK: vector.maskedstore
+
+// -----
+
+// As above, but for `vector.maskedload`.
+
+func.func @vector_maskedload_i8_block_arg_mask_negative(%arg0: memref<?xi8>, %mask: vector<8xi1>, %passthru: vector<8xi8>) -> vector<8xi8> {
+  %c0 = arith.constant 0 : index
+  %0 = vector.maskedload %arg0[%c0], %mask, %passthru : memref<?xi8>, vector<8xi1>, vector<8xi8> into vector<8xi8>
+  return %0 : vector<8xi8>
+}
+
+//  CHECK-LABEL: func @vector_maskedload_i8_block_arg_mask_negative
+//        CHECK: vector.maskedload