[Mlir-commits] [mlir] [mlir][x86] Add pattern to tile and unroll vector.contract (PR #229076)

[Adam Siemieniuk]

================
@@ -0,0 +1,809 @@
+//===- VectorContractMultiLevelUnroll.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/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
+
+#define DEBUG_TYPE "x86-vector-contract-multi-level-unroll"
+
+using namespace mlir;
+
+using llvm::all_of;
+using llvm::count;
+
+// Discardable attribute to control vector unroll patterns.
+static constexpr auto nativeShapeAttrName = "x86_vcmlu_native_shape";
+
+namespace {
+// Shared state for the pattern match and IR transformation.
+struct MLUCandidate {
+  // The CPU feature flag to target with the transformation, in the spelling
+  // used by llvm/include/llvm/TargetParser/X86TargetParser.def, e.g.
+  // "amx-bf16". Determines native shapes and supported data types for operands
+  // and the accumulator.
+  StringRef target;
+
+  // Anchor operation.
+  vector::ContractionOp contract;
+  bool isInVnniLayout = false;
+
+  // Shape of the original contraction.
+  VectorType accType;
+  VectorType lhsType;
+  VectorType rhsType;
+
+  // Shape of the register-tiled contraction. We'll insert a loop nest to cover
+  // the original shape.
+  VectorType accRegTileType;
+  VectorType lhsRegTileType;
+  VectorType rhsRegTileType;
+
+  // Shape of the native contraction (for the X86 nanokernel patterns). We'll
+  // unroll the register-tiled contraction to this shape.
+  VectorType accNativeType;
+  VectorType lhsNativeType;
+  VectorType rhsNativeType;
+  SmallVector<int64_t, 4> contractNativeShape;
+
+  // Bookkeeping for the accumulation loop and operands.
+  scf::ForOp accLoop;
+  Value accInitVal;
+
+  bool accIsZeroInit = false;
+  vector::TransferReadOp lhsRead;
+  vector::TransferReadOp rhsRead;
+
+  // Memory backing the accumulator.
+  Value accMemIn;
+  Value accMemOut;
+
+  // Reconstructed accumulation loop with native-shaped iter_args.
+  scf::ForOp slicedAccLoop;
+
+  // Keep track of operations that should be deleted after transformation
+  // (depending on which ops could be integrated into the new loop nest).
+  SmallVector<Operation *, 2> opsToDelete;
----------------
adam-smnk wrote:

nit: I think it's more generic as a struct field to just leave it with unspecified size `<Operation *>`

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