[MLIR] Primitive linkage lowering of FuncOp

[wsmoses]

Summary

FuncOp always lowers to an LLVM external linkage presently. This makes it impossible to define functions in mlir which are local to the current module. Until MLIR FuncOps have a more formal linkage specification, this commit allows funcop's that are already marked as private and not public to be internal upon lowering.

[jdoerfert]

Tests? (I would assume eventually funcOp needs all linkage kinds attached but this is an improvement nevertheless)

[wsmoses]

Change to llvm linkage rather than just private/public

[mehdi_amini]

mlir/test/Conversion/StandardToLLVM/convert-funcs.mlir

41

Can we use strings instead of integer here? This is just unreadable as is.

[ftynse]

mlir/test/Conversion/StandardToLLVM/convert-funcs.mlir

41

+1. We have helper functions LLVM::symbolize/stringifyLinkage to convert between integers and canonical string representations.

[wsmoses]

Trying to add the stringify/symbolize presents linkage issues, ironically. Both of those functions are (reasonable) defined only within the LLVM dialect. The parser/printer for funcop, however, is within the builtin dialect -- which means that things can't be linked if the helper code is used. Alternatively the llvm linkage attributes could be moved to the builtin dialect, but that's its own can of worms.

Thoughts on how to proceed?

[mehdi_amini]

I don't quite follow? The printer for func seems irrelevant to me here, this is about the parser/printer for the attribute itself: we need an attribute in the LLVM dialect instead of a builtin attribute and that should just work?

[ftynse]

In D108524#2978495, @wsmoses wrote:

The parser/printer for funcop, however, is within the builtin dialect -- which means that things can't be linked if the helper code is used. Alternatively the llvm linkage attributes could be moved to the builtin dialect, but that's its own can of worms.

You shouldn't be needing to modify the assembly of funcop, just attach llvm.linkage as an (unverified by op) string attribute. Then symbolize it in StandardToLLVM.cpp, which does depend on the LLVM dialect, just before passing it to LLVMFunc constructor. You can also add a verifier to the LLVM dialect that checks attributes with llvm.linkage name to only have supported strings.

In D108524#2978604, @mehdi_amini wrote:

I don't quite follow? The printer for func seems irrelevant to me here, this is about the parser/printer for the attribute itself: we need an attribute in the LLVM dialect instead of a builtin attribute and that should just work?

