[mlir][CFGToSCF] Visit subregions in CFGToSCF pass

[Hardcode84]

Summary

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

[zero9178]

There is one concern that I have with this change, which is that if the subregion were to contain an infinite loop, forcing the creation of an unreachable terminator, the interface implementation would currently create a func.return with a poison value which IIRC is only allowed to be nested within a func.func.
Since the interface implementation cannot know (I think no such interface exists at the moment) what kind of terminator is required, it'd lead to invalid IR.
Example IR:

func.func private  @foo()

func.func @nested_region() {
  scf.execute_region {
    cf.br ^bb0
  ^bb0:
    func.call @foo() : () -> ()
    cf.br ^bb0
  }
  return
}

A solution to this would ideally be a ub.unreachable op. An interface or trait returning or helping build an appropiate terminator for a region op would also help but a lot more intrusive I imagine.

[Hardcode84]

In D158349#4601725, @zero9178 wrote:

There is one concern that I have with this change, which is that if the subregion were to contain an infinite loop, forcing the creation of an unreachable terminator, the interface implementation would currently create a func.return with a poison value which IIRC is only allowed to be nested within a func.func.
Since the interface implementation cannot know (I think no such interface exists at the moment) what kind of terminator is required, it'd lead to invalid IR.
Example IR:

func.func private  @foo()

func.func @nested_region() {
  scf.execute_region {
    cf.br ^bb0
  ^bb0:
    func.call @foo() : () -> ()
    cf.br ^bb0
  }
  return
}

A solution to this would ideally be a ub.unreachable op. An interface or trait returning or helping build an appropiate terminator for a region op would also help but a lot more intrusive I imagine.

I've checked it and pass actually fails gracefully in this case, although error message is not very informative

Expected 'func.func' as top level operation

[Hardcode84]

So, I think if pass fails gracefully and we change error message to smth like Cannot create unreachable terminator for 'opname' it should be fine.

[zero9178]

In D158349#4601817, @Hardcode84 wrote:

So, I think if pass fails gracefully and we change error message to smth like Cannot create unreachable terminator for 'opname' it should be fine.

Yeah I think that the implementation of a ub.unreachable shouldn't be a block for this either and creating a nicer error message would be great.
I do think that this should be documented however and if possible, that overriding the interface and using the transformCFGToSCF function manually ideally mentioned as an alternative.
One of the reason the interface and separate function exist is because I figured it was impossible to fit every downstream scenario with just an upstream pass, so creating the tools for people to make their own passes was important.

[Hardcode84]

improve error message and pass docs

[zero9178]

LGTM

mlir/include/mlir/Conversion/Passes.td

310

mlir/include/mlir/Conversion/Passes.td

let description = [{
     Lifts ControlFlow operations to SCF dialect operations.
 
     This pass is prefixed with "lift" instead of "convert" as it is not always
     guaranteed to replace all ControlFlow ops.
     If a region contains only a single kind of return-like operation, all
     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",
                            "arith::ArithDialect",
                            "ub::UBDialect",
                            // TODO: This is only necessary until we have a
                            //       ub.unreachable op.
                            "func::FuncDialect"];

mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp

 FailureOr<Operation *>
 ControlFlowToSCFTransformation::createUnreachableTerminator(Location loc,
                                                             OpBuilder &builder,
                                                             Region &region) {
 
   // 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>(
           loc, llvm::map_to_vector(funcOp.getResultTypes(),
                                    [&](Type type) {
                                      return getUndefValue(loc, builder, type);
                                    }))
       .getOperation();
 }
 
 namespace {
 
 struct LiftControlFlowToSCF
     : public impl::LiftControlFlowToSCFPassBase<LiftControlFlowToSCF> {
 
   using Base::Base;
 
   void runOnOperation() override {
     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>());
+      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();
+
       return WalkResult::advance();
     });
     if (result.wasInterrupted())
       return signalPassFailure();
 
     if (!changed)
       markAllAnalysesPreserved();
   }
 };
 } // namespace

mlir/test/Conversion/ControlFlowToSCF/test.mlir

 // CHECK-NEXT:       scf.yield %[[C0]], %[[UB]], %[[C0]], %[[C1]]
 // CHECK:          scf.yield %[[IF_RES]]#0, %[[IF_RES]]#1, %[[IF_RES]]#2, %[[IF_RES]]#3
 // CHECK:        %[[COND:.*]] = arith.trunci %[[SWITCH]]#3
 // CHECK-NEXT:   scf.condition(%[[COND]]) %[[SWITCH]]#0, %[[SWITCH]]#1
 // CHECK:      do
 // 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