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 @@ -16,6 +16,7 @@ #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" @@ -451,8 +452,22 @@ // 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(resolvedCall.call->getParentOp()); + 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,17 @@ 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