[Mlir-commits] [mlir] [MLIR] Add no-inline attribute (PR #117392)

[William Moses]

https://github.com/wsmoses created https://github.com/llvm/llvm-project/pull/117392

None

>From 240672bc86ddbfb82981a16e1106271c5fbcab22 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Fri, 22 Nov 2024 17:42:12 -0500
Subject: [PATCH] Add no-inline attribute

---
 mlir/include/mlir/Dialect/Func/IR/FuncOps.td   |  4 ++++
 mlir/include/mlir/Interfaces/CallInterfaces.td | 10 ++++++++++
 mlir/lib/Transforms/Utils/Inliner.cpp          |  4 ++++
 mlir/test/Transforms/inlining.mlir             |  9 +++++++++
 4 files changed, 27 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
index 22efe15aa83a50..e0560c8e1e038a 100644
--- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
+++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
@@ -98,6 +98,10 @@ def CallOp : Func_Op<"call",
     void setCalleeFromCallable(CallInterfaceCallable callee) {
       (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
     }
+
+    bool legalToInline() {
+      return !(*this)->hasAttr("noinline");
+    }
   }];
 
   let assemblyFormat = [{
diff --git a/mlir/include/mlir/Interfaces/CallInterfaces.td b/mlir/include/mlir/Interfaces/CallInterfaces.td
index c6002da0d491ce..35f33bf4797f6a 100644
--- a/mlir/include/mlir/Interfaces/CallInterfaces.td
+++ b/mlir/include/mlir/Interfaces/CallInterfaces.td
@@ -80,6 +80,16 @@ def CallOpInterface : OpInterface<"CallOpInterface"> {
       /*methodBody=*/[{}], /*defaultImplementation=*/[{
         return ::mlir::call_interface_impl::resolveCallable($_op);
       }]
+    >,
+    InterfaceMethod<[{
+        Return whether a given call is legal to inline or not (assuming all other requirements are met,
+        like the callee is direct, the parent region supports multiple blocks and the resolved, etc).
+        This can be used to implement a `noinline`-like attribute.
+    }],
+      "bool", "legalToInline", (ins),
+      /*methodBody=*/[{}], /*defaultImplementation=*/[{
+        return true;
+      }]
     >
   ];
 }
diff --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 8acfc96d2b611b..1e329935b7807b 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -712,6 +712,10 @@ bool Inliner::Impl::shouldInline(ResolvedCall &resolvedCall) {
   if (resolvedCall.call->hasTrait<OpTrait::IsTerminator>())
     return false;
 
+  // Don't inline calls which explicitly forbit inlining.
+  if (!resolvedCall.call.legalToInline())
+    return false;
+
   // Don't allow inlining if the target is a self-recursive function.
   if (llvm::count_if(*resolvedCall.targetNode,
                      [&](CallGraphNode::Edge const &edge) -> bool {
diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 79a2936b104fa1..262047132064c7 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -19,6 +19,15 @@ 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) {"noinline"=""} : (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