[MLIR][LLVM] Only disallow inlining for selected function attributes.

[definelicht]

Summary

This loosens the requirement of no passthrough function attribute being
present to checking for specific attributes that prevent inlining. Since
these attributes are no longer strictly passthrough, they should
eventually be upgraded to some form of addressable attributes.

Drops the expensive StringSwitches over call and function attributes in
favor of selectively disallowing attributes that prevent inlining
(similiar to the LLVM inliner), thereby moving to a less conservative
approach.

[gysit]

LGTM!

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

2868–2891

nit: ; -> .?

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

struct LLVMInlinerInterface : public DialectInlinerInterface {
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
     if (!wouldBeCloned)
       return false;
     auto callOp = dyn_cast<LLVM::CallOp>(call);
     auto funcOp = dyn_cast<LLVM::LLVMFuncOp>(callable);
     if (!callOp || !funcOp)
       return false;
-    return isLegalToInlineCallAttributes(callOp) &&
-           isLegalToInlineFuncAttributes(funcOp);
+    // TODO: Handle argument and result attributes;
+    if (funcOp.getArgAttrs() || funcOp.getResAttrs())
+      return false;
+    // TODO: Handle exceptions.
+    if (funcOp.getPersonality())
+      return false;
+    if (funcOp.getPassthrough()) {
+      // TODO: Used attributes should not be passthrough.
+      DenseSet<StringAttr> disallowed(
+          {StringAttr::get(funcOp->getContext(), "noduplicate"),
+           StringAttr::get(funcOp->getContext(), "noinline"),
+           StringAttr::get(funcOp->getContext(), "optnone"),
+           StringAttr::get(funcOp->getContext(), "presplitcoroutine"),
+           StringAttr::get(funcOp->getContext(), "returns_twice"),
+           StringAttr::get(funcOp->getContext(), "strictfp")});
+      if (llvm::any_of(*funcOp.getPassthrough(), [&](Attribute attr) {
+            auto stringAttr = dyn_cast<StringAttr>(attr);
+            if (!stringAttr)
+              return false;
+            return disallowed.contains(stringAttr);
+          }))
+        return false;
+    }
+    return true;
   }
 
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
-  /// Conservative allowlist-based inlining of operations supported so far.
+  /// Conservative allowlist of operations supported so far.
   bool isLegalToInline(Operation *op, Region *, bool, IRMapping &) const final {
     if (isPure(op))
       return true;
     return llvm::TypeSwitch<Operation *, bool>(op)
         .Case<LLVM::LoadOp, LLVM::StoreOp>([&](auto memOp) {
           // Some attributes on load and store operations require handling
           // during inlining. Since this is not yet implemented, refuse to
           // inline memory operations that have any of these attributes.
void processInlinedCallBlocks(
       iterator_range<Region::iterator> inlinedBlocks) const override {
     // Alloca operations with a constant size that were in the entry block of
     // the callee should be moved to the entry block of the caller, as this will
     // fold into prologue/epilogue code during code generation.
     // This is not implemented as a standalone pattern because we need to know
     // which newly inlined block was previously the entry block of the callee.
     moveConstantAllocasToEntryBlock(inlinedBlocks);
   }
-
-private:
-  /// Returns true if all attributes of `callOp` are handled during inlining.
-  [[nodiscard]] static bool isLegalToInlineCallAttributes(LLVM::CallOp callOp) {
-    return all_of(callOp.getAttributeNames(), [&](StringRef attrName) {
-      return llvm::StringSwitch<bool>(attrName)
-          // TODO: Propagate and update branch weights.
-          .Case("branch_weights", !callOp.getBranchWeights())
-          .Case("callee", true)
-          .Case("fastmathFlags", true)
-          .Default(false);
-    });
-  }
-
-  /// Returns true if all attributes of `funcOp` are handled during inlining.
-  [[nodiscard]] static bool
-  isLegalToInlineFuncAttributes(LLVM::LLVMFuncOp funcOp) {
-    return all_of(funcOp.getAttributeNames(), [&](StringRef attrName) {
-      return llvm::StringSwitch<bool>(attrName)
-          .Case("CConv", true)
-          .Case("arg_attrs", ([&]() {
-                  if (!funcOp.getArgAttrs())
-                    return true;
-                  return llvm::all_of(funcOp.getArgAttrs().value(),
-                                      [&](Attribute) {
-                                        // TODO: Handle argument attributes.
-                                        return false;
-                                      });
-                })())
-          .Case("dso_local", true)
-          .Case("function_entry_count", true)
-          .Case("function_type", true)
-          // TODO: Once the garbage collector attribute is supported on
-          // LLVM::CallOp, make sure that the garbage collector matches.
-          .Case("garbageCollector", !funcOp.getGarbageCollector())
-          .Case("linkage", true)
-          .Case("memory", true)
-          .Case("passthrough", !funcOp.getPassthrough())
-          // Exception handling is not yet supported, so bail out if the
-          // personality is set.
-          .Case("personality", !funcOp.getPersonality())
-          // TODO: Handle result attributes.
-          .Case("res_attrs", !funcOp.getResAttrs())
-          .Case("sym_name", true)
-          .Default(false);
-    });
-  }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 void LLVMDialect::initialize() {

mlir/test/Dialect/LLVMIR/inlining.mlir

llvm.func internal fastcc @callee() -> (i32) attributes { function_entry_count = 42 : i64, dso_local } {
   llvm.return %0 : i32
 }
 
 // CHECK-LABEL: llvm.func @caller
 // CHECK-NEXT: %[[CST:.+]] = llvm.mlir.constant
 // CHECK-NEXT: llvm.return %[[CST]]
 llvm.func @caller() -> (i32) {
   // Include all call attributes that don't prevent inlining.
-  %0 = llvm.call @callee() { fastmathFlags = #llvm.fastmath<nnan, ninf> } : () -> (i32)
+  %0 = llvm.call @callee() { fastmathFlags = #llvm.fastmath<nnan, ninf>, branch_weights = dense<42> : vector<1xi32> } : () -> (i32)
   llvm.return %0 : i32
 }
 
 // -----
 
 llvm.func @foo() -> (i32) attributes { passthrough = ["noinline"] } {
   %0 = llvm.mlir.constant(0 : i32) : i32
   llvm.return %0 : i32

 // CHECK-NEXT: return
 llvm.func @caller() -> (i32) {
   %0 = llvm.call @callee() : () -> (i32)
   llvm.return %0 : i32
 }
 
 // -----
 
-llvm.func @callee() -> (i32) attributes { passthrough = ["foo"] } {
-  %0 = llvm.mlir.constant(42 : i32) : i32
-  llvm.return %0 : i32
+llvm.func @callee() attributes { passthrough = ["foo", "bar"] } {
+  llvm.return
 }
 
 // CHECK-LABEL: llvm.func @caller
-// CHECK-NEXT: llvm.call @callee
-// CHECK-NEXT: return
-llvm.func @caller() -> (i32) {
-  %0 = llvm.call @callee() : () -> (i32)
-  llvm.return %0 : i32
+// CHECK-NEXT: llvm.return
+llvm.func @caller() {
+  llvm.call @callee() : () -> ()
+  llvm.return
 }
 
 // -----
 
-llvm.func @callee() -> (i32) attributes { garbageCollector = "foo" } {
-  %0 = llvm.mlir.constant(42 : i32) : i32
-  llvm.return %0 : i32
-}
+llvm.func @callee_noinline() attributes { passthrough = ["noinline"] }
+llvm.func @callee_optnone() attributes { passthrough = ["optnone"] }
+llvm.func @callee_noduplicate() attributes { passthrough = ["noduplicate"] }
+llvm.func @callee_presplitcoroutine() attributes { passthrough = ["presplitcoroutine"] }
+llvm.func @callee_returns_twice() attributes { passthrough = ["returns_twice"] }
+llvm.func @callee_strictfp() attributes { passthrough = ["strictfp"] }
 
 // CHECK-LABEL: llvm.func @caller
-// CHECK-NEXT: llvm.call @callee
-// CHECK-NEXT: return
-llvm.func @caller() -> (i32) {
-  %0 = llvm.call @callee() : () -> (i32)
-  llvm.return %0 : i32
+// CHECK-NEXT: llvm.call @callee_noinline
+// CHECK-NEXT: llvm.call @callee_optnone
+// CHECK-NEXT: llvm.call @callee_noduplicate
+// CHECK-NEXT: llvm.call @callee_presplitcoroutine
+// CHECK-NEXT: llvm.call @callee_returns_twice
+// CHECK-NEXT: llvm.call @callee_strictfp
+// CHECK-NEXT: llvm.return
+llvm.func @caller() {
+  llvm.call @callee_noinline() : () -> ()
+  llvm.call @callee_optnone() : () -> ()
+  llvm.call @callee_noduplicate() : () -> ()
+  llvm.call @callee_presplitcoroutine() : () -> ()
+  llvm.call @callee_returns_twice() : () -> ()
+  llvm.call @callee_strictfp() : () -> ()
+  llvm.return
 }
 
 // -----
 
 llvm.func @callee(%ptr : !llvm.ptr {llvm.byval = !llvm.ptr}) -> (!llvm.ptr) {
   llvm.return %ptr : !llvm.ptr
 }
 
 // CHECK-LABEL: llvm.func @caller
 // CHECK-NEXT: llvm.call @callee
 // CHECK-NEXT: return
 llvm.func @caller(%ptr : !llvm.ptr) -> (!llvm.ptr) {
   %0 = llvm.call @callee(%ptr) : (!llvm.ptr) -> (!llvm.ptr)
   llvm.return %0 : !llvm.ptr
 }
 
 // -----
 
-llvm.func @callee() {
-  llvm.return
-}
-
-// CHECK-LABEL: llvm.func @caller
-// CHECK-NEXT: llvm.call @callee
-// CHECK-NEXT: llvm.return
-llvm.func @caller() {
-  llvm.call @callee() { branch_weights = dense<42> : vector<1xi32> } : () -> ()
-  llvm.return
-}
-
-// -----
-
 llvm.func @static_alloca() -> f32 {
   %0 = llvm.mlir.constant(4 : i32) : i32
   %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
   %2 = llvm.load %1 : !llvm.ptr -> f32
   llvm.return %2 : f32
 }
 
 llvm.func @dynamic_alloca(%size : i32) -> f32 {