diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h @@ -18,4 +18,13 @@ /// Include the generated interface declarations. #include "mlir/Interfaces/LoopLikeInterface.h.inc" +namespace mlir { + +/// Returns if a block is inside a loop (within the current function). This can +/// be either because the block is nested inside a LoopLikeInterface or because +/// the control flow graph is cyclic +bool blockIsInLoop(Block *block); + +} // namespace mlir + #endif // MLIR_INTERFACES_LOOPLIKEINTERFACE_H_ diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp --- a/mlir/lib/Interfaces/LoopLikeInterface.cpp +++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp @@ -7,8 +7,51 @@ //===----------------------------------------------------------------------===// #include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "llvm/ADT/DenseSet.h" using namespace mlir; /// Include the definitions of the loop-like interfaces. #include "mlir/Interfaces/LoopLikeInterface.cpp.inc" + +namespace mlir { +bool blockIsInLoop(Block *block) { + Operation *parent = block->getParentOp(); + + // The block could be inside a loop-like operation + if (isa(parent) || + parent->getParentOfType()) + return true; + + // This block might be nested inside another block, which is in a loop + if (!isa(parent)) { + if (blockIsInLoop(parent->getBlock())) { + return true; + } + } + + // Or the block could be inside a control flow graph loop: + // A block is in a control flow graph loop if it can reach itself in a graph + // traversal + llvm::DenseSet visited; + llvm::SmallVector stack; + stack.push_back(block); + while (!stack.empty()) { + Block *current = stack.pop_back_val(); + auto [it, inserted] = visited.insert(current); + if (!inserted) { + // loop detected + if (current == block) + return true; + continue; + } + + stack.reserve(stack.size() + current->getNumSuccessors()); + for (Block *successor : current->getSuccessors()) + stack.push_back(successor); + } + return false; +} + +} // namespace mlir diff --git a/mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir b/mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir @@ -0,0 +1,143 @@ +// 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: + +// check nested blocks + func.func @check_alloc_in_loop(%counter : i64) { + cf.br ^bb1(%counter: i64) + ^bb1(%lv : i64): + %cm1 = arith.constant -1 : i64 + %rem = arith.addi %lv, %cm1 : i64 + %zero = arith.constant 0 : i64 + %p = arith.cmpi eq, %rem, %zero : i64 + cf.cond_br %p, ^bb3, ^bb2 + ^bb2: + scf.execute_region -> () { + %c1 = arith.constant 1 : i64 + scf.yield + } + cf.br ^bb1(%rem: i64) + ^bb3: + return + } +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb0(%arg0: i64): +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb1(%0: i64) +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb0: +// CHECK-NEXT: %c1_i64 +// CHECK: Block is in a loop +// CHECK-NEXT: ^bb2: +// CHECK: Block is not in a loop +// CHECK-NEXT: ^bb3: +} diff --git a/mlir/test/lib/Interfaces/CMakeLists.txt b/mlir/test/lib/Interfaces/CMakeLists.txt --- a/mlir/test/lib/Interfaces/CMakeLists.txt +++ b/mlir/test/lib/Interfaces/CMakeLists.txt @@ -1 +1,2 @@ +add_subdirectory(LoopLikeInterface) add_subdirectory(TilingInterface) diff --git a/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt b/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt @@ -0,0 +1,9 @@ +add_mlir_library(MLIRLoopLikeInterfaceTestPasses + TestBlockInLoop.cpp + + EXCLUDE_FROM_LIBMLIR + + LINK_LIBS PUBLIC + MLIRPass + MLIRLoopLikeInterface + ) diff --git a/mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp b/mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp new file mode 100644 --- /dev/null +++ b/mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp @@ -0,0 +1,47 @@ +//===- TestBlockInLoop.cpp - Pass to test mlir::blockIsInLoop -------------===// +// +// 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 "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Interfaces/LoopLikeInterface.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 mlir::blockIsInLoop()"; + } + + void runOnOperation() override { + mlir::func::FuncOp func = getOperation(); + func.walk([](mlir::Block *block) { + llvm::outs() << "Block is "; + if (blockIsInLoop(block)) + 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 registerLoopLikeInterfaceTestPasses() { PassRegistration(); } +} // namespace mlir diff --git a/mlir/tools/mlir-opt/CMakeLists.txt b/mlir/tools/mlir-opt/CMakeLists.txt --- a/mlir/tools/mlir-opt/CMakeLists.txt +++ b/mlir/tools/mlir-opt/CMakeLists.txt @@ -21,6 +21,7 @@ MLIRFuncTestPasses MLIRGPUTestPasses MLIRLinalgTestPasses + MLIRLoopLikeInterfaceTestPasses MLIRMathTestPasses MLIRMemRefTestPasses MLIRNVGPUTestPasses 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 @@ -33,6 +33,7 @@ void registerCloneTestPasses(); void registerPassManagerTestPass(); void registerPrintSpirvAvailabilityPass(); +void registerLoopLikeInterfaceTestPasses(); void registerShapeFunctionTestPasses(); void registerSideEffectTestPasses(); void registerSliceAnalysisTestPass(); @@ -139,6 +140,7 @@ registerConvertToTargetEnvPass(); registerPassManagerTestPass(); registerPrintSpirvAvailabilityPass(); + registerLoopLikeInterfaceTestPasses(); registerShapeFunctionTestPasses(); registerSideEffectTestPasses(); registerSliceAnalysisTestPass();