[mlir] LLVMFuncOp: provide a capability to pass attributes through to LLVM IR

[ftynse]

Summary

LLVM IR functions can have arbitrary attributes attached to them, some of which
affect may affect code transformations. Until we can model all attributes
consistently, provide a pass-through mechanism that forwards attributes from
the LLVMFuncOp in MLIR to LLVM IR functions during translation. This mechanism
relies on LLVM IR being able to recognize string representations of the
attributes and performs some additional checking to avoid hitting assertions
within LLVM code.

[nicolasvasilache]

Thanks Alex, I can confirm that with this, the expected behavior is reached on the LLVM matrix front.
Also + @fhahn

[mehdi_amini]

mlir/docs/Dialects/LLVM.md

107

That seems a bit like a hack to me: when importing from LLVM are we gonna encode attributes this way? Are we gonna honor them?
We're getting into the "json of compiler" here.

[nicolasvasilache]

mlir/docs/Dialects/LLVM.md

107

Isn't the LLVM string = string attribute the issue here?
Would an alternative require giving more semantics than LLVM has?

In other words, until/unless LLVM force more semantics on those attributes, isn't it a raising problem to import from LLVM and guesstimate types?

[mehdi_amini]

mlir/docs/Dialects/LLVM.md

107

LLVM has string attributes only for target specific I believe, otherwise there are rules for "IR-level" attributes: https://github.com/llvm/llvm-project/blob/master/llvm/include/llvm/IR/Attributes.td
(I am not sure I totally understood what you meant though, so my answer may be off here)

[ftynse]

mlir/docs/Dialects/LLVM.md

107

I agree that it feels hacky, but so does LLVM's approach. You can add arbitrary quoted strings as function attributes, which some passes of various levels of experimental-ness rely on, and we got to support that. We have a need for that today to support some forms of vectorization. While we could restrict this feature to only support string attributes, we would need to encode a list of allowed attributes somewhere, which sounds equivalent to just building support for them, and we don't have time to invest in that right now.

I expect that we will eventually have first-class modeling for the equivalents of common attributes that may differ from what LLVM does with attributes (e.g., for inlining or FP math options). Until we do, we need to have a way to attach these attributes in the generated code for the purposes of code generation.

https://github.com/llvm/llvm-project/blob/master/llvm/include/llvm/IR/Attributes.td is of little help because it does not say which attributes are allowed on functions (as opposed to, e.g., function attributes) and which of them are allowed to have values. Furthermore, there are attributes allowed on functions listed in https://llvm.org/docs/LangRef.html#function-attributes that are not listed in Attributes.td.

[ftynse]

mlir/docs/Dialects/LLVM.md

107

(as opposed to, e.g., function attributes)

I meant function _argument_ attributes

[mehdi_amini]

mlir/docs/Dialects/LLVM.md

107

I agree that LLVM is not perfect, but this is much uglier and I'm gonna be concerned by any code that start to manipulate or emit this. It'll have to be very contained.

[ftynse]

mlir/docs/Dialects/LLVM.md

107

The code Nicolas has uses platform-specific attributes, which are anyway strings that are not registered anyway. I don't see how to "fix" that without imposing our choices further down the pipeline. If someone suddenly starts to use this for an actual function attribute listed in LangRef, we will be able to model that attribute specifically, without having to prematurely overdesign all of them.

