[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

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

[ftynse]

mlir/include/mlir/Analysis/LoopInfo.h

1

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

30

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

65

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

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

mlir/test/lib/Analysis/TestLoopInfo.cpp

45

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/LoopInfo.h

+//===- LoopInfo.h - LoopInfo analysis for regions ---------------*- 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 LoopInfo analysis for MLIR. The LoopInfo 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 llvm {
+
+// llvm::GraphTrait specializations that are required for the LLVM's generic
+// LoopInfo. Note that the const_casts are required because MLIR has no constant
+// accessors on IR constructs.
+
+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();
+  }
+};
+} // namespace llvm
+
+namespace mlir {
+class Loop : public llvm::LoopBase<mlir::Block, mlir::Loop> {
+private:
+  explicit Loop(mlir::Block *block);
+
+  friend class llvm::LoopBase<mlir::Block, Loop>;
+  friend class llvm::LoopInfoBase<mlir::Block, Loop>;
+};
+
+/// Instantiate a variant of LLVM LoopInfo that works on mlir::Block.
+class LoopInfo : public llvm::LoopInfoBase<mlir::Block, mlir::Loop> {
+public:
+  LoopInfo(const llvm::DominatorTreeBase<mlir::Block, false> &domTree);
+};
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_LOOPINFO_H

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   CallGraph.cpp
   DataLayoutAnalysis.cpp
   FlatLinearValueConstraints.cpp
   Liveness.cpp
+  LoopInfo.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
+  LoopInfo.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   DataFlow/ConstantPropagationAnalysis.cpp
   DataFlow/DeadCodeAnalysis.cpp
   DataFlow/DenseAnalysis.cpp
   DataFlow/IntegerRangeAnalysis.cpp

mlir/lib/Analysis/LoopInfo.cpp

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

mlir/test/Analysis/test-loopinfo.mlir

+// RUN: mlir-opt -pass-pipeline="builtin.module(any(test-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
+  TestLoopInfo.cpp
   TestMatchReduction.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestSlice.cpp
 
   DataFlow/TestDeadCodeAnalysis.cpp
   DataFlow/TestDenseDataFlowAnalysis.cpp

mlir/test/lib/Analysis/TestLoopInfo.cpp

+//===- TestLoopInfo.cpp - Test 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 LoopInfo analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/LoopInfo.h"
+#include "mlir/IR/FunctionInterfaces.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// A testing pass that applies the LoopInfo analysis on a region and prints the
+/// information it collected to llvm::errs().
+struct TestLoopInfo
+    : public PassWrapper<TestLoopInfo, InterfacePass<FunctionOpInterface>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopInfo)
+
+  StringRef getArgument() const final { return "test-loop-info"; }
+  StringRef getDescription() const final {
+    return "Test the loop info analysis.";
+  }
+
+  void runOnOperation() override;
+};
+} // namespace
+
+void TestLoopInfo::runOnOperation() {
+  auto func = getOperation();
+  DominanceInfo &domInfo = getAnalysis<DominanceInfo>();
+  Region &region = func.getFunctionBody();
+
+  // Prints the label of the test.
+  llvm::errs() << "Testing : " << func.getNameAttr() << "\n";
+
+  // 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::LoopInfo loopInfo(domTree);
+
+  if (loopInfo.getTopLevelLoops().empty())
+    llvm::errs() << "no loops\n";
+  else
+    loopInfo.print(llvm::errs());
+}
+
+namespace mlir {
+namespace test {
+void registerTestLoopInfoPass() { PassRegistration<TestLoopInfo>(); }
+} // 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 registerTestLoopInfoPass();
 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::registerTestLoopInfoPass();
   mlir::test::registerTestLoopMappingPass();
   mlir::test::registerTestLoopUnrollingPass();
   mlir::test::registerTestLowerToLLVM();
   mlir::test::registerTestMatchReductionPass();
   mlir::test::registerTestMathAlgebraicSimplificationPass();
   mlir::test::registerTestMathPolynomialApproximationPass();
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();