The dialect level attribute printer/parser is triggered for attribute types that belong to the dialect (e.g., #llvm.linkage), not attribute names. Enum attributes are just IntegerAttr that are sugared by the printer/parser of the operation that uses them most of the time.

[mehdi_amini]

In D108524#2978604, @mehdi_amini wrote:

I don't quite follow? The printer for func seems irrelevant to me here, this is about the parser/printer for the attribute itself: we need an attribute in the LLVM dialect instead of a builtin attribute and that should just work?

The dialect level attribute printer/parser is triggered for attribute types that belong to the dialect (e.g., #llvm.linkage), not attribute names. Enum attributes are just IntegerAttr that are sugared by the printer/parser of the operation that uses them most of the time.

I think we're saying the same thing? (I wrote we need an attribute in the LLVM dialect instead of a builtin attribute, the current diff shows an i64)

[ftynse]

In D108524#2980086, @mehdi_amini wrote:

In D108524#2978604, @mehdi_amini wrote:

I don't quite follow? The printer for func seems irrelevant to me here, this is about the parser/printer for the attribute itself: we need an attribute in the LLVM dialect instead of a builtin attribute and that should just work?

The dialect level attribute printer/parser is triggered for attribute types that belong to the dialect (e.g., #llvm.linkage), not attribute names. Enum attributes are just IntegerAttr that are sugared by the printer/parser of the operation that uses them most of the time.

I think we're saying the same thing? (I wrote we need an attribute in the LLVM dialect instead of a builtin attribute, the current diff shows an i64)

I'm fine with this being a string attribute (func @foo() attributes { llvm.linkage = "internal" }) as opposed to introducing a new type specifically for llvm linkage (which is currently an tablegen-backed i64 "enum" attribute pretty-printed as strings by llvm.func assembly, and changing all this is a lot of churn).

[mehdi_amini]

OK I misread what you were saying. I would prefer a first class attribute: string looks too much like JSON to me. It would also avoid having to do a conversion in StandardToLLVM which does not seem right for this.

[wsmoses]

Custom LLVM linkage attribute parsing

[wsmoses]

How does the following work as a scheme -- adding a custom printer/parser to the LLVM Linkage attribute itself. For a number of specific reasons, I'm not sure of a better way to do so than above. Specifically the tablegen-generated generatedAttributeParser and generatedAttributePrinter don't include any enum attributes since they are themselves subattributes of IntegerAttr and thus aren't registered to the dialect.

[mehdi_amini]

You can look at the llvm.fastmath flags for an example (it is defined as LLVM_FMFAttr in LLVMOps.td I believe.

[mehdi_amini]

(and void FMFAttr::print(DialectAsmPrinter &printer) and Attribute FMFAttr::parse(MLIRContext *context, DialectAsmParser &parser, Type type) in mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp)

[wsmoses]

Yeah I had tried reproducing the FMF flag variant originally, but it was unable to build in this case because the linkage is an IntegerAttr and the new Linkage wrapper attribute wasn’t allowed to have an IntegerStorage as a parameter. The reason the FMF flags get away with it is because it is a bitfield.

[mehdi_amini]

FMF is in the exact same situation. There is some tech debt in the infrastructure (we have a new member starting soon, we may ask them to look into this!), because in ODS we can't decouple generating the Enum and all the stringification associated with it, from the Attribute generation. So there is LLVM::FMFAttr and LLVM::FastMathFlagsAttr that are both generated at the moment. The latter is just an IntegerAttribute indeed: class FastmathFlagsAttr : public ::mlir::IntegerAttr {.
The trick is to just not use this at all and use another attribute.

So this is what I did in D109209 ; with a twist compared to FMF: to avoid the name collision but still use Linkage for the enum and LinkageAttr for the attribute, I moved the generation of the enum inside a sub-namespace and I inlined only the enum-related symbol and not the generated attribute class that derive from IntegerAttr.

[wsmoses]

Rebase with LLVM linkage printer

[mehdi_amini]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

270

Do we have a verifier that forbid anything else than a LinkageAttr here and makes the cast always valid?

[mehdi_amini]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

270

Ping here?

[wsmoses]

We do not but should. Apologies on the delay, will take a look at this week (was preparing for PACT and the start of the term).

[mehdi_amini]

In D108524#2995780, @wsmoses wrote:

We do not but should. Apologies on the delay, will take a look at this week (was preparing for PACT and the start of the term).

Did you get back to it since?

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 //===- StandardToLLVM.cpp - Standard to LLVM dialect conversion -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (ArrayAttr argAttrDicts = funcOp.getAllArgAttrs()) {
                "unexpected deletion of function argument");
         for (size_t j = 0; j < mapping->size; ++j)
           newArgAttrs[mapping->inputNo + j] = argAttrDicts[i];
       }
       attributes.push_back(
           rewriter.getNamedAttr(function_like_impl::getArgDictAttrName(),
                                 rewriter.getArrayAttr(newArgAttrs)));
     }
+    for (auto pair : llvm::enumerate(attributes)) {
+      if (pair.value().first == "llvm.linkage") {
+        attributes.erase(attributes.begin() + pair.index());
+        break;
+      }
+    }
 
     // Create an LLVM function, use external linkage by default until MLIR
     // functions have linkage.
+    LLVM::Linkage linkage = LLVM::Linkage::External;
+    if (funcOp->hasAttr("llvm.linkage")) {
+      linkage = funcOp->getAttr("llvm.linkage")
+                    .cast<mlir::LLVM::LinkageAttr>().getLinkage();
+    }
     auto newFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
-        funcOp.getLoc(), funcOp.getName(), llvmType, LLVM::Linkage::External,
+        funcOp.getLoc(), funcOp.getName(), llvmType, linkage,
         /*dsoLocal*/ false, attributes);
     rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                                 newFuncOp.end());
     if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), *typeConverter,
                                            &result)))
       return nullptr;
 
     return newFuncOp;

mlir/test/Conversion/StandardToLLVM/convert-funcs.mlir

// CHECK-NEXT:  llvm.br ^bb1(%arg0 : !llvm.ptr<func<void ()>>)
   br ^bb1(%arg0 : () -> ())
 
 //CHECK-NEXT: ^bb1(%0: !llvm.ptr<func<void ()>>):
 ^bb1(%bbarg: () -> ()):
 // CHECK-NEXT:  llvm.return %0 : !llvm.ptr<func<void ()>>
   return %bbarg : () -> ()
 }
 
+// CHECK-LABEL: llvm.func extern_weak @llvmlinkage(i32)
+func private @llvmlinkage(i32) attributes { "llvm.linkage" = #llvm.linkage<extern_weak> }
+
 // CHECK-LABEL: llvm.func @body(i32)
 func private @body(i32)
 
 // CHECK-LABEL: llvm.func @indirect_const_call
 // CHECK-SAME: (%[[ARG0:.*]]: i32) {
 func @indirect_const_call(%arg0: i32) {
 // CHECK-NEXT: %[[ADDR:.*]] = llvm.mlir.addressof @body : !llvm.ptr<func<void (i32)>>
   %0 = constant @body : (i32) -> ()