[MLIR][LLVM] Extend argument/result/function attribute support.

[definelicht]

Summary

Use a disallowlist approach to deselect attributes on arguments and
return values that are not yet supported, but allow the rest to be
present when inlining.

To achieve fast lookup in isLegalToInline, the set of unsupported
function attributes is cached on the inliner interface object during
construction. By using a set of StringAttrs this should just amount to a
pointer lookup.

Depends on D146633

[gysit]

LGTM

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

133

nit: isArgOrResAttrSupported ?

mlir/test/Dialect/LLVMIR/inlining.mlir

429

nit: ret_attr -> res_attr

[definelicht]

Address review comments by renaming Ret->Res.

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

Value copySize = builder.create<LLVM::ConstantOp>(
       func.getLoc(), builder.getI64Type(), builder.getI64IntegerAttr(typeSize));
   Value isVolatile = builder.create<LLVM::ConstantOp>(
       func.getLoc(), builder.getI1Type(), builder.getBoolAttr(false));
   builder.create<LLVM::MemcpyOp>(func.getLoc(), allocaOp, argument, copySize,
                                  isVolatile);
   return allocaOp;
 }
 
+/// Returns true if the given argument or result attribute is supported by the
+/// inliner, false otherwise.
+static bool isArgOrRetAttrSupported(NamedAttribute attr) {
+  if (attr.getName() == LLVM::LLVMDialect::getAlignAttrName())
+    return false;
+  if (attr.getName() == LLVM::LLVMDialect::getInAllocaAttrName())
+    return false;
+  if (attr.getName() == LLVM::LLVMDialect::getNoAliasAttrName())
+    return false;
+  return true;
+}
+
 namespace {
 struct LLVMInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
+  LLVMInlinerInterface(Dialect *dialect)
+      : DialectInlinerInterface(dialect),
+        // Cache set of StringAttrs for fast lookup in `isLegalToInline`.
+        disallowedFunctionAttrs({
+            StringAttr::get(dialect->getContext(), "noduplicate"),
+            StringAttr::get(dialect->getContext(), "noinline"),
+            StringAttr::get(dialect->getContext(), "optnone"),
+            StringAttr::get(dialect->getContext(), "presplitcoroutine"),
+            StringAttr::get(dialect->getContext(), "returns_twice"),
+            StringAttr::get(dialect->getContext(), "strictfp"),
+        }) {}
+
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
     if (!wouldBeCloned)
       return false;
     auto callOp = dyn_cast<LLVM::CallOp>(call);
     if (!callOp) {
       LLVM_DEBUG(llvm::dbgs()
                  << "Cannot inline: call is not an LLVM::CallOp\n");
       return false;
     }
     auto funcOp = dyn_cast<LLVM::LLVMFuncOp>(callable);
     if (!funcOp) {
       LLVM_DEBUG(llvm::dbgs()
                  << "Cannot inline: callable is not an LLVM::LLVMFuncOp\n");
       return false;
     }
     if (auto attrs = funcOp.getArgAttrs()) {
-      for (Attribute attr : *attrs) {
-        auto attrDict = cast<DictionaryAttr>(attr);
+      for (DictionaryAttr attrDict : attrs->getAsRange<DictionaryAttr>()) {
         for (NamedAttribute attr : attrDict) {
-          if (attr.getName() == LLVM::LLVMDialect::getByValAttrName())
-            continue;
-          // TODO: Handle all argument attributes;
+          if (!isArgOrRetAttrSupported(attr)) {
             LLVM_DEBUG(llvm::dbgs() << "Cannot inline " << funcOp.getSymName()
-                                  << ": unhandled argument attribute \""
-                                  << attr.getName() << "\"\n");
+                                    << ": unhandled argument attribute "
+                                    << attr.getName() << "\n");
             return false;
           }
         }
       }
-    // TODO: Handle result attributes;
-    if (funcOp.getResAttrs()) {
+    }
+    if (auto attrs = funcOp.getResAttrs()) {
+      for (DictionaryAttr attrDict : attrs->getAsRange<DictionaryAttr>()) {
+        for (NamedAttribute attr : attrDict) {
+          if (!isArgOrRetAttrSupported(attr)) {
             LLVM_DEBUG(llvm::dbgs() << "Cannot inline " << funcOp.getSymName()
-                              << ": unhandled result attribute\n");
+                                    << ": unhandled return attribute "
+                                    << attr.getName() << "\n");
             return false;
           }
+        }
+      }
+    }
     // TODO: Handle exceptions.
     if (funcOp.getPersonality()) {
       LLVM_DEBUG(llvm::dbgs() << "Cannot inline " << funcOp.getSymName()
                               << ": unhandled function personality\n");
       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;
-            if (disallowed.contains(stringAttr)) {
+            if (disallowedFunctionAttrs.contains(stringAttr)) {
               LLVM_DEBUG(llvm::dbgs()
                          << "Cannot inline " << funcOp.getSymName()
                          << ": found disallowed function attribute "
                          << stringAttr << "\n");
               return true;
             }
             return false;
           }))
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);
   }
