[mlir] add noinline attribute to func.func/call (#119970)

This allows for inlining to be somewhat controlled by the user instead
of always inlining everything. External heuristics may be used to place
`no_inline` attributes on invidiual calls or functions to prevent
inlining.

Author: ftynse

mlir/include/mlir/Dialect/Func/IR/FuncOps.td

@@ -49,7 +49,8 @@ def CallOp : Func_Op<"call",
     ```
   }];
 
-  let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands);
+  let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands,
+                       UnitAttr:$no_inline);
   let results = (outs Variadic<AnyType>);
 
   let builders = [
@@ -270,7 +271,8 @@ def FuncOp : Func_Op<"func", [
                        TypeAttrOf<FunctionType>:$function_type,
                        OptionalAttr<StrAttr>:$sym_visibility,
                        OptionalAttr<DictArrayAttr>:$arg_attrs,
-                       OptionalAttr<DictArrayAttr>:$res_attrs);
+                       OptionalAttr<DictArrayAttr>:$res_attrs,
+                       UnitAttr:$no_inline);
   let regions = (region AnyRegion:$body);
 
   let builders = [OpBuilder<(ins

mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp

@@ -27,18 +27,22 @@ 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<func::CallOp>(call);
+    auto funcOp = dyn_cast<func::FuncOp>(callable);
+    return !(callOp && callOp.getNoInline()) &&
+           !(funcOp && funcOp.getNoInline());
   }
 
   /// All operations can be inlined.
   bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
-  /// All functions can be inlined.
+  /// All function bodies can be inlined.
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }

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