mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

   1 //===- TestLinalgDecomposeOps.cpp - Test Linalg decomposition  ------------===//
   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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements a pass for testing decomposition of Linalg ops.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "mlir/Dialect/Affine/IR/AffineOps.h"
  14 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
  15 #include "mlir/Pass/Pass.h"
  16 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
  17
  18 using namespace mlir;
  19
  20 namespace {
  21 struct TestLinalgDecomposeOps
  22     : public PassWrapper<TestLinalgDecomposeOps, OperationPass<>> {
  23   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgDecomposeOps)
  24
  25   TestLinalgDecomposeOps() = default;
  26   TestLinalgDecomposeOps(const TestLinalgDecomposeOps &pass)
  27       : PassWrapper(pass){};
  28   void getDependentDialects(DialectRegistry &registry) const override {
  29     registry.insert<affine::AffineDialect, linalg::LinalgDialect>();
  30   }
  31   StringRef getArgument() const final { return "test-linalg-decompose-ops"; }
  32   StringRef getDescription() const final {
  33     return "Test Linalg decomposition patterns";
  34   }
  35
  36   Option<bool> removeDeadArgsAndResults{
  37       *this, "remove-dead-args-and-results",
  38       llvm::cl::desc("Test patterns to erase unused operands and results"),
  39       llvm::cl::init(false)};
  40
  41   void runOnOperation() override {
  42     MLIRContext *context = &this->getContext();
  43     RewritePatternSet decompositionPatterns(context);
  44     linalg::populateDecomposeLinalgOpsPattern(decompositionPatterns,
  45                                               removeDeadArgsAndResults);
  46     if (failed(applyPatternsAndFoldGreedily(
  47             getOperation(), std::move(decompositionPatterns)))) {
  48       return signalPassFailure();
  49     }
  50   }
  51 };
  52 } // namespace
  53
  54 namespace mlir {
  55 namespace test {
  56 void registerTestLinalgDecomposeOps() {
  57   PassRegistration<TestLinalgDecomposeOps>();
  58 }
  59 } // namespace test
  60 } // namespace mlir