[Mlir-commits] [llvm] [mlir] [mlir][X86] Fold VNNI operand unpacking into contractions (PR #232800)

[Adam Siemieniuk]

================
@@ -0,0 +1,318 @@
+//===- PrepareVectorContractLoops.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/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "x86-prepare-vector-contract-loops"
+
+using namespace mlir;
+
+namespace {
+vector::ContractionOp getCandidate(scf::ForOp loop) {
+  if (loop.getNumRegionIterArgs() != 1) {
+    LDBG() << "reject loop at " << loop.getLoc() << ": expected one iter_arg\n";
+    return {};
+  }
+
+  auto yield = cast<scf::YieldOp>(loop.getBody()->getTerminator());
+  for (Operation &op : loop.getBody()->without_terminator()) {
+    auto contract = dyn_cast<vector::ContractionOp>(&op);
+    if (!contract)
+      continue;
+    if (contract.getKind() != vector::CombiningKind::ADD ||
+        contract.getAcc() != loop.getRegionIterArgs().front() ||
+        !contract.getAcc().hasOneUse() || !contract.getResult().hasOneUse() ||
+        yield->getNumOperands() != 1 ||
+        yield->getOperand(0) != contract.getResult()) {
+      LDBG() << "reject contraction at " << contract.getLoc()
+             << ": not a single loop-carried accumulation\n";
+      continue;
+    }
+
+    auto iterators = contract.getIteratorTypesArray();
+    auto maps = contract.getIndexingMapsArray();
+    MLIRContext *ctx = contract.getContext();
+    auto m = getAffineDimExpr(0, ctx);
+    auto n = getAffineDimExpr(1, ctx);
+    auto k = getAffineDimExpr(2, ctx);
+    auto par = vector::IteratorType::parallel;
+    auto red = vector::IteratorType::reduction;
+    if (llvm::equal(iterators, ArrayRef{par, par, red}) &&
+        maps[0] == AffineMap::get(3, 0, {m, k}, ctx) &&
+        maps[1] == AffineMap::get(3, 0, {k, n}, ctx) &&
+        maps[2] == AffineMap::get(3, 0, {m, n}, ctx))
+      return contract;
+
+    auto v = getAffineDimExpr(3, ctx);
+    if (x86::isInVnniLayout(contract, maps) &&
+        llvm::equal(iterators, ArrayRef{par, par, red, red}) &&
+        maps[0] == AffineMap::get(4, 0, {m, k, v}, ctx) &&
+        maps[1] == AffineMap::get(4, 0, {k, n, v}, ctx) &&
+        maps[2] == AffineMap::get(4, 0, {m, n}, ctx))
+      return contract;
+    LDBG() << "reject contraction at " << contract.getLoc()
+           << ": noncanonical matmul layout or iterators\n";
+  }
+  return {};
+}
+
+bool isSimpleReadOutsideLoop(vector::TransferReadOp read, scf::ForOp loop) {
+  return read && read->getParentOp() == loop &&
+         isa<MemRefType>(read.getBase().getType()) &&
+         loop.isDefinedOutsideOfLoop(read.getBase()) &&
----------------
adam-smnk wrote:

Just curious, why does it have to be outside the loop?

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