From 52bc6b9770bb85cba00a6f075e58dc0694ca2881 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <git@ozinenko.com>
Date: Sat, 14 Dec 2024 15:02:18 +0100
Subject: [PATCH] [mlir] add noinline attribute to func.func/call

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.
---
 mlir/include/mlir/Dialect/Func/IR/FuncOps.td  |  6 +++--
 .../Func/Extensions/InlinerExtension.cpp      | 10 +++++---
 mlir/test/Transforms/inlining.mlir            | 23 +++++++++++++++++++
 3 files changed, 34 insertions(+), 5 deletions(-)

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<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
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<func::CallOp>(call);
+    auto funcOp = dyn_cast<func::FuncOp>(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