[Mlir-commits] [mlir] [mlir][x86] Decouple Accumulator from ADD-Based Vector Contractions (over loops) (PR #204327)

[Adam Siemieniuk]

================
@@ -289,27 +320,23 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
 
 // This function shuffles the vectors written by vector.contract operation
 // as a flat layout structure before they are stored.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter,
-                                       Operation *opA, Operation *opB,
-                                       int64_t nonUnitDimAcc,
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value opA,
+                                       Value opB, int64_t nonUnitDimAcc,
                                        VectorType accTy) {
-  // Helper to extract vector operand from write-like ops
-  auto getWrittenVector = [](Operation *op) -> Value {
-    if (auto write = dyn_cast<vector::TransferWriteOp>(op))
-      return write.getVector();
-    if (auto store = dyn_cast<vector::StoreOp>(op))
-      return store.getValueToStore();
-    return nullptr;
-  };
 
-  Value vecA = getWrittenVector(opA);
-  Value vecB = getWrittenVector(opB);
+  Value vecA = contractionUsersAfterYield(opA);
----------------
adam-smnk wrote:

nit: this is a bit of util regression by adding implicit constrain that ops have to be inside a loop.
Could the old and new behavior be both supported?

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