[mlir][interfaces] Add helpers for detecting repetitive regions

[springerm]

Summary

Add helper functions to check if an op may be executed multiple times based on RegionBranchOpInterface.

[ftynse]

mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp

60

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

private:
   /// if an upper bound is not known.
   Optional<unsigned> upper;
 };
 
 /// Return `true` if `a` and `b` are in mutually exclusive regions as per
 /// RegionBranchOpInterface.
 bool insideMutuallyExclusiveRegions(Operation *a, Operation *b);
 
+/// Return the first enclosing region of the given op that may be executed
+/// repetitively as per RegionBranchOpInterface or `nullptr` if no such region
+/// exists.
+Region *getEnclosingRepetitiveRegion(Operation *op);
+
+/// Return the first enclosing region of the given Value that may be executed
+/// repetitively as per RegionBranchOpInterface or `nullptr` if no such region
+/// exists.
+Region *getEnclosingRepetitiveRegion(Value value);
+
 //===----------------------------------------------------------------------===//
 // RegionBranchTerminatorOpInterface
 //===----------------------------------------------------------------------===//
 
 /// Returns true if the given operation is either annotated with the
 /// `ReturnLike` trait or implements the `RegionBranchTerminatorOpInterface`.
 bool isRegionReturnLike(Operation *operation);
 

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
 
   let extraClassDeclaration = [{
     /// Convenience helper in case none of the operands is known.
     void getSuccessorRegions(Optional<unsigned> index,
                              SmallVectorImpl<RegionSuccessor> &regions) {
        SmallVector<Attribute, 2> nullAttrs(getOperation()->getNumOperands());
        getSuccessorRegions(index, nullAttrs, regions);
     }
+
+    /// Return `true` if control flow originating from the given region may
+    /// eventually branch back to the same region. (Maybe after passing through
+    /// other regions.)
+    bool isRepetitiveRegion(unsigned index);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchTerminatorOpInterface
 //===----------------------------------------------------------------------===//
 
 def RegionBranchTerminatorOpInterface :

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

return !isRegionReachable(regionA, regionB) &&
            !isRegionReachable(regionB, regionA);
   }
 
   // Could not find a common RegionBranchOpInterface among a's and b's
   // ancestors.
   return false;
 }
 
+bool RegionBranchOpInterface::isRepetitiveRegion(unsigned index) {
+  SmallVector<bool> visited(getOperation()->getNumRegions(), false);
+  visited[index] = true;
+
+  // Retrieve all successors of the region and enqueue them in the worklist.
+  SmallVector<unsigned> worklist;
+  auto enqueueAllSuccessors = [&](unsigned index) {
+    SmallVector<RegionSuccessor> successors;
+    this->getSuccessorRegions(index, successors);
+    for (RegionSuccessor successor : successors)
+      if (!successor.isParent())
+        worklist.push_back(successor.getSuccessor()->getRegionNumber());
+  };
+  enqueueAllSuccessors(index);
+
+  // Process all regions in the worklist via DFS.
+  while (!worklist.empty()) {
+    unsigned nextRegion = worklist.pop_back_val();
+    if (nextRegion == index)
+      return true;
+    if (visited[nextRegion])
+      continue;
+    visited[nextRegion] = true;
+    enqueueAllSuccessors(nextRegion);
+  }
+
+  return false;
+}
+
+Region *mlir::getEnclosingRepetitiveRegion(Operation *op) {
+  while (Region *region = op->getParentRegion()) {
+    op = region->getParentOp();
+    if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op))
+      if (branchOp.isRepetitiveRegion(region->getRegionNumber()))
+        return region;
+  }
+  return nullptr;
+}
+
+Region *mlir::getEnclosingRepetitiveRegion(Value value) {
+  Region *region = value.getParentRegion();
+  while (region) {
+    Operation *op = region->getParentOp();
+    if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op))
+      if (branchOp.isRepetitiveRegion(region->getRegionNumber()))
+        return region;
+    region = op->getParentRegion();
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // RegionBranchTerminatorOpInterface
 //===----------------------------------------------------------------------===//
 
 /// Returns true if the given operation is either annotated with the
 /// `ReturnLike` trait or implements the `RegionBranchTerminatorOpInterface`.
 bool mlir::isRegionReturnLike(Operation *operation) {
   return dyn_cast<RegionBranchTerminatorOpInterface>(operation) ||

mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp

struct MutuallyExclusiveRegionsOp
   }
 
   // Regions have no successors.
   void getSuccessorRegions(Optional<unsigned> index,
                            ArrayRef<Attribute> operands,
                            SmallVectorImpl<RegionSuccessor> &regions) {}
 };
 
