[Mlir-commits] [mlir] [mlir][Vector] Guard scalable reduction dim in matmul lowering (PR #226024)

[Cheng-Chieh Liu]

https://github.com/jay87406 created https://github.com/llvm/llvm-project/pull/226024

`ContractionOpToMatmulOpLowering` previously rejected scalable `vector.contract` ops by checking only the result type. While this covers the M and N dimensions, it misses the K (reduction) dimension. For contractions with static M/N but a scalable K (e.g., `vector<2x[4]xf32>` times `vector<[4]x3xf32>`), the guard incorrectly passed. The pattern then attempted to build a flattened, non-scalable LHS/RHS type via `VectorType::get(lhsType.getNumElements(), ...)`, followed by a `vector.shape_cast`. This caused the `shape_cast` verifier to reject the conversion (source has a scalable dimension, result doesn't), crashing the pass instead of gracefully failing the match.

Also check the LHS/RHS operand types for a scalable dimension, and add a regression test.

>From 2ae840806d5cd7f004462478cbf239e83a62b380 Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Wed, 23 Sep 2026 22:45:53 -0500
Subject: [PATCH] [mlir][Vector] Guard scalable reduction dim in matmul
 lowering

ContractionOpToMatmulOpLowering rejects a scalable vector.contract by
checking only the result type, which reflects the M and N dims but
never the K (reduction) dim. For a contraction with a static M and N
but a scalable K, e.g. vector<2x[4]xf32> times vector<[4]x3xf32>, the
guard passes and the pattern proceeds to build a flattened, non-
scalable LHS/RHS type via VectorType::get(lhsType.getNumElements(),
...), then vector.shape_cast into it -- which the shape_cast verifier
rejects (source has a scalable dim, result doesn't), crashing the
pass instead of gracefully declining the match.

Also check the LHS/RHS operand types for a scalable dim, and add a
regression test.
---
 .../Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp    |  8 +++++---
 ...ector-contract-to-matrix-intrinsics-transforms.mlir | 10 ++++++++++
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index dd62377c2f6d3..5065104c9ba7a 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -2094,10 +2094,12 @@ FailureOr<Value> ContractionOpToMatmulOpLowering::matchAndRewriteMaskableOp(
 
   Type opResType = op.getType();
   VectorType vecType = dyn_cast<VectorType>(opResType);
-  if (vecType && vecType.isScalable()) {
-    // Note - this is sufficient to reject all cases with scalable vectors.
+  if (vecType && vecType.isScalable())
+    return failure();
+  // The result type only reflects the M and N dims, so a scalable K
+  // (reduction) dim would slip through the check above.
+  if (op.getLhsType().isScalable() || op.getRhsType().isScalable())
     return failure();
-  }
 
   Type elementType = op.getLhsType().getElementType();
   if (!elementType.isIntOrFloat())
diff --git a/mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
index bf9b2ace35622..995f29b7f8eef 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
@@ -75,3 +75,13 @@ func.func @matmul_scalable(%arg0: vector<2x4xf32>,
     : vector<2x4xf32>, vector<4x[3]xf32> into vector<2x[3]xf32>
   return %0 : vector<2x[3]xf32>
 }
+
+// CHECK-LABEL: func @matmul_scalable_reduction
+// CHECK-NOT: llvm.intr.matrix.multiply
+func.func @matmul_scalable_reduction(%arg0: vector<2x[4]xf32>,
+                                     %arg1: vector<[4]x3xf32>,
+                                     %arg2: vector<2x3xf32>) -> vector<2x3xf32> {
+  %0 = vector.contract #matmat_trait %arg0, %arg1, %arg2
+    : vector<2x[4]xf32>, vector<[4]x3xf32> into vector<2x3xf32>
+  return %0 : vector<2x3xf32>
+}