[Mlir-commits] [mlir] [mlir][x86] Support multi-level unroll of contractions with non-native K extent (PR #234593)

[Adam Siemieniuk]

================
@@ -487,22 +498,78 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
     return rewriter.notifyMatchFailure(
         contract, "LHS/RHS are not vector.transfer_read ops");
 
-  auto checkTransferOp = [&](vector::TransferReadOp transferOp) {
+  auto matchOperandRead = [&](vector::TransferReadOp transferOp, bool isLhs) {
+    using IK = MLUCandidate::IndexKind;
+
     if (!isSimpleTransferOp(transferOp))
-      return rewriter.notifyMatchFailure(contract, [&](Diagnostic &diag) {
-        diag << "transfer op has non-identity permutation, mask, or "
-                "out-of-bounds check: "
-             << *transferOp;
-      });
+      return rewriter.notifyMatchFailure(
+          contract, "operand read has non-identity permutation, mask, or "
+                    "out-of-bounds check");
+
+    auto indices = transferOp.getIndices();
+    auto ivIt =
+        std::find(indices.begin(), indices.end(), accLoop.getInductionVar());
+    if (ivIt == indices.end())
+      return rewriter.notifyMatchFailure(
+          contract, "operand read does not use loop index");
+    if (std::find(ivIt + 1, indices.end(), accLoop.getInductionVar()) !=
+        indices.end())
+      return rewriter.notifyMatchFailure(
+          contract, "operand read uses loop index more than once");
+
+    unsigned firstVecDim =
+        indices.size() - transferOp.getVectorType().getRank();
+    unsigned ivPos = std::distance(indices.begin(), ivIt);
+    unsigned kPos = firstVecDim + isLhs;
+    if (ivPos >= firstVecDim) {
+      if (ivPos != kPos)
+        return rewriter.notifyMatchFailure(
+            contract,
+            "loop index is not in expected K position in operand read");
+      if (candidate.operandIndexKind != IK::Unknown &&
+          candidate.operandIndexKind != IK::Offset)
+        return rewriter.notifyMatchFailure(contract,
+                                           "inconsistent operand index kind");
+      candidate.operandIndexKind = IK::Offset;
+    } else {
+      // Temporary limitation: having a zero index in the K position means that
+      // we can apply an offset without introducing additional operations in the
+      // accumulation loop.
+      if (!isZeroInteger(indices[kPos]))
+        return rewriter.notifyMatchFailure(
+            contract, "operand read K index must be zero for block index kind");
+      if (candidate.operandIndexKind != IK::Unknown &&
+          candidate.operandIndexKind != IK::Block)
+        return rewriter.notifyMatchFailure(contract,
+                                           "inconsistent operand index kind");
+      candidate.operandIndexKind = IK::Block;
----------------
adam-smnk wrote:

I feel these checks are pretty difficult to follow. Especially mix of validation and modification of candidate status (`operandIndexKind` read as property of contraction but seems modified per read inside lambda).

Could the kind be evaluated and assigned once before further validation?
Also, I think with a bit of cleanup, `Unknown` might be not needed at all.

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