[Mlir-commits] [mlir] [mlir][vector] Fix invalid IR in `ContractionOpLowering` (PR #78130)

[Nicolas Vasilache]

================
@@ -458,49 +450,70 @@ struct UnrolledOuterProductGenerator
     return res;
   }
 
+  std::optional<int64_t> getReductionSize(VectorType vecType,
+                                          int64_t reductionDim) {
+    // Cannot unroll scalable dimension.
+    if (vecType.getScalableDims()[reductionDim])
+      return std::nullopt;
+    int64_t reductionSize = vecType.getDimSize(reductionDim);
+    assert(reductionSize > 0 &&
+           "Reduction dim must be a known static size to allow unrolling");
+    return reductionSize;
+  }
+
   /// Two outer parallel, one inner reduction (matmat flavor).
   FailureOr<Value> matmat() {
     if (!iters({Par(), Par(), Red()}))
       return failure();
     // Set up the parallel/reduction structure in the right form.
     AffineExpr m, n, k;
     bindDims(rewriter.getContext(), m, n, k);
-    Value transposedMask = t(mask, {2, 0, 1});
+
     // Classical row-major matmul:  Just permute the lhs.
     if (layout({{m, k}, {k, n}, {m, n}}))
-      return outerProd(t(lhs), rhs, res, lhsType, /*reductionDim=*/1,
-                       transposedMask);
+      if (auto reductionSize = getReductionSize(lhsType, 1))
+        return outerProd(t(lhs), rhs, res, lhsType, *reductionSize,
+                         t(mask, {2, 0, 1}));
     // TODO: may be better to fail and use some vector<k> -> scalar reduction.
     if (layout({{m, k}, {n, k}, {m, n}})) {
-      Value tlhs = t(lhs);
-      return outerProd(tlhs, t(rhs), res, lhsType, /*reductionDim=*/1,
-                       transposedMask);
+      if (auto reductionSize = getReductionSize(lhsType, 1)) {
+        Value tlhs = t(lhs);
+        return outerProd(tlhs, t(rhs), res, lhsType, *reductionSize,
----------------
nicolasvasilache wrote:

These tests have now become non-deterministic across different systems as the order of the commas is not specified.

The line just before was a way to remedy this, you'll need to calculate the mask (or t(rhs)) ahead of time in each case.

we could sprinkle CHECK-DAG in various places like you did but this does not scale, every test writer will need to be aware of this or they will end up writing new tests that work on their system and be very puzzled when build bots come back red.

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