[Mlir-commits] [mlir] [mlir][x86] Support rank-reducing RHS reads in AMX online packing (PR #232279)

[Adam Siemieniuk]

================
@@ -981,15 +981,23 @@ struct VectorContractToAMXDotProduct
             });
         auto one = rewriter.getIndexAttr(1);
         SmallVector<OpFoldResult> strides(indexVals.size(), one);
+        // A rank-reducing read drops leading unit dims. Append them back for
+        // the subview.
+        SmallVector<int64_t> subviewShape(indexVals.size() - vecTy.getRank(),
+                                          1);
+        llvm::append_range(subviewShape, vecTy.getShape());
         SmallVector<OpFoldResult> sizes = getAsIndexOpFoldResult(
-            contractOp.getRhs().getDefiningOp()->getContext(),
-            vecTy.getShape());
+            contractOp.getRhs().getDefiningOp()->getContext(), subviewShape);
         auto subview = memref::SubViewOp::create(rewriter, loc, srcBuffRhs,
                                                  indexVals, sizes, strides);
+        // The source subview already accounts for the leading unit dimensions.
+        // Set them to zero for the later loads.
+        for (Value &index : MutableArrayRef<Value>(indicesRhs).drop_back(2))
----------------
adam-smnk wrote:

It's for flat layout.
I realized these offsets can be simplified further, made a small cleanup.

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