[mlir][CFGToSCF] Visit subregions in CFGToSCF pass

This is useful when user already have partially-scf'ed IR or other ops with nested regions (e.g. linalg.generic).

Also, improve error message and pass docs.

Differential Revision: https://reviews.llvm.org/D158349

Author: Hardcode84

mlir/include/mlir/Conversion/Passes.td

@@ -304,6 +304,12 @@ def LiftControlFlowToSCFPass : Pass<"lift-cf-to-scf"> {
     ControlFlow operations will be replaced successfully.
     Otherwise a single ControlFlow switch branching to one block per return-like
     operation kind remains.
+
+    This pass may need to create unreachable terminators in case of infinite
+    loops, which is only supported for 'func.func' for now. If you potentially
+    have infinite loops inside CFG regions not belonging to 'func.func',
+    consider using `transformCFGToSCF` function directly with corresponding
+    `CFGToSCFInterface::createUnreachableTerminator` implementation.
   }];
 
   let dependentDialects = ["scf::SCFDialect",

mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp

@@ -140,10 +140,11 @@ ControlFlowToSCFTransformation::createUnreachableTerminator(Location loc,
   // TODO: This should create a `ub.unreachable` op. Once such an operation
   //       exists to make the pass independent of the func dialect. For now just
   //       return poison values.
-  auto funcOp = dyn_cast<func::FuncOp>(region.getParentOp());
+  Operation *parentOp = region.getParentOp();
+  auto funcOp = dyn_cast<func::FuncOp>(parentOp);
   if (!funcOp)
-    return emitError(loc, "Expected '")
-           << func::FuncOp::getOperationName() << "' as top level operation";
+    return emitError(loc, "Cannot create unreachable terminator for '")
+           << parentOp->getName() << "'";
 
   return builder
       .create<func::ReturnOp>(
@@ -165,18 +166,29 @@ struct LiftControlFlowToSCF
     ControlFlowToSCFTransformation transformation;
 
     bool changed = false;
-    WalkResult result = getOperation()->walk([&](func::FuncOp funcOp) {
+    Operation *op = getOperation();
+    WalkResult result = op->walk([&](func::FuncOp funcOp) {
       if (funcOp.getBody().empty())
         return WalkResult::advance();
 
-      FailureOr<bool> changedFunc = transformCFGToSCF(
-          funcOp.getBody(), transformation,
-          funcOp != getOperation() ? getChildAnalysis<DominanceInfo>(funcOp)
-                                   : getAnalysis<DominanceInfo>());
-      if (failed(changedFunc))
+      auto &domInfo = funcOp != op ? getChildAnalysis<DominanceInfo>(funcOp)
+                                   : getAnalysis<DominanceInfo>();
+
+      auto visitor = [&](Operation *innerOp) -> WalkResult {
+        for (Region &reg : innerOp->getRegions()) {
+          FailureOr<bool> changedFunc =
+              transformCFGToSCF(reg, transformation, domInfo);
+          if (failed(changedFunc))
+            return WalkResult::interrupt();
+
+          changed |= *changedFunc;
+        }
+        return WalkResult::advance();
+      };
+
+      if (funcOp->walk<WalkOrder::PostOrder>(visitor).wasInterrupted())
         return WalkResult::interrupt();
 
-      changed |= *changedFunc;
       return WalkResult::advance();
     });
     if (result.wasInterrupted())

mlir/test/Conversion/ControlFlowToSCF/test.mlir

@@ -678,3 +678,35 @@ func.func @multi_entry_loop(%cond: i1) {
 // CHECK:        scf.yield
 // CHECK:      call @foo(%[[WHILE]]#1)
 // CHECK-NEXT: return
+
+// -----
+
+func.func @nested_region() {
+  scf.execute_region {
+    %cond = "test.test1"() : () -> i1
+    cf.cond_br %cond, ^bb1, ^bb2
+  ^bb1:
+    "test.test2"() : () -> ()
+    cf.br ^bb3
+  ^bb2:
+    "test.test3"() : () -> ()
+    cf.br ^bb3
+  ^bb3:
+    "test.test4"() : () -> ()
+    scf.yield
+  }
+  return
+}
+
+// CHECK-LABEL: func @nested_region
+// CHECK:      scf.execute_region {
+// CHECK:      %[[COND:.*]] = "test.test1"()
+// CHECK-NEXT: scf.if %[[COND]]
+// CHECK-NEXT:   "test.test2"()
+// CHECK-NEXT: else
+// CHECK-NEXT:   "test.test3"()
+// CHECK-NEXT: }
+// CHECK-NEXT: "test.test4"()
+// CHECK-NEXT: scf.yield
+// CHECK-NEXT: }
+// CHECK-NEXT: return