[MLIR] Add no-inline attribute

[wsmoses]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir-func

@llvm/pr-subscribers-mlir

Author: William Moses (wsmoses)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/117392.diff

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/Func/IR/FuncOps.td (+4)
• (modified) mlir/include/mlir/Interfaces/CallInterfaces.td (+10)
• (modified) mlir/lib/Transforms/Utils/Inliner.cpp (+4)
• (modified) mlir/test/Transforms/inlining.mlir (+9)

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

[River707]

I don't think this should be a hook on CallInterface. The Inliner interface already allows for expressing various types of legality that could reference attributes: e.g.

llvm-project/mlir/include/mlir/Transforms/InliningUtils.h

Line 59 in 24ced77

virtual bool isLegalToInline(Operation *call, Operation *callable,

[wsmoses]

@River707 I moved the interface, per suggestion (pardon the delay)

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff 76f0ff8210d56a050d2679926a2fdddd3a8c16d6 691eb770a65e4ad1e66d6f711d3fe92f360e4317 --extensions cpp -- mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp

View the diff from clang-format here.

diff --git a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
index 59e108c5b6..3e29ef5fd6 100644
--- a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
@@ -23,7 +23,7 @@ namespace {
 struct FuncInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
-  static constexpr const char* NoInlineAttrName = "noinline";
+  static constexpr const char *NoInlineAttrName = "noinline";
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
@@ -32,7 +32,8 @@ struct FuncInlinerInterface : public DialectInlinerInterface {
   /// All call operations can be inlined.
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
-    return !call->hasAttr(NoInlineAttrName) && !callable->hasAttr(NoInlineAttrName);
+    return !call->hasAttr(NoInlineAttrName) &&
+           !callable->hasAttr(NoInlineAttrName);
   }
 
   /// All operations can be inlined.

[ftynse]

Please fix clang-format.

[ftynse]

Hmm, would it make sense to have this as a proper ODS attribute on FuncOp and CallOp? LLVM dialect already has one

llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

Line 1569 in cc54a0c

OptionalAttr<UnitAttr>:$no_inline,

.

[River707]

Yeah, it'd be good to integrate this directly into the dialect if possible (makes it much cleaner to target). Otherwise you're at the whims of attribute propagation.