+/// All regions of this op call each other in a large circle.
+struct LoopRegionsOp
+    : public Op<LoopRegionsOp, RegionBranchOpInterface::Trait> {
+  using Op::Op;
+  static const unsigned kNumRegions = 3;
+
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
+
+  static StringRef getOperationName() { return "cftest.loop_regions_op"; }
+
+  void getSuccessorRegions(Optional<unsigned> index,
+                           ArrayRef<Attribute> operands,
+                           SmallVectorImpl<RegionSuccessor> &regions) {
+    if (index) {
+      if (*index == 1)
+        // This region also branches back to the parent.
+        regions.push_back(RegionSuccessor());
+      regions.push_back(
+          RegionSuccessor(&getOperation()->getRegion(*index % kNumRegions)));
+    }
+  }
+};
+
 /// Regions are executed sequentially.
 struct SequentialRegionsOp
     : public Op<SequentialRegionsOp, RegionBranchOpInterface::Trait> {
   using Op::Op;
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "cftest.sequential_regions_op"; }
 
   // Region 0 has Region 1 as a successor.
   void getSuccessorRegions(Optional<unsigned> index,
                            ArrayRef<Attribute> operands,
                            SmallVectorImpl<RegionSuccessor> &regions) {
     if (index == 0u) {
       Operation *thisOp = this->getOperation();
       regions.push_back(RegionSuccessor(&thisOp->getRegion(1)));
     }
   }
 };
 
 /// A dialect putting all the above together.
 struct CFTestDialect : Dialect {
   explicit CFTestDialect(MLIRContext *ctx)
       : Dialect(getDialectNamespace(), ctx, TypeID::get<CFTestDialect>()) {
-    addOperations<DummyOp, MutuallyExclusiveRegionsOp, SequentialRegionsOp>();
+    addOperations<DummyOp, MutuallyExclusiveRegionsOp, LoopRegionsOp,
+                  SequentialRegionsOp>();
   }
   static StringRef getDialectNamespace() { return "cftest"; }
 };
 
 TEST(RegionBranchOpInterface, MutuallyExclusiveOps) {
   const char *ir = R"MLIR(
 "cftest.mutually_exclusive_regions_op"() (
       {"cftest.dummy_op"() : () -> ()},  // op1
"cftest.mutually_exclusive_regions_op"() (
   Operation *op2 = &testOp->getRegion(1).front().front();
   Operation *op3 =
       &testOp->getRegion(0).front().front().getRegion(1).front().front();
 
   EXPECT_TRUE(insideMutuallyExclusiveRegions(op1, op2));
   EXPECT_TRUE(insideMutuallyExclusiveRegions(op3, op2));
   EXPECT_FALSE(insideMutuallyExclusiveRegions(op1, op3));
 }
+
+TEST(RegionBranchOpInterface, RecursiveRegions) {
+  const char *ir = R"MLIR(
+"cftest.loop_regions_op"() (
+      {"cftest.dummy_op"() : () -> ()},  // op1
+      {"cftest.dummy_op"() : () -> ()},  // op2
+      {"cftest.dummy_op"() : () -> ()}   // op3
+  ) : () -> ()
+  )MLIR";
+
+  DialectRegistry registry;
+  registry.insert<CFTestDialect>();
+  MLIRContext ctx(registry);
+
+  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
+  Operation *testOp = &module->getBody()->getOperations().front();
+  auto regionOp = cast<RegionBranchOpInterface>(testOp);
+  Operation *op1 = &testOp->getRegion(0).front().front();
+  Operation *op2 = &testOp->getRegion(1).front().front();
+  Operation *op3 = &testOp->getRegion(2).front().front();
+
+  EXPECT_TRUE(regionOp.isRepetitiveRegion(0));
+  EXPECT_TRUE(regionOp.isRepetitiveRegion(1));
+  EXPECT_TRUE(regionOp.isRepetitiveRegion(2));
+  EXPECT_NE(getEnclosingRepetitiveRegion(op1), nullptr);
+  EXPECT_NE(getEnclosingRepetitiveRegion(op2), nullptr);
+  EXPECT_NE(getEnclosingRepetitiveRegion(op3), nullptr);
+}
+
+TEST(RegionBranchOpInterface, NotRecursiveRegions) {
+  const char *ir = R"MLIR(
+"cftest.sequential_regions_op"() (
+      {"cftest.dummy_op"() : () -> ()},  // op1
+      {"cftest.dummy_op"() : () -> ()}   // op2
+  ) : () -> ()
+  )MLIR";
+
+  DialectRegistry registry;
+  registry.insert<CFTestDialect>();
+  MLIRContext ctx(registry);
+
+  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
+  Operation *testOp = &module->getBody()->getOperations().front();
+  Operation *op1 = &testOp->getRegion(0).front().front();
+  Operation *op2 = &testOp->getRegion(1).front().front();
+
+  EXPECT_EQ(getEnclosingRepetitiveRegion(op1), nullptr);
+  EXPECT_EQ(getEnclosingRepetitiveRegion(op2), nullptr);
+}