[Mlir-commits] [mlir] [MLIR][Vector] Move scalable dims check before IR creation in BroadcastOpLowering (PR #188953)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188953

BroadcastOpLowering::matchAndRewrite created ub::PoisonOp before checking whether a scalable outer dimension prevents the stretch-not-at-start case. When that check triggered, the pattern returned failure() after IR was already modified, violating MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

Fix: move the scalable dimension check to before the PoisonOp creation.

Assisted-by: Claude Code
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

>From 1da5c53b750ec922a36a22e9222406f616b1c5c6 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 15:57:45 -0700
Subject: [PATCH] [MLIR][Vector] Move scalable dims check before IR creation in
 BroadcastOpLowering

BroadcastOpLowering::matchAndRewrite created ub::PoisonOp before checking
whether a scalable outer dimension prevents the stretch-not-at-start case.
When that check triggered, the pattern returned failure() after IR was
already modified, violating MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

Fix: move the scalable dimension check to before the PoisonOp creation.

Assisted-by: Claude Code
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.
---
 .../Vector/Transforms/LowerVectorBroadcast.cpp     | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
index 61d9357e19bb4..66dd7c8f36e6b 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
@@ -117,6 +117,16 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
     VectorType resType =
         VectorType::get(dstType.getShape().drop_front(), eltType,
                         dstType.getScalableDims().drop_front());
+
+    // For "stretch not at start" with a scalable outer dimension we would need
+    // to emit an scf.for loop, which is not yet supported.  Check before
+    // creating any IR so that returning failure() does not violate the pattern
+    // API contract.
+    if (m != 0 && dstType.getScalableDims()[0]) {
+      // TODO: For scalable vectors we should emit an scf.for loop.
+      return failure();
+    }
+
     Value result = ub::PoisonOp::create(rewriter, loc, dstType);
     if (m == 0) {
       // Stetch at start.
@@ -126,10 +136,6 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
         result = vector::InsertOp::create(rewriter, loc, bcst, result, d);
     } else {
       // Stetch not at start.
-      if (dstType.getScalableDims()[0]) {
-        // TODO: For scalable vectors we should emit an scf.for loop.
-        return failure();
-      }
       for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d) {
         Value ext = vector::ExtractOp::create(rewriter, loc, op.getSource(), d);
         Value bcst = vector::BroadcastOp::create(rewriter, loc, resType, ext);