mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp

   1 //===- TestAffineLoopUnswitching.cpp - Test affine if/else hoisting -------===//
   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 to hoist affine if/else structures.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "mlir/Dialect/Affine/Analysis/Utils.h"
  14 #include "mlir/Dialect/Affine/IR/AffineOps.h"
  15 #include "mlir/Dialect/Affine/Utils.h"
  16 #include "mlir/Pass/Pass.h"
  17 #include "mlir/Transforms/Passes.h"
  18
  19 #define PASS_NAME "test-affine-loop-unswitch"
  20
  21 using namespace mlir;
  22 using namespace mlir::affine;
  23
  24 namespace {
  25
  26 /// This pass applies the permutation on the first maximal perfect nest.
  27 struct TestAffineLoopUnswitching
  28     : public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> {
  29   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopUnswitching)
  30
  31   StringRef getArgument() const final { return PASS_NAME; }
  32   StringRef getDescription() const final {
  33     return "Tests affine loop unswitching / if/else hoisting";
  34   }
  35   TestAffineLoopUnswitching() = default;
  36   TestAffineLoopUnswitching(const TestAffineLoopUnswitching &pass) = default;
  37
  38   void runOnOperation() override;
  39
  40   /// The maximum number of iterations to run this for.
  41   constexpr static unsigned kMaxIterations = 5;
  42 };
  43
  44 } // namespace
  45
  46 void TestAffineLoopUnswitching::runOnOperation() {
  47   // Each hoisting invalidates a lot of IR around. Just stop the walk after the
  48   // first if/else hoisting, and repeat until no more hoisting can be done, or
  49   // the maximum number of iterations have been run.
  50   Operation *op = getOperation();
  51   unsigned i = 0;
  52   do {
  53     auto walkFn = [](AffineIfOp op) {
  54       return succeeded(hoistAffineIfOp(op)) ? WalkResult::interrupt()
  55                                             : WalkResult::advance();
  56     };
  57     if (op->walk(walkFn).wasInterrupted())
  58       break;
  59   } while (++i < kMaxIterations);
  60 }
  61
  62 namespace mlir {
  63 void registerTestAffineLoopUnswitchingPass() {
  64   PassRegistration<TestAffineLoopUnswitching>();
  65 }
  66 } // namespace mlir