[mlir][Transforms] Inliner: Extra checks for unstructured control flow

[springerm]

Summary

Do not inline IR with multiple blocks into ops that may not support unstructured control flow.

This fixes #64978.

Depends On: D159078

[mehdi_amini]

I guess you're lookin into the SingleBlockRegion trait still?

mlir/lib/Transforms/Inliner.cpp

461

.size() is O(N), you can probably try &targetRegion->front() != &targetRegion->back() ?
(I don't know if we have a guarantee that targetRegion is not empty?)

471

I'd guarantee lazy evaluation here:

auto targetSupportsMultipleBlocks = [&]() {
    targetRegion->getParentOp()->getName() ==
        resolvedCall.call->getParentOp()->getName() ||
    isa<FunctionOpInterface>(resolvedCall.call->getParentOp());
};
if (sourceHasMultipleBlocks && !targetSupportsMultipleBlocks())

[springerm]

In D159072#4624270, @mehdi_amini wrote:

I guess you're lookin into the SingleBlockRegion trait still?

Yes, I'm waiting for comments on the Discourse post, then will update this revision.

[springerm]

address comments

[springerm]

better variable names

mlir/lib/Transforms/Inliner.cpp

 // incremental propagation of inlining decisions from the leafs to the roots of
 // the callgraph.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/Passes.h"
 
 #include "mlir/Analysis/CallGraph.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Threading.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/DebugStringHelper.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/Support/Debug.h"

 static bool shouldInline(ResolvedCall &resolvedCall) {
   // Don't allow inlining terminator calls. We currently don't support this
   // case.
   if (resolvedCall.call->hasTrait<OpTrait::IsTerminator>())
     return false;
 
   // Don't allow inlining if the target is an ancestor of the call. This
   // prevents inlining recursively.
-  if (resolvedCall.targetNode->getCallableRegion()->isAncestor(
-          resolvedCall.call->getParentRegion()))
+  Region *targetRegion = resolvedCall.targetNode->getCallableRegion();
+  if (targetRegion->isAncestor(resolvedCall.call->getParentRegion()))
+    return false;
+
+  // Don't allow inlining if the source has multiple blocks (unstructured
+  // control flow) but we cannot be sure that the target supports that.
+  bool sourceHasMultipleBlocks = targetRegion->getBlocks().size() > 1;
+  // Functions typically support multiple blocks. Furthermore, if both parent
+  // ops have the same type, it is safe to inline.
+  // TODO: This does not account for ops like scf.execute_region, into which
+  // multiple blocks could be inlined.
+  bool targetSupportsMultipleBlocks =
+      targetRegion->getParentOp()->getName() ==
+          resolvedCall.call->getParentOp()->getName() ||
+      isa<FunctionOpInterface>(resolvedCall.call->getParentOp());
+  if (sourceHasMultipleBlocks && !targetSupportsMultipleBlocks)
     return false;
 
   // Otherwise, inline.
   return true;
 }
 
 /// Attempt to inline calls within the given scc. This function returns
 /// success if any calls were inlined, failure otherwise.
 static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,

mlir/test/Transforms/inlining.mlir

 }
 
 // CHECK-LABEL: func @func_with_block_args_location_callee2
 func.func @func_with_block_args_location_callee2(%arg0 : i32) {
   call @func_with_block_args_location(%arg0) : (i32) -> ()
   return
 }
 
+// CHECK-LABEL: func @func_with_block_args_location_callee3
+func.func @func_with_block_args_location_callee3(%arg0 : i32) {
+  "test.dummy_op"() ({
+    // Call cannot be inlined because "test.dummy" may not support unstructured
+    // control flow in its body.
+    // CHECK: call @func_with_block_args_location
+    call @func_with_block_args_location(%arg0) : (i32) -> ()
+  }) : () -> ()
+  return
+}
+
 // Check that we can handle argument and result attributes.
 test.conversion_func_op @handle_attr_callee_fn_multi_arg(%arg0 : i16, %arg1 : i16 {"test.handle_argument"}) -> (i16 {"test.handle_result"}, i16) {
   %0 = arith.addi %arg0, %arg1 : i16
   %1 = arith.subi %arg0, %arg1 : i16
   "test.return"(%0, %1) : (i16, i16) -> ()
 }
 test.conversion_func_op @handle_attr_callee_fn(%arg0 : i32 {"test.handle_argument"}) -> (i32 {"test.handle_result"}) {
   "test.return"(%arg0) : (i32) -> ()