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

[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>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp (+12) 
- (modified) mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 2a152ef0634d71..830544c552f8c0 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 d00fe588f2b5bf..c96c930e0e1799 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

``````````

</details>


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