[mlir] Add support for LLVM's dso_local attr

[fdeazeve]

Summary

This patch brings support for setting runtime preemption specifiers of
LLVM's GlobalValues. In LLVM semantics, if the dso_local attribute
is not explicitly requested, then it is inferred based on linkage and
visibility. We model this same behavior with a UnitAttribute: if it is
present, then we explicitly request the GlobalValue to marked as
dso_local, otherwise we rely on the GlobalValue itself to make this
decision.

[fdeazeve]

This was not meant to be a new review, but an update to D104983. Not sure what I did wrong with ARC :(

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

 
 def LLVM_GlobalOp : LLVM_Op<"mlir.global",
     [IsolatedFromAbove, SingleBlockImplicitTerminator<"ReturnOp">, Symbol]> {
   let arguments = (ins
     TypeAttr:$type,
     UnitAttr:$constant,
     StrAttr:$sym_name,
     Linkage:$linkage,
+    UnitAttr:$dso_local,
     OptionalAttr<AnyAttr>:$value,
     OptionalAttr<I64Attr>:$alignment,
     DefaultValuedAttr<Confined<I32Attr, [IntNonNegative]>, "0">:$addr_space,
     OptionalAttr<UnnamedAddr>:$unnamed_addr,
     OptionalAttr<StrAttr>:$section
   );
   let summary = "LLVM dialect global.";
   let description = [{
def LLVM_GlobalOp : LLVM_Op<"mlir.global",
   }];
   let regions = (region AnyRegion:$initializer);
 
   let builders = [
     OpBuilder<(ins "Type":$type, "bool":$isConstant, "Linkage":$linkage,
       "StringRef":$name, "Attribute":$value,
       CArg<"uint64_t", "0">:$alignment,
       CArg<"unsigned", "0">:$addrSpace,
+      CArg<"bool", "false">:$dsoLocal,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let extraClassDeclaration = [{
     /// Return the LLVM type of the global.
     Type getType() {
       return type();
     }
let description = [{
     // A function with `internal` linkage.
     llvm.func internal @internal_func() {
       llvm.return
     }
     ```
   }];
 
   let arguments = (ins DefaultValuedAttr<Linkage, "Linkage::External">:$linkage,
+                   UnitAttr:$dso_local,
                    OptionalAttr<FlatSymbolRefAttr>:$personality,
                    OptionalAttr<ArrayAttr>:$passthrough);
 
   let regions = (region AnyRegion:$body);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<(ins "StringRef":$name, "Type":$type,
       CArg<"Linkage", "Linkage::External">:$linkage,
+      CArg<"bool", "false">:$dsoLocal,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs,
       CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs)>
   ];
 
   let extraClassDeclaration = [{
     // Add an entry block to an empty function, and set up the block arguments
     // to match the signature of the function.
     Block *addEntryBlock();

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp,
   }
   // Add a dialect specific kernel attribute in addition to GPU kernel
   // attribute. The former is necessary for further translation while the
   // latter is expected by gpu.launch_func.
   if (gpuFuncOp.isKernel())
     attributes.emplace_back(kernelAttributeName, rewriter.getUnitAttr());
   auto llvmFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
       gpuFuncOp.getLoc(), gpuFuncOp.getName(), funcType,
-      LLVM::Linkage::External, attributes);
+      LLVM::Linkage::External, /*dsoLocal*/ false, attributes);
 
   {
     // Insert operations that correspond to converted workgroup and private
     // memory attributions to the body of the function. This must operate on
     // the original function, before the body region is inlined in the new
     // function to maintain the relation between block arguments and the
     // parent operation that assigns their semantics.
     OpBuilder::InsertionGuard guard(rewriter);

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

static void wrapForExternalCallers(OpBuilder &rewriter, Location loc,
   filterFuncAttributes(funcOp->getAttrs(), /*filterArgAttrs=*/false,
                        attributes);
   Type wrapperFuncType;
   bool resultIsNowArg;
   std::tie(wrapperFuncType, resultIsNowArg) =
       typeConverter.convertFunctionTypeCWrapper(type);
   auto wrapperFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
       loc, llvm::formatv("_mlir_ciface_{0}", funcOp.getName()).str(),
-      wrapperFuncType, LLVM::Linkage::External, attributes);
+      wrapperFuncType, LLVM::Linkage::External, /*dsoLocal*/ false, attributes);
 
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPointToStart(wrapperFuncOp.addEntryBlock());
 
   SmallVector<Value, 8> args;
   size_t argOffset = resultIsNowArg ? 1 : 0;
   for (auto &en : llvm::enumerate(type.getInputs())) {
     Value arg = wrapperFuncOp.getArgument(en.index() + argOffset);
static void wrapExternalFunction(OpBuilder &builder, Location loc,
 
   SmallVector<NamedAttribute, 4> attributes;
   filterFuncAttributes(funcOp->getAttrs(), /*filterArgAttrs=*/false,
                        attributes);
 
   // Create the auxiliary function.
   auto wrapperFunc = builder.create<LLVM::LLVMFuncOp>(
       loc, llvm::formatv("_mlir_ciface_{0}", funcOp.getName()).str(),
-      wrapperType, LLVM::Linkage::External, attributes);
+      wrapperType, LLVM::Linkage::External, /*dsoLocal*/ false, attributes);
 
   builder.setInsertionPointToStart(newFuncOp.addEntryBlock());
 
   // Get a ValueRange containing arguments.
   FunctionType type = funcOp.getType();
   SmallVector<Value, 8> args;
   args.reserve(type.getNumInputs());
   ValueRange wrapperArgsRange(newFuncOp.getArguments());
if (ArrayAttr argAttrDicts = funcOp.getAllArgAttrs()) {
           rewriter.getNamedAttr(function_like_impl::getArgDictAttrName(),
                                 rewriter.getArrayAttr(newArgAttrs)));
     }
 
     // Create an LLVM function, use external linkage by default until MLIR
     // functions have linkage.
     auto newFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
         funcOp.getLoc(), funcOp.getName(), llvmType, LLVM::Linkage::External,
-        attributes);
+        /*dsoLocal*/ false, attributes);
     rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                                 newFuncOp.end());
     if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), *typeConverter,
                                            &result)))
       return nullptr;
 
     return newFuncOp;
   }

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

 
 /// Returns the name used for the unnamed_addr attribute. This *must* correspond
 /// to the name of the attribute in ODS.
 static StringRef getUnnamedAddrAttrName() { return "unnamed_addr"; }
 
 void GlobalOp::build(OpBuilder &builder, OperationState &result, Type type,
                      bool isConstant, Linkage linkage, StringRef name,
                      Attribute value, uint64_t alignment, unsigned addrSpace,
-                     ArrayRef<NamedAttribute> attrs) {
+                     bool dsoLocal, ArrayRef<NamedAttribute> attrs) {
   result.addAttribute(SymbolTable::getSymbolAttrName(),
                       builder.getStringAttr(name));
   result.addAttribute("type", TypeAttr::get(type));
   if (isConstant)
     result.addAttribute("constant", builder.getUnitAttr());
   if (value)
     result.addAttribute("value", value);
+  if (dsoLocal)
+    result.addAttribute("dso_local", builder.getUnitAttr());
 
   // Only add an alignment attribute if the "alignment" input
   // is different from 0. The value must also be a power of two, but
   // this is tested in GlobalOp::verify, not here.
   if (alignment != 0)
     result.addAttribute("alignment", builder.getI64IntegerAttr(alignment));
 
   result.addAttribute(getLinkageAttrName(),
Block *LLVMFuncOp::addEntryBlock() {
   LLVMFunctionType type = getType();
   for (unsigned i = 0, e = type.getNumParams(); i < e; ++i)
     entry->addArgument(type.getParamType(i));
   return entry;
 }
 
 void LLVMFuncOp::build(OpBuilder &builder, OperationState &result,
                        StringRef name, Type type, LLVM::Linkage linkage,
-                       ArrayRef<NamedAttribute> attrs,
+                       bool dsoLocal, ArrayRef<NamedAttribute> attrs,
                        ArrayRef<DictionaryAttr> argAttrs) {
   result.addRegion();
   result.addAttribute(SymbolTable::getSymbolAttrName(),
                       builder.getStringAttr(name));
   result.addAttribute("type", TypeAttr::get(type));
   result.addAttribute(getLinkageAttrName(),
                       builder.getI64IntegerAttr(static_cast<int64_t>(linkage)));
   result.attributes.append(attrs.begin(), attrs.end());
+  if (dsoLocal)
+    result.addAttribute("dso_local", builder.getUnitAttr());
   if (argAttrs.empty())
     return;
 
   assert(type.cast<LLVMFunctionType>().getNumParams() == argAttrs.size() &&
          "expected as many argument attribute lists as arguments");
   function_like_impl::addArgAndResultAttrs(builder, result, argAttrs,
                                            /*resultAttrs=*/llvm::None);
 }

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 /// initializer.
 static bool shouldDropGlobalInitializer(llvm::GlobalValue::LinkageTypes linkage,
                                         llvm::Constant *cst) {
   return (linkage == llvm::GlobalVariable::ExternalLinkage &&
           isa<llvm::UndefValue>(cst)) ||
          linkage == llvm::GlobalVariable::ExternalWeakLinkage;
 }
 
+/// Sets the runtime preemption specifier of `gv` to dso_local if
+/// `dsoLocalRequested` is true, otherwise it is left unchanged.
+static void addRuntimePreemptionSpecifier(bool dsoLocalRequested,
+                                          llvm::GlobalValue *gv) {
+  if (dsoLocalRequested)
+    gv->setDSOLocal(true);
+}
+
 /// Create named global variables that correspond to llvm.mlir.global
 /// definitions.
 LogicalResult ModuleTranslation::convertGlobals() {
   for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>()) {
     llvm::Type *type = convertType(op.getType());
     llvm::Constant *cst = llvm::UndefValue::get(type);
     if (op.getValueOrNull()) {
       // String attributes are treated separately because they cannot appear as
auto *var = new llvm::GlobalVariable(
         /*InsertBefore=*/nullptr, llvm::GlobalValue::NotThreadLocal, addrSpace);
 
     if (op.unnamed_addr().hasValue())
       var->setUnnamedAddr(convertUnnamedAddrToLLVM(*op.unnamed_addr()));
 
     if (op.section().hasValue())
       var->setSection(*op.section());
 
+    addRuntimePreemptionSpecifier(op.dso_local(), var);
+
     Optional<uint64_t> alignment = op.alignment();
     if (alignment.hasValue())
       var->setAlignment(llvm::MaybeAlign(alignment.getValue()));
 
     globalsMapping.try_emplace(op, var);
   }
 
   // Convert global variable bodies. This is done after all global variables
LogicalResult ModuleTranslation::convertFunctionSignatures() {
   // call graph with cycles, or global initializers that reference functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvm::FunctionCallee llvmFuncCst = llvmModule->getOrInsertFunction(
         function.getName(),
         cast<llvm::FunctionType>(convertType(function.getType())));
     llvm::Function *llvmFunc = cast<llvm::Function>(llvmFuncCst.getCallee());
     llvmFunc->setLinkage(convertLinkageToLLVM(function.linkage()));
     mapFunction(function.getName(), llvmFunc);
+    addRuntimePreemptionSpecifier(function.dso_local(), llvmFunc);
 
     // Forward the pass-through attributes to LLVM.
     if (failed(forwardPassthroughAttributes(function.getLoc(),
                                             function.passthrough(), llvmFunc)))
       return failure();
   }
 
   return success();

mlir/test/Target/LLVMIR/llvmir.mlir

 // CHECK: @no_unnamed_addr = private constant i64 42
 llvm.mlir.global private constant @no_unnamed_addr(42 : i64) : i64
 // CHECK: @local_unnamed_addr = private local_unnamed_addr constant i64 42
 llvm.mlir.global private local_unnamed_addr constant @local_unnamed_addr(42 : i64) : i64
 // CHECK: @unnamed_addr = private unnamed_addr constant i64 42
 llvm.mlir.global private unnamed_addr constant @unnamed_addr(42 : i64) : i64
 
 //
+// dso_local attribute.
+//
+
+llvm.mlir.global @has_dso_local(42 : i64) {dso_local} : i64
+// CHECK: @has_dso_local = dso_local global i64 42
+
+//
 // Section attribute.
 //
 
 // CHECK: @sectionvar = internal constant [10 x i8] c"teststring", section ".mysection"
 llvm.mlir.global internal constant @sectionvar("teststring")  {section = ".mysection"}: !llvm.array<10 x i8>
 
 //
 // Declarations of the allocation functions to be linked against. These are

 //
 
 // CHECK: define internal void @func_internal
 llvm.func internal @func_internal() {
   llvm.return
 }
 
 //
+// dso_local attribute.
+//
+
+// CHECK: define dso_local void @dso_local_func
+llvm.func @dso_local_func() attributes {dso_local} {
+  llvm.return
+}
+
+//
 // MemRef type conversion, allocation and communication with functions.
 //
 
 // CHECK-LABEL: define void @memref_alloc()
 llvm.func @memref_alloc() {
 // CHECK-NEXT: %{{[0-9]+}} = call i8* @malloc(i64 400)
 // CHECK-NEXT: %{{[0-9]+}} = bitcast i8* %{{[0-9]+}} to float*
 // CHECK-NEXT: %{{[0-9]+}} = insertvalue { float* } undef, float* %{{[0-9]+}}, 0