[Mlir-commits] [mlir] [mlir][x86] Adds extra pattern for AMX lowering (PR #213228)

[Arun Thangamani]

https://github.com/arun-thmn created https://github.com/llvm/llvm-project/pull/213228

None

>From f889774c3e00fffef7c00c4a30a6c248c2cf3c6a Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Fri, 31 Jul 2026 02:02:53 -0700
Subject: [PATCH] adds extra pattern for AMX lowering

---
 .../VectorContractToAMXDotProduct.cpp         | 35 ++++++++++++-------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 238d03ce4f6af..28b8542d79cfb 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -495,14 +495,16 @@ static SmallVector<Value> createTileZeros(OpBuilder &rewriter, Location loc,
 static Value getIndxToLoadStoreFromPckBuffer(
     OpBuilder &rewriter, Location loc, Value ivInnerLoop, Value ivOuterLoop,
     bool isInnerLoopUBHasOddQuot, bool isInnerLoopUBLarger, bool pack,
-    unsigned int blockingFactor) {
+    Value blockStride) {
 
   Value c2 = arith::ConstantIndexOp::create(rewriter, loc, 2);
-  Value packOffset =
-      arith::ConstantIndexOp::create(rewriter, loc, (16 * blockingFactor));
 
+  // `blockStride` is the reduction (K) loop step, i.e. the amount by which the
+  // induction variable advances for one K-block. Dividing the induction value
+  // by it yields the K-block index regardless of whether the loop counts
+  // K-elements (step == 16*blockingFactor) or pre-blocked K-tiles (step == 1).
   Value quotientInnerLoop =
-      arith::DivUIOp::create(rewriter, loc, ivInnerLoop, packOffset);
+      arith::DivUIOp::create(rewriter, loc, ivInnerLoop, blockStride);
   Value remInnerLoop = arith::RemUIOp::create(
       rewriter, loc, rewriter.getIndexType(), quotientInnerLoop, c2);
 
@@ -587,8 +589,7 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
                                                      nLoadIndx, ivNewInnerLoop);
               indxToStoreInBuffer = getIndxToLoadStoreFromPckBuffer(
                   rewriter, loc, ivNewInnerLoop, ivOuterLoop,
-                  isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack,
-                  blockingFactor);
+                  isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack, step);
               Value indxToLoadFromMatB =
                   arith::AddIOp::create(rewriter, loc, indxToStoreInBuffer, c1);
               indxToLoadFromBuffer =
@@ -653,14 +654,16 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
         if (!isVnni) {
           if (outerLoop) {
             if (!pack) {
-              Value nLoadIndx = arith::ConstantIndexOp::create(
-                  rewriter, locNewInnerLoop, offset);
               matB = Value();
               indxToLoadFromBuffer = c0;
+              // Use the real spill-block induction value (== spillInnerLoop)
+              // together with the loop step so the computed ping-pong slot
+              // matches the prefetch store side for any number of register
+              // blocks, including odd counts (e.g. 96 = 3 blocks). Passing a
+              // constant here mis-parities the slot for odd block counts.
               indxToLoadFromBuffer = getIndxToLoadStoreFromPckBuffer(
-                  rewriter, loc, nLoadIndx, ivOuterLoop,
-                  isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack,
-                  blockingFactor);
+                  rewriter, loc, ivNewInnerLoop, ivOuterLoop,
+                  isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack, step);
             }
           } else {
             if (!pack) {
@@ -1045,10 +1048,18 @@ struct VectorContractToAMXDotProduct
     while (true) {
       Operation *parent = current->getParentOfType<scf::ForOp>();
 
-      if (!parent)
+      if (!parent) {
+        // The accumulator initialization can be hoisted above an enclosing
+        // parallel region (scf.parallel/scf.forall) when the register tile
+        // matches the problem size and the M/N register loops fold away. In
+        // that case the reduction loop(s) collected so far are still valid to
+        // rewrite, so stop climbing instead of bailing out.
+        if (!loopLists.empty())
+          break;
         return rewriter.notifyMatchFailure(
             contractOp,
             "Accumulator read and contract op not within scf.for op");
+      }
 
       loopLists.push_back(dyn_cast<scf::ForOp>(parent));