[Pass] Add matmul-to-mmt4d

lib/Transform/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(Affine)
-add_subdirectory(Arith)
+add_subdirectory(Arith)
+add_subdirectory(Linalg)

lib/Transform/Linalg/CMakeLists.txt

@@ -0,0 +1,8 @@
+add_mlir_dialect_library(MatmulToMmt4d
+    MatmulToMmt4d.h
+    MatmulToMmt4d.cpp
+
+    ${PROJECT_SOURCE_DIR}/lib/Transform/Linalg/
+    ADDITIONAL_HEADER_DIRS
+    LINK_LIBS PUBLIC
+)

lib/Transform/Linalg/MatmulToMmt4d.cpp

@@ -0,0 +1,134 @@
+#include "lib/Transform/Linalg/MatmulToMmt4d.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallVector.h"
+#include <cstdint>
+
+namespace mlir {
+namespace dummy {
+
+static llvm::cl::opt<int64_t>
+    clMTile("dummy-m-tile", llvm::cl::desc("Inner tile size of M dimension"),
+            llvm::cl::init(32));
+
+static llvm::cl::opt<int64_t>
+    clNTile("dummy-n-tile", llvm::cl::desc("Inner tile size of N dimension"),
+            llvm::cl::init(32));
+
+static llvm::cl::opt<int64_t>
+    clKTile("dummy-k-tile", llvm::cl::desc("Inner tile size of K dimension"),
+            llvm::cl::init(32));
+
+struct Matmul : public OpRewritePattern<linalg::MatmulOp> {
+    Matmul(mlir::MLIRContext *context)
+        : OpRewritePattern<linalg::MatmulOp>(context, /*benefit=*/1) {}
+
+    LogicalResult matchAndRewrite(linalg::MatmulOp op,
+                                  PatternRewriter &rewriter) const override {
+        int64_t M0 = clMTile;
+        int64_t N0 = clNTile;
+        int64_t K0 = clKTile;
+
+        // DPS here means Destination Passing Style
+        // retrieves the input operands
+        auto inputs = op.getDpsInputOperands();
+        // retrieves the DPS accumulator/init
+        auto outputs = op.getDpsInits();
+
+        // gets the type of given tensor by casting it to RankedTensorType
+        auto lhsType = cast<RankedTensorType>(inputs[0]->get().getType());
+        auto rhsType = cast<RankedTensorType>(inputs[1]->get().getType());
+        auto resultType = cast<RankedTensorType>(outputs[0].getType());
+
+        Location loc = op.getLoc();
+        Value paddingValue = rewriter.create<arith::ConstantOp>(
+            loc, rewriter.getZeroAttr(lhsType.getElementType()));
+
+        // returns the dimension of given tensor value
+        llvm::SmallVector<OpFoldResult> lhsSourceDims =
+            tensor::getMixedSizes(rewriter, loc, inputs[0]->get());
+        // returns the ArrayAttr as a OpFoldResult
+        llvm::SmallVector<OpFoldResult> lhsTileSizes =
+            getAsOpFoldResult(rewriter.getI64ArrayAttr({M0, K0}));
+        SmallVector<int64_t> lhsInnerDimsPos = {0, 1};
+        // returns the shape that the pack result would result in
+        SmallVector<OpFoldResult> lhsResultDims =
+            linalg::PackOp::getResultShape(rewriter, loc, lhsSourceDims,
+                                           lhsTileSizes, lhsInnerDimsPos,
+                                           lhsInnerDimsPos);
+        tensor::EmptyOp emptyOp0 = rewriter.create<tensor::EmptyOp>(
+            loc, lhsResultDims, lhsType.getElementType());
+        linalg::PackOp lhsPack = rewriter.create<linalg::PackOp>(
+            loc, inputs[0]->get(), emptyOp0, lhsInnerDimsPos, lhsTileSizes,
+            paddingValue, lhsInnerDimsPos);
+
+        llvm::SmallVector<OpFoldResult> rhsSourceDims =
+            tensor::getMixedSizes(rewriter, loc, inputs[1]->get());
+        llvm::SmallVector<OpFoldResult> rhsTileSizes =
+            getAsOpFoldResult(rewriter.getI64ArrayAttr({N0, K0}));
+        SmallVector<int64_t> rhsInnerDimsPos = {1, 0};
+        SmallVector<OpFoldResult> rhsResultDims =
+            linalg::PackOp::getResultShape(rewriter, loc, rhsSourceDims,
+                                           rhsTileSizes, rhsInnerDimsPos,
+                                           rhsInnerDimsPos);
+        tensor::EmptyOp emptyOp1 = rewriter.create<tensor::EmptyOp>(
+            loc, rhsResultDims, rhsType.getElementType());
+        linalg::PackOp rhsPack = rewriter.create<linalg::PackOp>(
+            loc, inputs[1]->get(), emptyOp1, rhsInnerDimsPos, rhsTileSizes,
+            paddingValue, rhsInnerDimsPos);
+
+        llvm::SmallVector<OpFoldResult> resSourceDims =
+            tensor::getMixedSizes(rewriter, loc, outputs[0]);
+        llvm::SmallVector<OpFoldResult> resTileSizes =
+            getAsOpFoldResult(rewriter.getI64ArrayAttr({M0, N0}));
+        SmallVector<int64_t> resInnerDimsPos = {0, 1};
+        SmallVector<OpFoldResult> resResultDims =
+            linalg::PackOp::getResultShape(rewriter, loc, resSourceDims,
+                                           resTileSizes, resInnerDimsPos,
+                                           resInnerDimsPos);
+        tensor::EmptyOp emptyOp2 = rewriter.create<tensor::EmptyOp>(
+            loc, resResultDims, resultType.getElementType());
+        linalg::PackOp resPack = rewriter.create<linalg::PackOp>(
+            loc, outputs[0], emptyOp2, resInnerDimsPos, resTileSizes,
+            paddingValue, resInnerDimsPos);
+
+        // ValueRange is just a view over the underlying data
+        // It does not hold the actual ownership of the data
+        linalg::Mmt4DOp mmt4d = rewriter.create<linalg::Mmt4DOp>(
+            loc, resPack.getResult().getType(),
+            ValueRange{lhsPack->getResult(0), rhsPack->getResult(0)},
+            ValueRange{resPack->getResult(0)});
+
+        llvm::SmallVector<OpFoldResult> mmt4dDims =
+            tensor::getMixedSizes(rewriter, loc, mmt4d.getDpsInits()[0]);
+        tensor::EmptyOp emptyOp3 = rewriter.create<tensor::EmptyOp>(
+            loc, resSourceDims, resultType.getElementType());
+        linalg::UnPackOp unpack = rewriter.create<linalg::UnPackOp>(
+            loc, mmt4d->getResult(0), emptyOp3, resInnerDimsPos, resTileSizes,
+            resInnerDimsPos);
+
+        // This repalces the uses of MatmulOp with UnpackOp
+        rewriter.replaceAllOpUsesWith(op, unpack);
+        // erases the MatmulOp
+        rewriter.eraseOp(op);
+
+        return success();
+    }
+};
+
+void MatmulToMmt4dPass::runOnOperation() {
+    mlir::RewritePatternSet patterns(&getContext());
+    patterns.insert<Matmul>(&getContext());
+    (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+}
+} // namespace dummy
+} // namespace mlir

lib/Transform/Linalg/MatmulToMmt4d.h

@@ -0,0 +1,25 @@
+#ifndef LIB_TRANSFORM_ARITH_MATMULTOMMT4D_H_
+#define LIB_TRANSFORM_ARITH_MATMULTOMMT4D_H_
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace dummy {
+
+class MatmulToMmt4dPass
+    : public PassWrapper<MatmulToMmt4dPass, OperationPass<mlir::func::FuncOp>> {
+  private:
+    void runOnOperation() override;
+
+    StringRef getArgument() const final { return "matmul-to-mmt4d"; }
+
+    StringRef getDescription() const final {
+        return "Convert linalg.matmul to linalg.mmt4d";
+    }
+};
+
+} // namespace dummy
+} // namespace mlir
+
+#endif // LIB_TRANSFORM_ARITH_MATMULTOMMT4D_H_