+
+  // Keeping this (immutable) state on the interface allows us to look up
+  // StringAttrs instead of looking up strings, since StringAttrs are bound to
+  // the current context and thus cannot be initialized as static fields.
+  const DenseSet<StringAttr> disallowedFunctionAttrs;
 };
 
 } // end anonymous namespace
 
 void LLVM::detail::addLLVMInlinerInterface(LLVM::LLVMDialect *dialect) {
   dialect->addInterfaces<LLVMInlinerInterface>();
 }

mlir/test/Dialect/LLVMIR/inlining.mlir

 // CHECK-LABEL: llvm.func @test_byval_write_only
 // CHECK-SAME: %[[PTR:[a-zA-Z0-9_]+]]: !llvm.ptr
 // CHECK: %[[ALLOCA:.+]] = llvm.alloca %{{.+}} x f64
 // CHECK: "llvm.intr.memcpy"(%[[ALLOCA]], %[[PTR]]
 llvm.func @test_byval_write_only(%ptr : !llvm.ptr) {
   llvm.call @with_byval_arg(%ptr) : (!llvm.ptr) -> ()
   llvm.return
 }
+
+// -----
+
+llvm.func @ignored_attrs(%ptr : !llvm.ptr { llvm.inreg, llvm.nocapture, llvm.nofree, llvm.preallocated = i32, llvm.returned, llvm.alignstack = 32 : i64, llvm.writeonly, llvm.noundef, llvm.nonnull }, %x : i32 { llvm.zeroext }) -> (!llvm.ptr { llvm.noundef, llvm.inreg, llvm.nonnull }) {
+  llvm.return %ptr : !llvm.ptr
+}
+
+// CHECK-LABEL: @test_ignored_attrs
+// CHECK-NOT: llvm.call
+// CHECK-NEXT: llvm.return
+llvm.func @test_ignored_attrs(%ptr : !llvm.ptr, %x : i32) {
+  llvm.call @ignored_attrs(%ptr, %x) : (!llvm.ptr, i32) -> (!llvm.ptr)
+  llvm.return
+}
+
+// -----
+
+llvm.func @disallowed_arg_attr(%ptr : !llvm.ptr { llvm.align = 16 : i32 }) {
+  llvm.return
+}
+
+// CHECK-LABEL: @test_disallow_arg_attr
+// CHECK-NEXT: llvm.call
+llvm.func @test_disallow_arg_attr(%ptr : !llvm.ptr) {
+  llvm.call @disallowed_arg_attr(%ptr) : (!llvm.ptr) -> ()
+  llvm.return
+}
+
+// -----
+
+llvm.func @disallowed_ret_attr(%ptr : !llvm.ptr) -> (!llvm.ptr { llvm.noalias }) {
+  llvm.return %ptr : !llvm.ptr
+}
+
+// CHECK-LABEL: @test_disallow_ret_attr
+// CHECK-NEXT: llvm.call
+llvm.func @test_disallow_ret_attr(%ptr : !llvm.ptr) {
+  llvm.call @disallowed_ret_attr(%ptr) : (!llvm.ptr) -> (!llvm.ptr)
+  llvm.return
+}