mlir/test/lib/IR/TestRegions.cpp

   1 //===- TestRegions.cpp - Pass to test Region's methods --------------------===//
   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 #include "TestDialect.h"
  10 #include "mlir/Dialect/Func/IR/FuncOps.h"
  11 #include "mlir/IR/BuiltinOps.h"
  12 #include "mlir/Pass/Pass.h"
  13
  14 using namespace mlir;
  15
  16 namespace {
  17 /// This is a test pass that tests Region's takeBody method by making the first
  18 /// function take the body of the second.
  19 struct TakeBodyPass
  20     : public PassWrapper<TakeBodyPass, OperationPass<ModuleOp>> {
  21   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TakeBodyPass)
  22
  23   StringRef getArgument() const final { return "test-take-body"; }
  24   StringRef getDescription() const final { return "Test Region's takeBody"; }
  25
  26   void runOnOperation() override {
  27     auto module = getOperation();
  28
  29     SmallVector<func::FuncOp> functions =
  30         llvm::to_vector(module.getOps<func::FuncOp>());
  31     if (functions.size() != 2) {
  32       module.emitError("Expected only two functions in test");
  33       signalPassFailure();
  34       return;
  35     }
  36
  37     functions[0].getBody().takeBody(functions[1].getBody());
  38   }
  39 };
  40
  41 } // namespace
  42
  43 namespace mlir {
  44 void registerRegionTestPasses() { PassRegistration<TakeBodyPass>(); }
  45 } // namespace mlir