[Mlir-commits] [mlir] [mlir][Vector] Reject scalable dst dim in contract-to-dot lowering (PR #224793)

[Cheng-Chieh Liu]

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

`ContractionOpToDotLowering` unrolls the free (parallel) dimensions of a `vector.contract` into per-row/column `vector.extract`/`vector.reduction` ops, using `dstType.getShape()[0]/[1]` as the loop trip count. However, for scalable free dimensions, this only uses the static minimum rather than the actual `vscale`-scaled extent. As a result, lanes beyond the fixed trip count were silently dropped when `vscale > 1`, producing incorrect data without any compilation error.

Reject scalable free dimensions before generating IR, and add negative lit test coverage for both `matvec` and `matmat` shapes.

>From d738d1b2be16f513816f32e5638a9bb2ce667d23 Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Fri, 18 Sep 2026 21:30:20 -0500
Subject: [PATCH] [mlir][Vector] Reject scalable dst dim in contract-to-dot
 lowering

ContractionOpToDotLowering unrolls the free (parallel) dimensions of a
vector.contract into per-row/column vector.extract/vector.reduction
ops, using dstType.getShape()[0]/[1] as the loop trip count. For a
scalable free dimension, that's only the static minimum, not the
actual vscale-scaled extent, so lanes beyond the fixed trip count were
silently dropped from the result whenever vscale > 1 -- no error, just
wrong data.

Reject a scalable free dimension before creating any IR, and add
negative lit coverage for both the matvec and matmat shapes.
---
 .../Vector/Transforms/LowerVectorContract.cpp | 12 +++++++++++
 .../vector-contract-to-dot-transforms.mlir    | 20 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 2a152ef0634d7..830544c552f8c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -644,6 +644,18 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
   if (vectorContractLowering != vector::VectorContractLowering::Dot)
     return failure();
 
+  // A scalable dst dim can't be unrolled with a static trip count; check
+  // before any IR is created below.
+  if (auto resultVecType = dyn_cast<VectorType>(op.getResultType())) {
+    if (resultVecType.getRank() >= 1) {
+      ArrayRef<bool> scalableDims = resultVecType.getScalableDims();
+      if (scalableDims[0] ||
+          (resultVecType.getRank() == 2 && scalableDims[1]))
+        return rewriter.notifyMatchFailure(
+            op, "cannot unroll a scalable dimension in dot-product lowering");
+    }
+  }
+
   auto iteratorTypes = op.getIteratorTypes().getValue();
   static constexpr std::array<int64_t, 2> perm = {1, 0};
   Location loc = op.getLoc();
diff --git a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
index d00fe588f2b5b..c96c930e0e179 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
@@ -434,6 +434,26 @@ func.func @batch_contract_fmf(%arg0: vector<2x2xf32>,
   return %0 : vector<2xf32>
 }
 
+// CHECK-LABEL: @negative_contract_to_dot_matvec_scalable
+//       CHECK: vector.contract
+func.func @negative_contract_to_dot_matvec_scalable(%arg0: vector<[2]x3xf32>,
+                        %arg1: vector<3xf32>,
+                        %arg2: vector<[2]xf32>) -> vector<[2]xf32> {
+  %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2
+    : vector<[2]x3xf32>, vector<3xf32> into vector<[2]xf32>
+  return %0 : vector<[2]xf32>
+}
+
+// CHECK-LABEL: @negative_contract_to_dot_matmat_scalable
+//       CHECK: vector.contract
+func.func @negative_contract_to_dot_matmat_scalable(%lhs: vector<[2]x2xf32>,
+                        %rhs: vector<2x[2]xf32>,
+                        %init: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
+  %res = vector.contract #matmat_trait %lhs, %rhs, %init
+    : vector<[2]x2xf32>, vector<2x[2]xf32> into vector<[2]x[2]xf32>
+  return %res : vector<[2]x[2]xf32>
+}
+
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
     %f = transform.structured.match ops{["func.func"]} in %module_op