[Mlir-commits] [mlir] [mlir][Vector] Reject scalable dst dim in contract-to-dot lowering (PR #224793)
Cheng-Chieh Liu
llvmlistbot at llvm.org
Sun Sep 20 18:16:58 PDT 2026
https://github.com/jay87406 updated https://github.com/llvm/llvm-project/pull/224793
>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 1/2] [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
>From 33f47cacbaf8a30a0e7b76eaddff8e0a4c8f45a6 Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Sun, 20 Sep 2026 19:56:09 -0500
Subject: [PATCH 2/2] Simplify the scalable check, add more dimension coverage
Per review: use VectorType::isScalable() instead of manually checking
each dim, since dstType is always rank 1 or 2 here. Also add a
negative test where only the inner (N) dst dim is scalable, and a
positive test where only the reduction (K) dim is scalable -- the
latter still lowers fine since vector.reduction handles a scalable
operand natively, without unrolling.
---
.../Vector/Transforms/LowerVectorContract.cpp | 13 ++---
.../vector-contract-to-dot-transforms.mlir | 53 +++++++++++++++++++
2 files changed, 57 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 830544c552f8c..4da031c7e58cb 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -646,15 +646,10 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
// 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");
- }
- }
+ if (auto dstType = dyn_cast<VectorType>(op.getResultType());
+ dstType && dstType.isScalable())
+ 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};
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 c96c930e0e179..cade4609b94a4 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
@@ -454,6 +454,59 @@ func.func @negative_contract_to_dot_matmat_scalable(%lhs: vector<[2]x2xf32>,
return %res : vector<[2]x[2]xf32>
}
+// Only the inner (N) dst dim is scalable; the outer (M) dim is still static
+// and gets peeled off by the progressive contraction lowering, but the
+// resulting per-row contraction still can't be unrolled and stays as
+// vector.contract.
+// CHECK-LABEL: @negative_contract_to_dot_matmat_inner_scalable
+// CHECK-SAME: %[[LHS:.+]]: vector<2x3xf32>, %[[RHS:.+]]: vector<3x[2]xf32>, %[[INIT:.+]]: vector<2x[2]xf32>
+// CHECK: %[[LHS0:.+]] = vector.extract %[[LHS]][0] : vector<3xf32> from vector<2x3xf32>
+// CHECK: %[[INIT0:.+]] = vector.extract %[[INIT]][0] : vector<[2]xf32> from vector<2x[2]xf32>
+// CHECK: vector.contract {{.*}} %[[LHS0]], %[[RHS]], %[[INIT0]] : vector<3xf32>, vector<3x[2]xf32> into vector<[2]xf32>
+// CHECK: %[[LHS1:.+]] = vector.extract %[[LHS]][1] : vector<3xf32> from vector<2x3xf32>
+// CHECK: %[[INIT1:.+]] = vector.extract %[[INIT]][1] : vector<[2]xf32> from vector<2x[2]xf32>
+// CHECK: vector.contract {{.*}} %[[LHS1]], %[[RHS]], %[[INIT1]] : vector<3xf32>, vector<3x[2]xf32> into vector<[2]xf32>
+func.func @negative_contract_to_dot_matmat_inner_scalable(%lhs: vector<2x3xf32>,
+ %rhs: vector<3x[2]xf32>,
+ %init: vector<2x[2]xf32>) -> vector<2x[2]xf32> {
+ %res = vector.contract #matmat_trait %lhs, %rhs, %init
+ : vector<2x3xf32>, vector<3x[2]xf32> into vector<2x[2]xf32>
+ return %res : vector<2x[2]xf32>
+}
+
+// Only the reduction (K) dim is scalable; vector.reduction handles that
+// natively without unrolling, so this still lowers.
+// CHECK-LABEL: func @contract_to_dot_matvec_scalable_reduction
+// CHECK-SAME: %[[A:.*0]]: vector<4x[8]xf32>,
+// CHECK-SAME: %[[B:.*1]]: vector<[8]xf32>,
+// CHECK-SAME: %[[C:.*2]]: vector<4xf32>
+// CHECK: %[[R:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>
+// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<[8]xf32> from vector<4x[8]xf32>
+// CHECK: %[[T2:.*]] = arith.mulf %[[T0]], %[[B]] : vector<[8]xf32>
+// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<[8]xf32> into f32
+// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : f32 into vector<4xf32>
+// CHECK: %[[T5:.*]] = vector.extract %[[A]][1] : vector<[8]xf32> from vector<4x[8]xf32>
+// CHECK: %[[T7:.*]] = arith.mulf %[[T5]], %[[B]] : vector<[8]xf32>
+// CHECK: %[[T8:.*]] = vector.reduction <add>, %[[T7]] : vector<[8]xf32> into f32
+// CHECK: %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : f32 into vector<4xf32>
+// CHECK: %[[T10:.*]] = vector.extract %[[A]][2] : vector<[8]xf32> from vector<4x[8]xf32>
+// CHECK: %[[T12:.*]] = arith.mulf %[[T10]], %[[B]] : vector<[8]xf32>
+// CHECK: %[[T13:.*]] = vector.reduction <add>, %[[T12]] : vector<[8]xf32> into f32
+// CHECK: %[[T14:.*]] = vector.insert %[[T13]], %[[T9]] [2] : f32 into vector<4xf32>
+// CHECK: %[[T15:.*]] = vector.extract %[[A]][3] : vector<[8]xf32> from vector<4x[8]xf32>
+// CHECK: %[[T17:.*]] = arith.mulf %[[T15]], %[[B]] : vector<[8]xf32>
+// CHECK: %[[T18:.*]] = vector.reduction <add>, %[[T17]] : vector<[8]xf32> into f32
+// CHECK: %[[T19:.*]] = vector.insert %[[T18]], %[[T14]] [3] : f32 into vector<4xf32>
+// CHECK: %[[T20:.*]] = arith.addf %[[T19]], %[[C]] : vector<4xf32>
+// CHECK: return %[[T20]] : vector<4xf32>
+func.func @contract_to_dot_matvec_scalable_reduction(%arg0: vector<4x[8]xf32>,
+ %arg1: vector<[8]xf32>,
+ %arg2: vector<4xf32>) -> vector<4xf32> {
+ %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2
+ : vector<4x[8]xf32>, vector<[8]xf32> into vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
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
More information about the Mlir-commits
mailing list