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

[Adam Siemieniuk]

================
@@ -0,0 +1,114 @@
+//===- MoveAccumulatorForContractLoop.cpp ---------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+#include "mlir/Dialect/X86/X86Dialect.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Casting.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+using namespace mlir::x86;
+
+namespace {
+// Transforms vector.contract(A, B, Acc) into vector.contract(A, B, 0) + Acc
+// to decouple the contraction computation from the accumulator update.
+struct MoveAccumulatorForContractLoop
+    : public OpRewritePattern<vector::ContractionOp> {
+  using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
+                                PatternRewriter &rewriter) const override {
+
+    if (contractOp.getKind() != vector::CombiningKind::ADD)
+      return rewriter.notifyMatchFailure(contractOp,
+                                         "Expects add combining kind.");
+
+    Operation *accReadOp =
+        traceToVectorReadLikeParentOperation(contractOp.getAcc());
+
+    Value contractValue = contractionUsersAfterYield(contractOp.getResult());
+
+    if (!contractValue)
+      return rewriter.notifyMatchFailure(
+          contractOp, "Final acc write might have multiple users.");
+
+    Operation *resultUserOp = *contractValue.getUsers().begin();
+
+    if (!accReadOp || !resultUserOp)
+      return rewriter.notifyMatchFailure(
+          contractOp, "Read from acc matrix is not by "
+                      "transfer_read/load/constant_zero or multiple users of "
+                      "contract operation.");
+
+    if (dyn_cast<arith::ConstantOp>(accReadOp))
+      return rewriter.notifyMatchFailure(
+          contractOp,
+          "The input acc to contract is already a constant vector.");
+
+    if ((accReadOp->getBlock() == contractOp->getBlock()) ||
+        (resultUserOp->getBlock() == contractOp->getBlock()))
+      return rewriter.notifyMatchFailure(
+          contractOp, "Acc read/write should be in a separate block.");
+
+    // Replace acc of a contraction operation with vector constant.
+    Value accValue = accReadOp->getResult(0);
+    Operation *firstUser = *accValue.getUsers().begin();
+    rewriter.setInsertionPoint(firstUser);
+
+    auto vecTy = llvm::dyn_cast<VectorType>(accValue.getType());
+    if (!vecTy)
+      return rewriter.notifyMatchFailure(contractOp, "Expects vector type.");
+
+    Location loc = accReadOp->getLoc();
+    Type elemTy = vecTy.getElementType();
+
+    Value zeroVec = arith::ConstantOp::create(
+        rewriter, loc,
+        DenseElementsAttr::get(vecTy, rewriter.getZeroAttr(elemTy)));
+
+    accValue.replaceAllUsesWith(zeroVec);
----------------
adam-smnk wrote:

That's not safe when there are other users.

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