[Mlir-commits] [mlir] [MLIR][XeGPU] Add unroll patterns for XeGPU (1/N) (PR #137010)

[Chao Chen]

================
@@ -0,0 +1,497 @@
+//===- XeGPUUnroll.cpp - patterns to do unrolling ---------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains patterns for unrolling XeGPU operations. It follows a
+// similar concept and design as vector unroll patterns, serving as a complement
+// to them.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
+
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
+#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include <numeric>
+
+namespace mlir {
+namespace xegpu {
+#define GEN_PASS_DEF_XEGPUUNROLL
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
+} // namespace xegpu
+} // namespace mlir
+
+#define DEBUG_TYPE "xegpu-unroll"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
+using namespace mlir;
+
+namespace {
+
+template <typename SourceOp>
+struct UnrollPattern : public OpRewritePattern<SourceOp> {
+  UnrollPattern(MLIRContext *context, const xegpu::UnrollOptions &options,
+                PatternBenefit benefit = 1)
+      : OpRewritePattern<SourceOp>(context, benefit), options(options) {}
+
+protected:
+  /// Return the target shape for the given `op`. Return std::nullopt if the
+  /// op shouldn't be or cannot be unrolled.
+  std::optional<SmallVector<int64_t>> getTargetShape(Operation *op) const {
+    LDBG("");
+    LDBG("Get unroll shape for: " << *op);
+
+    if (options.filterConstraint && failed(options.filterConstraint(op))) {
+      LDBG("--no filter constraint -> BAIL");
+      return std::nullopt;
+    }
+
+    assert(options.nativeShape &&
+           "expects the native shape for native shape call back function.");
+    auto nativeShape = options.nativeShape(op);
+    return nativeShape;
+  }
+
+  std::optional<SmallVector<int64_t>>
+  computeGrids(llvm::ArrayRef<int64_t> shape,
+               llvm::ArrayRef<int64_t> subShape) const {
+    // if the shape == subshape, we don't need to unroll.
+    if (shape == subShape) {
+      LDBG("shape == subshape, no unroll");
+      return std::nullopt;
+    }
+    return computeShapeRatio(shape, subShape);
+  }
+
+  // copy the layout attribte and drops the inst_data field.
+  xegpu::LayoutAttr getLaneLevelAttrsOnly(Attribute attr) const {
+    auto layout = dyn_cast_if_present<xegpu::LayoutAttr>(attr);
+    if (!layout || layout.getLaneLayout() == nullptr)
+      return xegpu::LayoutAttr();
+    return layout.dropInstData();
+  };
+
+  SmallVector<Type> getUnrolledTypes(ShapedType type,
+                                     ArrayRef<int64_t> blockSize) const {
+    auto elemTy = type.getElementType();
+    Type newTy;
+    // TensorDescType needs to drop the inst_data field in the layout attribute
----------------
chencha3 wrote:

After unroll, the data shape is supposed to instruction size, so the inst_data is meaningless. 

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