[mlir] Add `mlir::blockIsInLoop()`

[tblah]

Summary

This method returns whether a block is nested inside of a loop. There
can be two kinds of loop:

1. The block is nested inside of a LoopLikeOpInterface
2. There is a cycle in the control flow graph

This will be useful for Flang's stack arrays pass, which moves array
allocations from the heap to the stack. Special handling is needed when
allocations occur inside of loops to ensure additional stack space is
not allocated on each loop iteration.

[rriddle]

We generally try to keep the core IR API free from things specific to certain interfaces/compilation flows/etc. This utility is something that should generally live near the LoopLikeOpInterface.

[tblah]

Moved to a stand alone function in LoopLikeInterface.h

[DavidTruby]

mlir/lib/Interfaces/LoopLikeInterface.cpp

31

nit: this could be

return llvm::any_of(current->getSuccessors(), [&](const Block* successor) {
   return isInLoopImpl(target, successor, visited);
});

you can ignore this suggestion though.

mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp

25

Nit: these should also say override

[tblah]

Updated to use llvm::any_of instead of a for loop and an if.

[tblah]

Thanks for review @DavidTruby

mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp

25

clang-tidy tells me that override is already implied by final.

[tblah]

@rriddle ping

[DavidTruby]

mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp

25

Interesting, I don't think that's true in the general case. It might be true here though as these functions don't say they're virtual (so the virtualness must come from a base class)
You can think of a case where a virtual function at the very base of the hierarchy is marked final; this function clearly wouldn't imply override as it isn't overriding anything.
Probably just defer to clang-tidy here though!

[tblah]

mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp

25

https://clang.llvm.org/extra/clang-tidy/checks/modernize/use-override.html#cmdoption-arg-allowoverrideandfinal

[tblah]

Updated a comment and documentation string in the test pass to make it clear
the loop detection is no longer a method in mlir::Block

[kiranchandramohan]

Thanks @tblah for this patch. I have a question/comment inline.

mlir/lib/Interfaces/LoopLikeInterface.cpp

40–43

I was thinking of a case where there is no loop in the current region, but there is a cf style loop in the outer-region. A case is presented below. I think we will not currently catch this case. May be we have to do this check at the parent levels all the way to the function-level op.

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 -> () {
      //Assume alloc is here
      %c1 = arith.constant 1 : i64
      scf.yield
    }
    cf.br ^bb1(%rem: i64)
  ^bb3:
    return
}

[tblah]

Handle case where blocks are nested inside of eachother

[tblah]

mlir/lib/Interfaces/LoopLikeInterface.cpp

40–43

Thanks for pointing out this case. I've updated the patch and added this to the tests.

[kiranchandramohan]

LGTM. I have not accepted patches in this area before. If @rriddle is busy, it will be great to get another approval from one of the other reviewers on the list before submitting.

mlir/lib/Interfaces/LoopLikeInterface.cpp

37–41

Nit: I think the braces are not required.
https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

45

Nit: Now that we are recursing to all the parents, I guess the getParentOfType<LoopLikeOpInterface> is not required.

[tblah]

Ping for any MLIR reviewers

[mehdi_amini]

mlir/lib/Interfaces/LoopLikeInterface.cpp

29

Can you please remove the recursion and use an iterative algorithm?
(Guideline in https://mlir.llvm.org/getting_started/DeveloperGuide/)

[mehdi_amini]

mlir/lib/Interfaces/LoopLikeInterface.cpp

41

I'm not sure why we should hardcode func::FuncOp, for example this wouldn't work with a llvm::FuncOp. But also what is this supposed to handle that isn't handled later?

I suspect you could try to first filter out if there is a parent implementing the interface, before looking for a loop in the CFG at every level in parent regions:

 if (parent && getParentOfType<LoopLikeOpInterface>(parent))
   return true;

// Or the block could be inside a control flow graph loop:
 llvm::DenseSet<Block *> visited;
 while (block) {
   visited.clear();
   if ( isInLoopImpl(block, block, visited)) return true;
   if (Operation *parent = block->getParentOp()) block = parent->getBlock();
   else break;
 }
 return false;

[tblah]

• Check for a parent implementing LoopLikeOpInterface before searching for a CFG loop
• Don't use recursion when searching for CFG cycles
• Stop hierarchical search at anything implementing FunctionOpInterface, not just func.func

[tblah]

mlir/lib/Interfaces/LoopLikeInterface.cpp

40–43

I was thinking of a case where there is no loop in the current region, but there is a cf style loop in the outer-region. A case is presented below. I think we will not currently catch this case. May be we have to do this check at the parent levels all the way to the function-level op.

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 -> () {
      //Assume alloc is here
      %c1 = arith.constant 1 : i64
      scf.yield
    }
    cf.br ^bb1(%rem: i64)
  ^bb3:
    return
}

