diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td index 22efe15aa83a5..237a825c19104 100644 --- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td +++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td @@ -49,7 +49,8 @@ def CallOp : Func_Op<"call", ``` }]; - let arguments = (ins FlatSymbolRefAttr:$callee, Variadic:$operands); + let arguments = (ins FlatSymbolRefAttr:$callee, Variadic:$operands, + UnitAttr:$no_inline); let results = (outs Variadic); let builders = [ @@ -270,7 +271,8 @@ def FuncOp : Func_Op<"func", [ TypeAttrOf:$function_type, OptionalAttr:$sym_visibility, OptionalAttr:$arg_attrs, - OptionalAttr:$res_attrs); + OptionalAttr:$res_attrs, + UnitAttr:$no_inline); let regions = (region AnyRegion:$body); let builders = [OpBuilder<(ins diff --git a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp index 719a74a29a622..3328d58551bff 100644 --- a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp +++ b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp @@ -27,10 +27,14 @@ struct FuncInlinerInterface : public DialectInlinerInterface { // Analysis Hooks //===--------------------------------------------------------------------===// - /// All call operations can be inlined. + /// Call operations can be inlined unless specified otherwise by attributes + /// on either the call or the callbale. bool isLegalToInline(Operation *call, Operation *callable, bool wouldBeCloned) const final { - return true; + auto callOp = dyn_cast(call); + auto funcOp = dyn_cast(callable); + return !(callOp && callOp.getNoInline()) && + !(funcOp && funcOp.getNoInline()); } /// All operations can be inlined. @@ -38,7 +42,7 @@ struct FuncInlinerInterface : public DialectInlinerInterface { return true; } - /// All functions can be inlined. + /// All function bodies can be inlined. bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final { return true; } diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir index 79a2936b104fa..65ffaad1fa859 100644 --- a/mlir/test/Transforms/inlining.mlir +++ b/mlir/test/Transforms/inlining.mlir @@ -19,6 +19,29 @@ func.func @inline_with_arg(%arg0 : i32) -> i32 { return %0 : i32 } +// CHECK-LABEL: func @noinline_with_arg +func.func @noinline_with_arg(%arg0 : i32) -> i32 { + // CHECK-NEXT: func_with_arg + // CHECK-NEXT: return + + %0 = call @func_with_arg(%arg0) {no_inline} : (i32) -> i32 + return %0 : i32 +} + +func.func @non_inlinable_func_with_arg(%c : i32) -> i32 attributes {no_inline} { + %b = arith.addi %c, %c : i32 + return %b : i32 +} + +// CHECK-LABEL: func @noinline_with_func_arg +func.func @noinline_with_func_arg(%arg0 : i32) -> i32 { + // CHECK-NEXT: non_inlinable_func_with_arg + // CHECK-NEXT: return + + %0 = call @non_inlinable_func_with_arg(%arg0) : (i32) -> i32 + return %0 : i32 +} + // Inline a function that has multiple return operations. func.func @func_with_multi_return(%a : i1) -> (i32) { cf.cond_br %a, ^bb1, ^bb2