diff --git a/mlir/lib/Transforms/Inliner.cpp b/mlir/lib/Transforms/Inliner.cpp --- a/mlir/lib/Transforms/Inliner.cpp +++ b/mlir/lib/Transforms/Inliner.cpp @@ -451,8 +451,23 @@ // 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 = llvm::hasNItemsOrMore(*targetRegion, /*N=*/2); + // If both parent ops have the same type, it is safe to inline. Otherwise, + // decide based on whether the op has the SingleBlock trait or not. + // Note: This check does currently not account for SizedRegion/MaxSizedRegion. + auto targetSupportsMultipleBlocks = [&]() { + return targetRegion->getParentOp()->getName() == + resolvedCall.call->getParentOp()->getName() || + !resolvedCall.call->getParentOp() + ->mightHaveTrait(); + }; + if (sourceHasMultipleBlocks && !targetSupportsMultipleBlocks()) return false; // Otherwise, inline. diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir --- a/mlir/test/Transforms/inlining.mlir +++ b/mlir/test/Transforms/inlining.mlir @@ -227,6 +227,18 @@ 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) -> () + "test.terminator"() : () -> () + }) : () -> () + 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