[mlir:llvm] Pass arguments attributes to LLVM function arguments

[ftynse]

Please keep in mind that "passthrough" is a terrible hack primarily intended for discovering which attributes are must model first. We should be adding proper modeling for most of them, but there are just too many. Specifically, we already have inalloca, noalias and sret that are relevant for pointers, we can easily add nocapture if this is what you actually need.

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

2343–2344

Nit: please expand auto unless the type is obvious from line-level context.

2346

Ultra-nit: emitOpError prints the message of the style foo.bar op <your-message>, which would doesn't read well with this message, using emitError instead will just print <your-message>. Both are still emitted at the location of the op. Alternatively, rephrase the error message to start with "expects" so the whole text is `foo.bar op expects attribute to be an array".

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

765

[ezhulenev]

I'll abandon it in favor of proper modeling of nocapture

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

void LLVMFuncOp::print(OpAsmPrinter &p) {
   if (!body.empty()) {
     p << ' ';
     p.printRegion(body, /*printEntryBlockArgs=*/false,
                   /*printBlockTerminators=*/true);
   }
 }
 
 // Verifies LLVM- and implementation-specific properties of the LLVM func Op:
+// - arguments passthrough attributes defined as an array attribute
 // - functions don't have 'common' linkage
 // - external functions have 'external' or 'extern_weak' linkage;
 // - vararg is (currently) only supported for external functions;
 LogicalResult LLVMFuncOp::verify() {
+  for (unsigned i = 0; i < getNumArguments(); ++i) {
+    auto attrs = getArgAttrDict(i);
+    auto passthrough = attrs ? attrs.get("llvm.passthrough") : Attribute();
+    if (passthrough && !passthrough.isa<ArrayAttr>())
+      return emitOpError() << "`llvm.passthrough` attribute must be an array";
+  }
+
   if (getLinkage() == LLVM::Linkage::Common)
     return emitOpError() << "functions cannot have '"
                          << stringifyLinkage(LLVM::Linkage::Common)
                          << "' linkage";
 
   // Check to see if this function has a void return with a result attribute to
   // it. It isn't clear what semantics we would assign to that.
   if (getFunctionType().getReturnType().isa<LLVMVoidType>() &&

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

if (!value.empty())
     return emitError(loc) << "LLVM attribute '" << key
                           << "' does not expect a value, found '" << value
                           << "'";
 
   llvmFunc->addFnAttr(kind);
   return success();
 }
 
+/// Attempts to add an attribute identified by `key` to LLVM function argument
+/// `llvmArg`. Reports errors at `loc` if any. If the attribute has a kind known
+/// to LLVM IR, create the attribute of this kind, otherwise returns an error.
+static LogicalResult checkedAddLLVMArgAttribute(Location loc,
+                                                llvm::Argument *llvmArg,
+                                                StringRef key) {
+  auto kind = llvm::Attribute::getAttrKindFromName(key);
+  if (kind == llvm::Attribute::None)
+    return emitError(loc) << "unknown LLVM attribute '" << key << "'";
+
+  llvmArg->addAttr(kind);
+  return success();
+}
+
 /// Attaches the attributes listed in the given array attribute to `llvmFunc`.
 /// Reports error to `loc` if any and returns immediately. Expects `attributes`
 /// to be an array attribute containing either string attributes, treated as
 /// value-less LLVM attributes, or array attributes containing two string
 /// attributes, with the first string being the name of the corresponding LLVM
 /// attribute and the second string beings its value. Note that even integer
 /// attributes are expected to have their values expressed as strings.
 static LogicalResult
for (Attribute attr : *attributes) {
 
     if (failed(checkedAddLLVMFnAttribute(loc, llvmFunc, keyAttr.getValue(),
                                          valueAttr.getValue())))
       return failure();
   }
   return success();
 }
 
+/// Attaches the attributes listed in the given array attribute to `llvmArg`.
+/// Reports error to `loc` if any and returns immediately. Expects `attributes`
+/// to be an array attribute containing string attributes, treated as value-less
+/// LLVM attributes
+static LogicalResult
+forwardPassthroughAttributes(Location loc, Optional<ArrayAttr> attributes,
+                             llvm::Argument *llvmArg) {
+  if (!attributes)
+    return success();
+
+  for (Attribute attr : *attributes) {
+    if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
+      if (failed(
+              checkedAddLLVMArgAttribute(loc, llvmArg, stringAttr.getValue())))
+        return failure();
+      continue;
+    }
+
+    return emitError(loc)
+           << "expected 'llvm.passthrough' to contain string attributes";
+  }
+
+  return success();
+}
+
 LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   // Clear the block, branch value mappings, they are only relevant within one
   // function.
   blockMapping.clear();
   valueMapping.clear();
   branchMapping.clear();
   llvm::Function *llvmFunc = lookupFunction(func.getName());
 
for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvm::FunctionCallee llvmFuncCst = llvmModule->getOrInsertFunction(
         function.getName(),
         cast<llvm::FunctionType>(convertType(function.getFunctionType())));
     llvm::Function *llvmFunc = cast<llvm::Function>(llvmFuncCst.getCallee());
     llvmFunc->setLinkage(convertLinkageToLLVM(function.getLinkage()));
     mapFunction(function.getName(), llvmFunc);
     addRuntimePreemptionSpecifier(function.getDsoLocal(), llvmFunc);
 
-    // Forward the pass-through attributes to LLVM.
+    // Forward the pass-through function attributes to LLVM.
     if (failed(forwardPassthroughAttributes(
             function.getLoc(), function.getPassthrough(), llvmFunc)))
       return failure();
+
+    // Forward the pass-through function arguments attributes to LLVM.
+    for (unsigned i = 0; i < function.getNumArguments(); ++i) {
+      auto attrs = function.getArgAttrDict(i);
+      auto passthrough = attrs ? attrs.get("llvm.passthrough") : Attribute();
+      if (passthrough && failed(forwardPassthroughAttributes(
+                             function.getLoc(), passthrough.cast<ArrayAttr>(),
+                             llvmFunc->getArg(i))))
+        return failure();
+    }
   }
 
   return success();
 }
 
 LogicalResult ModuleTranslation::convertFunctions() {
   // Convert functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {

mlir/test/Target/LLVMIR/llvmir.mlir

llvm.func @callFenceInst() {
   // CHECK: fence release
   llvm.fence release
   // CHECK: fence release
   llvm.fence syncscope("") release
   llvm.return
 }
 
 // CHECK-LABEL: @passthrough
-// CHECK: #[[ATTR_GROUP:[0-9]*]]
-llvm.func @passthrough() attributes {passthrough = ["noinline", ["alignstack", "4"], "null_pointer_is_valid", ["foo", "bar"]]} {
+// CHECK-SAME: (ptr nocapture %[[ARG:.*]])
+// CHECK-SAME: #[[ATTR_GROUP:[0-9]*]]
+llvm.func @passthrough(%arg0: !llvm.ptr {llvm.passthrough = ["nocapture"]}) attributes {passthrough = ["noinline", ["alignstack", "4"], "null_pointer_is_valid", ["foo", "bar"]]} {
   llvm.return
 }
 
 // CHECK: attributes #[[ATTR_GROUP]] = {
 // CHECK-DAG: noinline
 // CHECK-DAG: alignstack=4
 // CHECK-DAG: null_pointer_is_valid
 // CHECK-DAG: "foo"="bar"