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

[Frank Schlimbach]

================
@@ -0,0 +1,315 @@
+//===- 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()) &&
+         read.getPermutationMap().isMinorIdentity() && !read.getMask() &&
+         !read.hasOutOfBoundsDim();
+}
+
+struct ExpandedRead {
+  vector::TransferReadOp read;
+  MemRefType type;
+  SmallVector<ReassociationIndices> reassociation;
+  int64_t lastIndex;
+  // Static view of a dynamic last dimension, or null.
+  MemRefType view;
+};
+
+std::optional<ExpandedRead> getExpandableRead(vector::TransferReadOp read,
+                                              unsigned factor,
+                                              scf::ForOp loop) {
+  if (!isSimpleReadOutsideLoop(read, loop))
+    return std::nullopt;
+  auto type = cast<MemRefType>(read.getBase().getType());
+  if (type.getElementType() != read.getVectorType().getElementType() ||
+      type.getRank() < 2)
+    return std::nullopt;
+  auto lastIndex = getConstantIntValue(read.getIndices().back());
+  if (!lastIndex || *lastIndex < 0 || *lastIndex % factor != 0)
+    return std::nullopt;
+
+  MemRefType view;
+  if (type.isDynamicDim(type.getRank() - 1)) {
+    SmallVector<int64_t> strides;
+    int64_t offset;
+    if (read.getVectorType().isScalable() ||
+        failed(type.getStridesAndOffset(strides, offset)))
+      return std::nullopt;
+    // The in-bounds read covers this prefix of the last dimension.
----------------
fschlimb wrote:

clearer now?

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