[mlir][Analysis] Introduce LoopInfo in mlir

[Dinistro]

Summary

This commit introduces an instantiation of LLVM's LoopInfo for CFGs in
MLIR. To test the LoopInfo, a test pass is added the checks the analysis
results for a set of CFGs.

[Dinistro]

mlir/include/mlir/Analysis/LoopInfo.h

30 ↗

(On Diff #510011)

I'm not sure if we should move these traits to `mlir/IR/RegionGraphTraits.h" file.

[ftynse]

mlir/include/mlir/Analysis/LoopInfo.h

1 ↗

(On Diff #510011)

Nit: I'm confused by how this is an analysis for regions, maybe for region bodies?

30 ↗

(On Diff #510011)

Yep, it already has specializations for non-const versions.

65 ↗

(On Diff #510011)

Please call this something like BlockLoop or CFGLoop, MLIR has many other kinds of loops and we don't want these things to be confused. Also provide documentation for all top-level entities.

mlir/lib/Analysis/LoopInfo.cpp

13 ↗

(On Diff #510011)

Nit: no need to prefix with ::mlir:: within MLIR's proper code base.

mlir/test/lib/Analysis/TestLoopInfo.cpp

45 ↗

(On Diff #510011)

Nit: check that the region is not empty and error out gracefully so there's no assertion on unexpected IR.

[Dinistro]

Renamed to CFGLoopInfo and addressed additional reviewer comments.

mlir/include/mlir/Analysis/CFGLoopInfo.h

+//===- CFGLoopInfo.h - LoopInfo analysis for region bodies ------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the CFGLoopInfo analysis for MLIR. The CFGLoopInfo is used
+// to identify natural loops and determine the loop depth of various nodes of a
+// CFG.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_LOOPINFO_H
+#define MLIR_ANALYSIS_LOOPINFO_H
+
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/RegionGraphTraits.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopInfoImpl.h"
+
+namespace mlir {
+
+/// Representation of a single loop formed by blocks. The inherited LoopBase
+/// class provides accessors to the loop analysis.
+class CFGLoop : public llvm::LoopBase<mlir::Block, mlir::CFGLoop> {
+private:
+  explicit CFGLoop(mlir::Block *block);
+
+  friend class llvm::LoopBase<mlir::Block, CFGLoop>;
+  friend class llvm::LoopInfoBase<mlir::Block, CFGLoop>;
+};
+
+/// An LLVM LoopInfo instantiation for MLIR that provides access to CFG loops
+/// found in the dominator tree.
+class CFGLoopInfo : public llvm::LoopInfoBase<mlir::Block, mlir::CFGLoop> {
+public:
+  CFGLoopInfo(const llvm::DominatorTreeBase<mlir::Block, false> &domTree);
+};
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_LOOPINFO_H

mlir/include/mlir/IR/RegionGraphTraits.h

static inline ChildIteratorType child_begin(NodeRef node) {
     return node->pred_begin();
   }
   static inline ChildIteratorType child_end(NodeRef node) {
     return node->pred_end();
   }
 };
 
 template <>
+struct GraphTraits<const mlir::Block *> {
+  using ChildIteratorType = mlir::Block::succ_iterator;
+  using Node = const mlir::Block;
+  using NodeRef = Node *;
+
+  static NodeRef getEntryNode(NodeRef node) { return node; }
+
+  static ChildIteratorType child_begin(NodeRef node) {
+    return const_cast<mlir::Block *>(node)->succ_begin();
+  }
+  static ChildIteratorType child_end(NodeRef node) {
+    return const_cast<mlir::Block *>(node)->succ_end();
+  }
+};
+
+template <>
+struct GraphTraits<Inverse<const mlir::Block *>> {
+  using ChildIteratorType = mlir::Block::pred_iterator;
+  using Node = const mlir::Block;
+  using NodeRef = Node *;
+
+  static NodeRef getEntryNode(Inverse<NodeRef> inverseGraph) {
+    return inverseGraph.Graph;
+  }
+
+  static ChildIteratorType child_begin(NodeRef node) {
+    return const_cast<mlir::Block *>(node)->pred_begin();
+  }
+  static ChildIteratorType child_end(NodeRef node) {
+    return const_cast<mlir::Block *>(node)->pred_end();
+  }
+};
+
+template <>
 struct GraphTraits<mlir::Region *> : public GraphTraits<mlir::Block *> {
   using GraphType = mlir::Region *;
   using NodeRef = mlir::Block *;
 
   static NodeRef getEntryNode(GraphType fn) { return &fn->front(); }
 
   using nodes_iterator = pointer_iterator<mlir::Region::iterator>;
   static nodes_iterator nodes_begin(GraphType fn) {

mlir/lib/Analysis/CFGLoopInfo.cpp

+//===- CFGLoopInfo.cpp - LoopInfo analysis for region bodies --------------===//
+//
+// 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/Analysis/CFGLoopInfo.h"
+
+using namespace mlir;
+
+CFGLoop::CFGLoop(mlir::Block *block)
+    : llvm::LoopBase<mlir::Block, CFGLoop>(block) {}
+
+CFGLoopInfo::CFGLoopInfo(
+    const llvm::DominatorTreeBase<mlir::Block, false> &domTree) {
+  analyze(domTree);
+}

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   CallGraph.cpp
   DataLayoutAnalysis.cpp
   FlatLinearValueConstraints.cpp
   Liveness.cpp
+  CFGLoopInfo.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   DataFlow/ConstantPropagationAnalysis.cpp
   DataFlow/DeadCodeAnalysis.cpp
   DataFlow/DenseAnalysis.cpp
   DataFlow/IntegerRangeAnalysis.cpp
   DataFlow/SparseAnalysis.cpp
   )
 
 add_subdirectory(Presburger)
 
 add_mlir_library(MLIRAnalysis
   AliasAnalysis.cpp
   CallGraph.cpp
   DataFlowFramework.cpp
   DataLayoutAnalysis.cpp
   FlatLinearValueConstraints.cpp
   Liveness.cpp
+  CFGLoopInfo.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   DataFlow/ConstantPropagationAnalysis.cpp
   DataFlow/DeadCodeAnalysis.cpp
   DataFlow/DenseAnalysis.cpp
   DataFlow/IntegerRangeAnalysis.cpp

mlir/test/Analysis/test-cfg-loop-info.mlir

+// RUN: mlir-opt -pass-pipeline="builtin.module(any(test-cfg-loop-info))" --split-input-file %s 2>&1 | FileCheck %s
+
+// CHECK-LABEL: Testing : "no_loop_single_block"
+// CHECK: no loops
+func.func @no_loop_single_block() {
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Testing : "no_loop"
+// CHECK: no loops
+func.func @no_loop() {
+  cf.br ^bb1
+^bb1:
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Testing : "simple_loop"
+// CHECK-NEXT: Blocks : ^[[BB0:.*]], ^[[BB1:.*]], ^[[BB2:.*]], ^[[BB3:.*]]
+// CHECK: Loop at depth 1 containing:
+// CHECK-SAME: ^[[BB1]]<header><exiting>
+// CHECK-SAME: ^[[BB2]]<latch>
+func.func @simple_loop(%c: i1) {
+  cf.br ^bb1
+^bb1:
+  cf.cond_br %c, ^bb2, ^bb3
+^bb2:
+  cf.br ^bb1
+^bb3:
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Testing : "single_block_loop"
+// CHECK-NEXT: Blocks : ^[[BB0:.*]], ^[[BB1:.*]], ^[[BB2:.*]]
+// CHECK: Loop at depth 1 containing:
+// CHECK-SAME: ^[[BB1]]<header><latch><exiting>
+func.func @single_block_loop(%c: i1) {
+  cf.br ^bb1
+^bb1:
+  cf.cond_br %c, ^bb1, ^bb2
+^bb2:
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Testing : "nested_loop"
+// CHECK-NEXT: Blocks : ^[[BB0:.*]], ^[[BB1:.*]], ^[[BB2:.*]], ^[[BB3:.*]], ^[[BB4:.*]]
+// CHECK: Loop at depth 1
+// CHECK-SAME: ^[[BB1]]<header><exiting>
+// CHECK-SAME: ^[[BB2]]<latch>
+// CHECK-SAME: ^[[BB3]]
+// CHECK: Loop at depth 2
+// CHECK-SAME: ^[[BB2]]<header><exiting>
+// CHECK-SAME: ^[[BB3]]<latch>
+func.func @nested_loop(%c: i1) {
+  cf.br ^bb1
+^bb1:
+  cf.cond_br %c, ^bb2, ^bb4
+^bb2:
+  cf.cond_br %c, ^bb1, ^bb3
+^bb3:
+  cf.br ^bb2
+^bb4:
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Testing : "multi_latch"
+// CHECK-NEXT: Blocks : ^[[BB0:.*]], ^[[BB1:.*]], ^[[BB2:.*]], ^[[BB3:.*]], ^[[BB4:.*]]
+// CHECK: Loop at depth 1
+// CHECK-SAME: ^[[BB1]]<header><exiting>
+// CHECK-SAME: ^[[BB2]]<latch>
+// CHECK-SAME: ^[[BB3]]<latch>
+func.func @multi_latch(%c: i1) {
+  cf.br ^bb1
+^bb1:
+  cf.cond_br %c, ^bb4, ^bb2
+^bb2:
+  cf.cond_br %c, ^bb1, ^bb3
+^bb3:
+  cf.br ^bb1
+^bb4:
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Testing : "multiple_loops"
+// CHECK-NEXT: Blocks : ^[[BB0:.*]], ^[[BB1:.*]], ^[[BB2:.*]], ^[[BB3:.*]], ^[[BB4:.*]], ^[[BB5:.*]]
+// CHECK: Loop at depth 1
+// CHECK-SAME: ^[[BB3]]<header><exiting>
+// CHECK-SAME: ^[[BB4]]<latch>
+// CHECK: Loop at depth 1
+// CHECK-SAME: ^[[BB1]]<header>
+// CHECK-SAME: ^[[BB2]]<latch><exiting>
+func.func @multiple_loops(%c: i1) {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.cond_br %c, ^bb3, ^bb1
+^bb3:
+  cf.cond_br %c, ^bb5, ^bb4
+^bb4:
+  cf.br ^bb3
+^bb5:
+  return
+}

mlir/test/lib/Analysis/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestAnalysis
   TestAliasAnalysis.cpp
   TestCallGraph.cpp
   TestDataFlowFramework.cpp
   TestLiveness.cpp
+  TestCFGLoopInfo.cpp
   TestMatchReduction.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestSlice.cpp
 
   DataFlow/TestDeadCodeAnalysis.cpp
   DataFlow/TestDenseDataFlowAnalysis.cpp

mlir/test/lib/Analysis/TestCFGLoopInfo.cpp

+//===- TestCFGLoopInfo.cpp - Test CFG loop info analysis ------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements logic for testing the CFGLoopInfo analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/CFGLoopInfo.h"
+#include "mlir/IR/FunctionInterfaces.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// A testing pass that applies the CFGLoopInfo analysis on a region and prints
+/// the information it collected to llvm::errs().
+struct TestCFGLoopInfo
+    : public PassWrapper<TestCFGLoopInfo, InterfacePass<FunctionOpInterface>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCFGLoopInfo)
+
+  StringRef getArgument() const final { return "test-cfg-loop-info"; }
+  StringRef getDescription() const final {
+    return "Test the loop info analysis.";
+  }
+
+  void runOnOperation() override;
+};
+} // namespace
+
+void TestCFGLoopInfo::runOnOperation() {
+  auto func = getOperation();
+  DominanceInfo &domInfo = getAnalysis<DominanceInfo>();
+  Region &region = func.getFunctionBody();
+
+  // Prints the label of the test.
+  llvm::errs() << "Testing : " << func.getNameAttr() << "\n";
+  if (region.empty()) {
+    llvm::errs() << "empty region\n";
+    return;
+  }
+
+  // Print all the block identifiers first such that the tests can match them.
+  llvm::errs() << "Blocks : ";
+  region.front().printAsOperand(llvm::errs());
+  for (auto &block : region.getBlocks()) {
+    llvm::errs() << ", ";
+    block.printAsOperand(llvm::errs());
+  }
+  llvm::errs() << "\n";
+
+  if (region.getBlocks().size() == 1) {
+    llvm::errs() << "no loops\n";
+    return;
+  }
+
+  llvm::DominatorTreeBase<mlir::Block, false> &domTree =
+      domInfo.getDomTree(&region);
+  mlir::CFGLoopInfo loopInfo(domTree);
+
+  if (loopInfo.getTopLevelLoops().empty())
+    llvm::errs() << "no loops\n";
+  else
+    loopInfo.print(llvm::errs());
+}
+
+namespace mlir {
+namespace test {
+void registerTestCFGLoopInfoPass() { PassRegistration<TestCFGLoopInfo>(); }
+} // namespace test
+} // namespace mlir

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

 void registerTestInterfaces();
 void registerTestLastModifiedPass();
 void registerTestLinalgDecomposeOps();
 void registerTestLinalgElementwiseFusion();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
+void registerTestCFGLoopInfoPass();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestLowerToLLVM();
 void registerTestMatchReductionPass();
 void registerTestMathAlgebraicSimplificationPass();
 void registerTestMathPolynomialApproximationPass();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
#endif
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLastModifiedPass();
   mlir::test::registerTestLinalgDecomposeOps();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgGreedyFusion();
   mlir::test::registerTestLinalgTransforms();
   mlir::test::registerTestLivenessPass();
   mlir::test::registerTestLoopFusion();
+  mlir::test::registerTestCFGLoopInfoPass();
   mlir::test::registerTestLoopMappingPass();
   mlir::test::registerTestLoopUnrollingPass();
   mlir::test::registerTestLowerToLLVM();
   mlir::test::registerTestMatchReductionPass();
   mlir::test::registerTestMathAlgebraicSimplificationPass();
   mlir::test::registerTestMathPolynomialApproximationPass();
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();