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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Cheng-Chieh Liu (jay87406)

<details>
<summary>Changes</summary>

`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.

---
Full diff: https://github.com/llvm/llvm-project/pull/226024.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp (+5-3) 
- (modified) mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir (+10) 


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

``````````

</details>


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