diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h --- a/mlir/include/mlir/IR/Block.h +++ b/mlir/include/mlir/IR/Block.h @@ -55,6 +55,9 @@ /// Return if this block is the entry block in the parent region. bool isEntryBlock(); + /// Return if this block is part of a loop (within the current function) + bool isInLoop(); + /// Insert this block (which must not already be in a region) right before /// the specified block. void insertBefore(Block *block); diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp --- a/mlir/lib/IR/Block.cpp +++ b/mlir/lib/IR/Block.cpp @@ -9,7 +9,9 @@ #include "mlir/IR/Block.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Operation.h" +#include "mlir/Interfaces/LoopLikeInterface.h" #include "llvm/ADT/BitVector.h" +#include "llvm/ADT/DenseSet.h" using namespace mlir; //===----------------------------------------------------------------------===// @@ -34,6 +36,35 @@ /// Return if this block is the entry block in the parent region. bool Block::isEntryBlock() { return this == &getParent()->front(); } +/// A block is in a control flow graph loop if it can reach itself in a graph +/// traversal +static bool isInLoopImpl(Block *target, Block *current, + llvm::DenseSet &visited) { + auto [it, inserted] = visited.insert(current); + if (!inserted) + // loop detected + return current == target; + + for (Block *successor : current->getSuccessors()) + if (isInLoopImpl(target, successor, visited)) + return true; + + return false; +} + +bool Block::isInLoop() { + // The block could be inside a loop-like operation: + Operation *parent = getParentOp(); + if (isa(parent) || + parent->getParentOfType()) + return true; + + // Or the block could be inside a control flow graph loop: + llvm::DenseSet visited; + visited.reserve(this->getParent()->getBlocks().size()); + return isInLoopImpl(this, this, visited); +} + /// Insert this block (which must not already be in a region) right before the /// specified block. void Block::insertBefore(Block *block) { diff --git a/mlir/test/IR/test-block-loop.mlir b/mlir/test/IR/test-block-loop.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/IR/test-block-loop.mlir @@ -0,0 +1,113 @@ +// RUN: mlir-opt %s --mlir-disable-threading -test-block-is-in-loop 2>&1 | FileCheck %s + +module { + // Test function with only one bb + func.func @simple() { + func.return + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0: + + // Test simple loop bb0 -> bb0 + func.func @loopForever() { + ^bb0: + cf.br ^bb1 + ^bb1: + cf.br ^bb1 + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb1: + + // Test bb0 -> bb1 -> bb2 -> bb1 + func.func @loopForever2() { + ^bb0: + cf.br ^bb1 + ^bb1: + cf.br ^bb2 + ^bb2: + cf.br ^bb1 + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb1: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb2: + + // Test conditional branch without loop + // bb0 -> bb1 -> {bb2, bb3} + func.func @noLoop(%arg0: i1) { + cf.br ^bb1 + ^bb1: + cf.cond_br %arg0, ^bb2, ^bb3 + ^bb2: + func.return + ^bb3: + func.return + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0(%arg0: i1) +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb1: +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb2: +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb3: + + // test multiple loops + // bb0 -> bb1 -> bb2 -> bb3 { -> bb2} -> bb4 { -> bb1 } -> bb5 + func.func @multipleLoops(%arg0: i1, %arg1: i1) { + cf.br ^bb1 + ^bb1: + cf.br ^bb2 + ^bb2: + cf.br ^bb3 + ^bb3: + cf.cond_br %arg0, ^bb4, ^bb2 + ^bb4: + cf.cond_br %arg1, ^bb1, ^bb5 + ^bb5: + return + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0(%arg0: i1, %arg1: i1) +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb1: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb2: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb3: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb4: +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb5: + + // test derived from real Flang output + func.func @_QPblockTest0(%arg0: i1, %arg1: i1) { + cf.br ^bb1 + ^bb1: // 2 preds: ^bb0, ^bb4 + cf.cond_br %arg0, ^bb2, ^bb5 + ^bb2: // pred: ^bb1 + cf.cond_br %arg1, ^bb3, ^bb4 + ^bb3: // pred: ^bb2 + return + ^bb4: // pred: ^bb2 + cf.br ^bb1 + ^bb5: // pred: ^bb1 + return + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0(%arg0: i1, %arg1: i1) +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb1: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb2: +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb3: +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb4: +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb5: +} diff --git a/mlir/test/lib/IR/CMakeLists.txt b/mlir/test/lib/IR/CMakeLists.txt --- a/mlir/test/lib/IR/CMakeLists.txt +++ b/mlir/test/lib/IR/CMakeLists.txt @@ -1,5 +1,6 @@ # Exclude tests from libMLIR.so add_mlir_library(MLIRTestIR + TestBlocks.cpp TestBuiltinAttributeInterfaces.cpp TestClone.cpp TestDiagnostics.cpp diff --git a/mlir/test/lib/IR/TestBlocks.cpp b/mlir/test/lib/IR/TestBlocks.cpp new file mode 100644 --- /dev/null +++ b/mlir/test/lib/IR/TestBlocks.cpp @@ -0,0 +1,46 @@ +//===- TestBlocks.cpp - Pass to test Region's methods ---------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "TestDialect.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "llvm/Support/raw_ostream.h" + +using namespace mlir; + +namespace { +/// This is a test pass that tests Blocks's isInLoop method by checking if each +/// block in a function is in a loop and outputing if it is +struct IsInLoopPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IsInLoopPass) + + StringRef getArgument() const final { return "test-block-is-in-loop"; } + StringRef getDescription() const final { return "Test Block's isInLoop"; } + + void runOnOperation() override { + mlir::func::FuncOp func = getOperation(); + + for (Block &block : func.getBlocks()) { + llvm::outs() << "Block is "; + if (block.isInLoop()) + llvm::outs() << "in a loop\n"; + else + llvm::outs() << "not in a loop\n"; + block.print(llvm::outs()); + llvm::outs() << "\n"; + } + } +}; + +} // namespace + +namespace mlir { +void registerBlockTestPasses() { PassRegistration(); } +} // namespace mlir diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -29,6 +29,7 @@ // Defined in the test directory, no public header. namespace mlir { +void registerBlockTestPasses(); void registerConvertToTargetEnvPass(); void registerCloneTestPasses(); void registerPassManagerTestPass(); @@ -135,6 +136,7 @@ #ifdef MLIR_INCLUDE_TESTS void registerTestPasses() { + registerBlockTestPasses(); registerCloneTestPasses(); registerConvertToTargetEnvPass(); registerPassManagerTestPass();