[Mlir-commits] [mlir] [mlir][linalg] Add pass to transpose matmul op (PR #89075)

Benjamin Maxwell llvmlistbot at llvm.org
Fri Apr 19 09:33:03 PDT 2024


================
@@ -0,0 +1,161 @@
+//===- TransposeMatmul.cpp - Convert Linalg matmul to transposed variants -===//
+//
+// 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 is intended to be a simple high-level (target-agnostic) matmul
+// transposition transformation.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_LINALGTRANSPOSEMATMULPASS
+#include "mlir/Dialect/Linalg/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "linalg-transpose-matmul"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+/// Pattern to replace
+///
+///   linalg.matmul(a, b)
+///
+/// with
+///
+///   linalg.matmul_transpose_a(linalg.transpose(a), b)
+///
+/// By default A is transposed. If `transposeA` is set to false then B is
+/// transposed.
+struct TransposeMatmul final : public OpRewritePattern<linalg::MatmulOp> {
+  TransposeMatmul(MLIRContext *ctx, bool transposeA, PatternBenefit benefit = 1)
+      : OpRewritePattern(ctx, benefit), transposeA(transposeA) {}
----------------
MacDue wrote:

Nit: There's nowhere in this pass (or populate patterns) to set the benefit, so you could just drop it.
```suggestion
  TransposeMatmul(MLIRContext *ctx, bool transposeA)
      : OpRewritePattern(ctx), transposeA(transposeA) {}
```

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


More information about the Mlir-commits mailing list