41

Thanks for reviewing.

I have changed to check to see if the parent implements FunctionOpInterface and now check for a parent implementing LoopLikeOpInterface first.

This condition is supposed to handle where there is no loop in the region that the current block belongs to, but that whole region is inside a loop, for example in @kiranchandramohan's example:

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 -> () {
      //Assume alloc is here
      %c1 = arith.constant 1 : i64
      scf.yield
    }
    cf.br ^bb1(%rem: i64)
  ^bb3:
    return
}

[tblah]

@mehdi_amini @rriddle ping

[rriddle]

mlir/include/mlir/Interfaces/LoopLikeInterface.h

26 ↗

(On Diff #492040)

I'm slightly nervous about how generally this function is named. Could we make it slightly more specific to LoopLikeInterface, or even make it a static method on the interface class?

mlir/lib/Interfaces/LoopLikeInterface.cpp

20

Please use full resolution of methods:

https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions

37–38

Do you need the llvm:: on these? We re-export a bunch of stuff from llvm into the mlir namespace.

[tblah]

Thanks for the review.

Changes:

• Move to a static method in LoopLikeOpInterface
• Remove redundant namespace specification on re-exported LLVM clases

[tblah]

Fix for failed shared library builds is at https://reviews.llvm.org/D143658

[mehdi_amini]

Is there a link to a log for the failure? I can't repro on MacOS

[mehdi_amini]

Actually that may be it:

: && /Library/Developer/CommandLineTools/usr/bin/c++ -fPIC -fvisibility-inlines-hidden -Werror=date-time -Werror=unguarded-availability-new -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wc++98-compat-extra-semi -Wimplicit-fallthrough -Wcovered-switch-default -Wno-noexcept-type -Wnon-virtual-dtor -Wdelete-non-virtual-dtor -Wsuggest-override -Wstring-conversion -Wmisleading-indentation -Wctad-maybe-unsupported -fdiagnostics-color -Werror=mismatched-tags -O3 -DNDEBUG -arch arm64 -isysroot /Library/Developer/CommandLineTools/SDKs/MacOSX13.1.sdk -dynamiclib -Wl,-headerpad_max_install_names -Wl,-dead_strip -o lib/libMLIRLoopLikeInterfaceTestPasses.dylib -install_name @rpath/libMLIRLoopLikeInterfaceTestPasses.dylib tools/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeFiles/MLIRLoopLikeInterfaceTestPasses.dir/TestBlockInLoop.cpp.o  -Wl,-rpath,@loader_path/../lib  lib/libMLIRPass.dylib  lib/libMLIRAnalysis.dylib  lib/libMLIRLoopLikeInterface.dylib  lib/libMLIRCallInterfaces.dylib  lib/libMLIRControlFlowInterfaces.dylib  lib/libMLIRDataLayoutInterfaces.dylib  lib/libMLIRInferIntRangeInterface.dylib  lib/libMLIRInferTypeOpInterface.dylib  lib/libMLIRSideEffectInterfaces.dylib  lib/libMLIRViewLikeInterface.dylib  lib/libMLIRIR.dylib  lib/libMLIRSupport.dylib  lib/libLLVMSupport.dylib && :
Undefined symbols for architecture arm64:
  "mlir::detail::TypeIDResolver<mlir::func::FuncOp, void>::id", referenced from:
      (anonymous namespace)::IsInLoopPass::runOnOperation() in TestBlockInLoop.cpp.o
ld: symbol(s) not found for architecture arm64

mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt

9

The missing dependency on the FuncDialect is here, and not on the interface I believe.

[rkayaith]

mlir/include/mlir/Interfaces/LoopLikeInterface.td

102–105

mlir/lib/Interfaces/LoopLikeInterface.cpp

10

this include doesn't seem necessary

[tblah]

mlir/include/mlir/Interfaces/LoopLikeInterface.td

102–105

Thanks!

mlir/lib/Interfaces/LoopLikeInterface.cpp

10

It was for a reference to FunctionOpInterface, I have changed it to the more specific mlir/IR/FunctionInterfaces.h

mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt

9

Thanks!

[tblah]

Updated to remove dependency upon func dialect plus fix comment.

[tblah]

@rkayaith, @mehdi_amini thanks for taking a look. Am I okay to commit this with the fixes?

[rkayaith]

In D141401#4118237, @tblah wrote:

@rkayaith, @mehdi_amini thanks for taking a look. Am I okay to commit this with the fixes?

Looks good to me, assuming the build issue is fixed.

mlir/include/mlir/Interfaces/LoopLikeInterface.td

InterfaceMethod<[{
       /*methodName=*/"getSingleUpperBound",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return std::nullopt;
       }]
     >,
   ];
