102 lines · cpp
1//===- TestLinalgFusionTransforms.cpp - Test Linalg fusion patterns -------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements logic for testing Linalg fusion patterns.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/IR/AffineOps.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Dialect/Linalg/Transforms/Transforms.h"16#include "mlir/Dialect/SCF/Transforms/Patterns.h"17#include "mlir/Pass/Pass.h"18#include "mlir/Pass/PassManager.h"19#include "mlir/Transforms/GreedyPatternRewriteDriver.h"20#include "mlir/Transforms/Passes.h"21 22using namespace mlir;23using namespace mlir::linalg;24 25static LogicalResult fuseLinalgOpsGreedily(func::FuncOp f) {26 OpBuilder b(f);27 28 // Save original Linalg ops, we only want to make a pass over those.29 SmallVector<LinalgOp, 8> linalgOps;30 f.walk([&](LinalgOp op) {31 // TODO: support multi-results.32 if (op->getNumResults() <= 1)33 linalgOps.push_back(op);34 });35 36 // Tile and Fuse for tensors inputs (TODO: all tensor operands).37 bool changed = false;38 for (LinalgOp linalgOp : llvm::reverse(linalgOps)) {39 for (OpOperand &opOperand : linalgOp->getOpOperands()) {40 if (isa<MemRefType>(opOperand.get().getType()))41 continue;42 if (isa<RankedTensorType>(opOperand.get().getType())) {43 // Tile and Fuse tensor input.44 if (opOperand.getOperandNumber() >= linalgOp.getNumDpsInputs())45 continue;46 auto info = fuseProducerOfTensor(b, opOperand);47 if (failed(info))48 continue;49 auto *originalOp = info->originalProducer.getOperation();50 auto *originalOpInLinalgOpsVector = llvm::find(linalgOps, originalOp);51 *originalOpInLinalgOpsVector = info->fusedProducer;52 // Don't mark for erasure in the tensor case, let DCE handle this.53 changed = true;54 }55 }56 }57 58 return changed ? success() : failure();59}60 61namespace {62struct TestLinalgGreedyFusion63 : public PassWrapper<TestLinalgGreedyFusion, OperationPass<func::FuncOp>> {64 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgGreedyFusion)65 66 void getDependentDialects(DialectRegistry ®istry) const override {67 registry.insert<affine::AffineDialect, linalg::LinalgDialect,68 memref::MemRefDialect, scf::SCFDialect>();69 }70 StringRef getArgument() const final { return "test-linalg-greedy-fusion"; }71 StringRef getDescription() const final {72 return "Test Linalg fusion by applying a greedy test transformation.";73 }74 void runOnOperation() override {75 MLIRContext *context = &getContext();76 RewritePatternSet patterns =77 linalg::getLinalgTilingCanonicalizationPatterns(context);78 patterns.add<ExtractSliceOfPadTensorSwapPattern>(context);79 scf::populateSCFForLoopCanonicalizationPatterns(patterns);80 FrozenRewritePatternSet frozenPatterns(std::move(patterns));81 OpPassManager pm(func::FuncOp::getOperationName());82 pm.addPass(createLoopInvariantCodeMotionPass());83 pm.addPass(createCanonicalizerPass());84 pm.addPass(createCSEPass());85 do {86 (void)applyPatternsGreedily(getOperation(), frozenPatterns);87 if (failed(runPipeline(pm, getOperation())))88 this->signalPassFailure();89 } while (succeeded(fuseLinalgOpsGreedily(getOperation())));90 }91};92} // namespace93 94namespace mlir {95namespace test {96void registerTestLinalgGreedyFusion() {97 PassRegistration<TestLinalgGreedyFusion>();98}99 100} // namespace test101} // namespace mlir102