mlir/docs/Dialects/LLVM.md

 
 // A function with `internal` linkage.
 llvm.func internal @internal_func() {
   llvm.return
 }
 
 ```
 
+#### Attribute pass-through
+
+An LLVM IR dialect function provides a mechanism to forward function-level
+attributes to LLVM IR using the `passthrough` attribute. This is an array
+attribute containing either string attributes or array attributes. In the former
+case, the value of the string is interpreted as the name of LLVM IR function
+attribute. In the latter case, the array is expected to contain exactly two
+string attributes, the first corresponding to the name of LLVM IR function
+attribute, and the second corresponding to its value. Note that even integer
+LLVM IR function attributes have their value represented in the string form.
+
+Example:
+
+```mlir
+llvm.func @func() attributes {
+  passthrough = ["noinline",           // value-less attribute
+                 ["alignstack", "4"],  // integer attribute with value
+                 ["other", "attr"]]    // attrbute unknown to LLVM
+} {
+  llvm.return
+}
+```
+
+If the attribute is not known to LLVM IR, it will be attached as a string
+attribute.
+
 ### LLVM IR operations
 
 The following operations are currently supported. The semantics of these
 operations corresponds to the semantics of the similarly-named LLVM IR
 instructions.
 
 #### Integer binary arithmetic operations
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_GlobalOp
   let parser = "return parseGlobalOp(parser, result);";
   let verifier = "return ::verify(*this);";
 }
 
 def LLVM_LLVMFuncOp
     : LLVM_ZeroResultOp<"func", [IsolatedFromAbove, FunctionLike, Symbol]>,
       Arguments<(ins DefaultValuedAttr<Linkage,
                                        "Linkage::External">:$linkage,
-                     OptionalAttr<FlatSymbolRefAttr>:$personality)> {
+                     OptionalAttr<FlatSymbolRefAttr>:$personality,
+                     OptionalAttr<ArrayAttr>:$passthrough)> {
   let summary = "LLVM dialect function, has wrapped LLVM IR function type";
 
   let regions = (region AnyRegion:$body);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, StringRef name, "

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (Block &b : f.getBlocks()) {
     if (blocks.count(&b) == 0)
       topologicalSortImpl(blocks, &b);
   }
   assert(blocks.size() == f.getBlocks().size() && "some blocks are not sorted");
 
   return blocks;
 }
 
+/// Attempts to add an attribute identified by `key`, optionally with the given
+/// `value` to LLVM function `llvmFunc`. Reports errors at `loc` if any. If the
+/// attribute has a kind known to LLVM IR, create the attribute of this kind,
+/// otherwise keep it as a string attribute. Performs additional checks for
+/// attributes known to have or not have a value in order to avoid assertions
+/// inside LLVM upon construction.
+static LogicalResult checkedAddLLVMFnAttribute(Location loc,
+                                               llvm::Function *llvmFunc,
+                                               StringRef key,
+                                               StringRef value = StringRef()) {
+  auto kind = llvm::Attribute::getAttrKindFromName(key);
+  if (kind == llvm::Attribute::None) {
+    llvmFunc->addFnAttr(key, value);
+    return success();
+  }
+
+  if (llvm::Attribute::doesAttrKindHaveArgument(kind)) {
+    if (value.empty())
+      return emitError(loc) << "LLVM attribute '" << key << "' expects a value";
+
+    int result;
+    if (!value.getAsInteger(/*Radix=*/0, result))
+      llvmFunc->addFnAttr(
+          llvm::Attribute::get(llvmFunc->getContext(), kind, result));
+    else
+      llvmFunc->addFnAttr(key, value);
+    return success();
+  }
+
+  if (!value.empty())
+    return emitError(loc) << "LLVM attribute '" << key
+                          << "' does not expect a value, found '" << value
+                          << "'";
+
+  llvmFunc->addFnAttr(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
+forwardPassthroughAttributes(Location loc, Optional<ArrayAttr> attributes,
+                             llvm::Function *llvmFunc) {
+  if (!attributes)
+    return success();
+
+  for (Attribute attr : *attributes) {
+    if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
+      if (failed(
+              checkedAddLLVMFnAttribute(loc, llvmFunc, stringAttr.getValue())))
+        return failure();
+      continue;
+    }
+
+    auto arrayAttr = attr.dyn_cast<ArrayAttr>();
+    if (!arrayAttr || arrayAttr.size() != 2)
+      return emitError(loc)
+             << "expected 'passthrough' to contain string or array attributes";
+
+    auto keyAttr = arrayAttr[0].dyn_cast<StringAttr>();
+    auto valueAttr = arrayAttr[1].dyn_cast<StringAttr>();
+    if (!keyAttr || !valueAttr)
+      return emitError(loc)
+             << "expected arrays within 'passthrough' to contain two strings";
+
+    if (failed(checkedAddLLVMFnAttribute(loc, llvmFunc, keyAttr.getValue(),
+                                         valueAttr.getValue())))
+      return failure();
+  }
+  return success();
+}
+
 LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   // Clear the block and value mappings, they are only relevant within one
   // function.
   blockMapping.clear();
   valueMapping.clear();
   llvm::Function *llvmFunc = functionMapping.lookup(func.getName());
 
   // Translate the debug information for this function.

 
 LogicalResult ModuleTranslation::convertFunctions() {
   // Declare all functions first because there may be function calls that form a
   // call graph with cycles.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvm::FunctionCallee llvmFuncCst = llvmModule->getOrInsertFunction(
         function.getName(),
         cast<llvm::FunctionType>(function.getType().getUnderlyingType()));
-    assert(isa<llvm::Function>(llvmFuncCst.getCallee()));
-    functionMapping[function.getName()] =
-        cast<llvm::Function>(llvmFuncCst.getCallee());
+    llvm::Function *llvmFunc = cast<llvm::Function>(llvmFuncCst.getCallee());
+    functionMapping[function.getName()] = llvmFunc;
+
+    // Forward the pass-through attributes to LLVM.
+    if (failed(forwardPassthroughAttributes(function.getLoc(),
+                                            function.passthrough(), llvmFunc)))
+      return failure();
   }
 
   // Convert functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     // Ignore external functions.
     if (function.isExternal())
       continue;
 

mlir/test/Target/llvmir-invalid.mlir

llvm.func @no_nested_struct() -> !llvm<"[2 x [2 x [2 x {i32}]]]"> {
   %0 = llvm.mlir.constant(dense<[[[1, 2], [3, 4]], [[42, 43], [44, 45]]]> : tensor<2x2x2xi32>) : !llvm<"[2 x [2 x [2 x {i32}]]]">
   llvm.return %0 : !llvm<"[2 x [2 x [2 x {i32}]]]">
 }
 
 // -----
 
 // expected-error @+1 {{unsupported constant value}}
 llvm.mlir.global internal constant @test([2.5, 7.4]) : !llvm<"[2 x double]">
+
+// -----
+
+// expected-error @+1 {{LLVM attribute 'noinline' does not expect a value}}
+llvm.func @passthrough_unexpected_value() attributes {passthrough = [["noinline", "42"]]}
+
+// -----
+
+// expected-error @+1 {{LLVM attribute 'alignstack' expects a value}}
+llvm.func @passthrough_expected_value() attributes {passthrough = ["alignstack"]}
+
+// -----
+
+// expected-error @+1 {{expected 'passthrough' to contain string or array attributes}}
+llvm.func @passthrough_wrong_type() attributes {passthrough = [42]}
+
+// -----
+
+// expected-error @+1 {{expected arrays within 'passthrough' to contain two strings}}
+llvm.func @passthrough_wrong_type() attributes {passthrough = [[42, 42]]}

mlir/test/Target/llvmir.mlir

llvm.func @callFenceInst() {
   // CHECK: fence syncscope("agent") release
   llvm.fence syncscope("agent") release
   // 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"]]} {
+  llvm.return
+}
+
+// CHECK: attributes #[[ATTR_GROUP]] = {
+// CHECK-DAG: noinline
+// CHECK-DAG: alignstack=4
+// CHECK-DAG: "null-pointer-is-valid"
+// CHECK-DAG: "foo"="bar"