+
+  let extraClassDeclaration = [{
+    /// Returns if a block is inside a loop (within the current function). This
+    /// can either be because the block is nested inside a LoopLikeInterface, or
+    /// because the control flow graph is cyclic
+    static bool blockIsInLoop(Block *block);
+  }];
 }
 
 #endif // MLIR_INTERFACES_LOOPLIKEINTERFACE

mlir/lib/Interfaces/LoopLikeInterface.cpp

 //===- LoopLikeInterface.cpp - Loop-like operations in MLIR ---------------===//
 //
 // 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/Interfaces/LoopLikeInterface.h"
+#include "mlir/IR/FunctionInterfaces.h"
+#include "llvm/ADT/DenseSet.h"
 
 using namespace mlir;
 
 /// Include the definitions of the loop-like interfaces.
 #include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
+
+bool LoopLikeOpInterface::blockIsInLoop(Block *block) {
+  Operation *parent = block->getParentOp();
+
+  // The block could be inside a loop-like operation
+  if (isa<LoopLikeOpInterface>(parent) ||
+      parent->getParentOfType<LoopLikeOpInterface>())
+    return true;
+
+  // This block might be nested inside another block, which is in a loop
+  if (!isa<FunctionOpInterface>(parent))
+    if (mlir::Block *parentBlock = parent->getBlock())
+      if (blockIsInLoop(parentBlock))
+        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
+  DenseSet<Block *> visited;
+  SmallVector<Block *> 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;
+}

mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir

+// 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:
+}

mlir/test/lib/Interfaces/CMakeLists.txt

+add_subdirectory(LoopLikeInterface)
 add_subdirectory(TilingInterface)

mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt

+add_mlir_library(MLIRLoopLikeInterfaceTestPasses
+  TestBlockInLoop.cpp
+
+  EXCLUDE_FROM_LIBMLIR
+
+  LINK_LIBS PUBLIC
+  MLIRPass
+  MLIRLoopLikeInterface
+  MLIRFuncDialect
+  )

mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp

+//===- 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<IsInLoopPass, OperationPass<func::FuncOp>> {
+  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 (LoopLikeOpInterface::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<IsInLoopPass>(); }
+} // namespace mlir

mlir/tools/mlir-opt/CMakeLists.txt

set(test_libs
     MLIRAffineTransformsTestPasses
     MLIRArithTestPasses
     MLIRBufferizationTestPasses
     MLIRControlFlowTestPasses
     MLIRDLTITestPasses
     MLIRFuncTestPasses
     MLIRGPUTestPasses
     MLIRLinalgTestPasses
+    MLIRLoopLikeInterfaceTestPasses
     MLIRMathTestPasses
     MLIRMemRefTestPasses
     MLIRNVGPUTestPasses
     MLIRSCFTestPasses
     MLIRShapeTestPasses
     MLIRSPIRVTestPasses
     MLIRTensorTestPasses
     MLIRTestAnalysis

mlir/tools/mlir-opt/mlir-opt.cpp

 using namespace mlir;
 
 // Defined in the test directory, no public header.
 namespace mlir {
 void registerConvertToTargetEnvPass();
 void registerCloneTestPasses();
 void registerPassManagerTestPass();
 void registerPrintSpirvAvailabilityPass();
+void registerLoopLikeInterfaceTestPasses();
 void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerRegionTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();

 } // namespace test
 
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
   registerCloneTestPasses();
   registerConvertToTargetEnvPass();
   registerPassManagerTestPass();
   registerPrintSpirvAvailabilityPass();
+  registerLoopLikeInterfaceTestPasses();
   registerShapeFunctionTestPasses();
   registerSideEffectTestPasses();
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
   registerRegionTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();