tests/matmul_to_mmt4d.mlir

@@ -0,0 +1,14 @@
+// RUN: dummy-opt --matmul-to-mmt4d %s | FileCheck %s
+
+// CHECK-LABEL: matmul_f32
+func.func @matmul_f32(%lhs: tensor<?x?xf32>, %rhs: tensor<?x?xf32>, %acc: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %result = linalg.matmul ins(%lhs, %rhs: tensor<?x?xf32>, tensor<?x?xf32>) outs(%acc: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %result: tensor<?x?xf32>
+}
+
+// CHECK: linalg.pack
+// CHECK: linalg.pack
+// CHECK: linalg.pack
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.mmt4d
+// CHECK: linalg.unpack

tools/CMakeLists.txt

@@ -8,6 +8,7 @@ set (LIBS
     ${conversion_libs}
     MLIRPoly10x
     AffineFullUnroll
+    MatmulToMmt4d
     MulToAdd
     MLIROptLib
     MLIRPass

tools/dummy-opt.cpp

@@ -6,6 +6,7 @@
 
 #include "lib/Transform/Affine/Passes.h"
 #include "lib/Transform/Arith/MulToAdd.h"
+#include "lib/Transform/Linalg/MatmulToMmt4d.h"
 
 #include "lib/Dialect/Poly10x/Poly10xDialect.h"
 
@@ -23,6 +24,7 @@ int main(int argc, char **argv) {
     mlir::dummy::registerAffinePasses();
     // register hand-authored passes
     mlir::PassRegistration<mlir::dummy::MulToAddPass>();
+    mlir::PassRegistration<mlir::dummy::MatmulToMmt4dPass>();
 
     return mlir::asMainReturnCode(
         mlir::MlirOptMain(argc, argv, "Dummy Pass Driver", registry));