Add support for thread-local globals

Author: bitwalker

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

@@ -538,6 +538,26 @@ def Linkage : LLVM_EnumAttr<
   let cppNamespace = "::mlir::LLVM";
 }
 
+def ThreadLocalModeNotThreadLocal
+    : LLVM_EnumAttrCase<"NotThreadLocal", "not_thread_local", "NotThreadLocal", 0>;
+def ThreadLocalModeGeneralDynamic
+    : LLVM_EnumAttrCase<"GeneralDynamic", "general_dynamic", "GeneralDynamicTLSModel", 1>;
+def ThreadLocalModeLocalDynamic
+    : LLVM_EnumAttrCase<"LocalDynamic", "local_dynamic", "LocalDynamicTLSModel", 2>;
+def ThreadLocalModeInitialExec
+    : LLVM_EnumAttrCase<"InitialExec", "initial_exec", "InitialExecTLSModel", 3>;
+def ThreadLocalModeLocalExec
+    : LLVM_EnumAttrCase<"LocalExec", "local_exec", "LocalExecTLSModel", 4>;
+
+def ThreadLocalMode : LLVM_EnumAttr<
+    "ThreadLocalMode",
+    "::llvm::GlobalValue::ThreadLocalMode",
+    "LLVM thread local mode",
+    [ThreadLocalModeNotThreadLocal,
+     ThreadLocalModeGeneralDynamic, ThreadLocalModeLocalDynamic,
+     ThreadLocalModeInitialExec, ThreadLocalModeLocalExec]> {
+  let cppNamespace = "::mlir::LLVM";
+}
 
 def LLVM_AddressOfOp
     : LLVM_OneResultOp<"mlir.addressof">,
@@ -572,6 +592,7 @@ def LLVM_GlobalOp
                          SingleBlockImplicitTerminator<"ReturnOp">, Symbol]>,
       Arguments<(ins TypeAttr:$type, UnitAttr:$constant, StrAttr:$sym_name,
                  Linkage:$linkage,
+                 ThreadLocalMode:$tls_mode,
                  OptionalAttr<AnyAttr>:$value,
                  DefaultValuedAttr<NonNegativeI32Attr, "0">:$addr_space)> {
   let summary = "LLVM dialect global.";
@@ -593,9 +614,25 @@ def LLVM_GlobalOp
 
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, LLVMType type, "
-              "bool isConstant, Linkage linkage, StringRef name, "
+              "bool isConstant, Linkage linkage, ThreadLocalMode tlsMode,"
+              "StringRef name,"
+              "Attribute value, unsigned addrSpace = 0, "
+              "ArrayRef<NamedAttribute> attrs = {}">,
+    OpBuilder<"Builder *builder, OperationState &result, LLVMType type, "
+              "bool isConstant, Linkage linkage,"
+              "StringRef name,"
               "Attribute value, unsigned addrSpace = 0, "
-              "ArrayRef<NamedAttribute> attrs = {}">
+              "ArrayRef<NamedAttribute> attrs = {}",
+              [{ build(builder,
+                       result,
+                       type,
+                       isConstant,
+                       linkage,
+                       ThreadLocalMode::NotThreadLocal,
+                       name,
+                       value,
+                       addrSpace,
+                       attrs); }]>,
   ];
 
   let extraClassDeclaration = [{

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

@@ -912,9 +912,14 @@ static LogicalResult verify(AddressOfOp op) {
 /// the name of the attribute in ODS.
 static StringRef getLinkageAttrName() { return "linkage"; }
 
+/// Returns the name used for the TLS attribute. This *must* correspond to
+/// the name of the attribute in ODS.
+static StringRef getThreadLocalModeAttrName() { return "tls_mode"; }
+
+
 void GlobalOp::build(Builder *builder, OperationState &result, LLVMType type,
-                     bool isConstant, Linkage linkage, StringRef name,
-                     Attribute value, unsigned addrSpace,
+                     bool isConstant, Linkage linkage, ThreadLocalMode tlsMode,
+                     StringRef name, Attribute value, unsigned addrSpace,
                      ArrayRef<NamedAttribute> attrs) {
   result.addAttribute(SymbolTable::getSymbolAttrName(),
                       builder->getStringAttr(name));
@@ -925,6 +930,8 @@ void GlobalOp::build(Builder *builder, OperationState &result, LLVMType type,
     result.addAttribute("value", value);
   result.addAttribute(getLinkageAttrName(), builder->getI64IntegerAttr(
                                                 static_cast<int64_t>(linkage)));
+  result.addAttribute(getThreadLocalModeAttrName(), builder->getI64IntegerAttr(
+                                                static_cast<int64_t>(tlsMode)));
   if (addrSpace != 0)
     result.addAttribute("addr_space", builder->getI32IntegerAttr(addrSpace));
   result.attributes.append(attrs.begin(), attrs.end());
@@ -933,6 +940,9 @@ void GlobalOp::build(Builder *builder, OperationState &result, LLVMType type,
 
 static void printGlobalOp(OpAsmPrinter &p, GlobalOp op) {
   p << op.getOperationName() << ' ' << stringifyLinkage(op.linkage()) << ' ';
+  auto tlsMode = op.tls_mode();
+  if (tlsMode != ThreadLocalMode::NotThreadLocal)
+    p << stringifyThreadLocalMode(tlsMode) << ' ';
   if (op.constant())
     p << "constant ";
   p.printSymbolName(op.sym_name());
@@ -942,7 +952,7 @@ static void printGlobalOp(OpAsmPrinter &p, GlobalOp op) {
   p << ')';
   p.printOptionalAttrDict(op.getAttrs(),
                           {SymbolTable::getSymbolAttrName(), "type", "constant",
-                           "value", getLinkageAttrName()});
+                           "value", getLinkageAttrName(), getThreadLocalModeAttrName()});
 
   // Print the trailing type unless it's a string global.
   if (op.getValueOrNull().dyn_cast_or_null<StringAttr>())
@@ -976,14 +986,15 @@ template <typename Ty> struct EnumTraits {};
   }
 
 REGISTER_ENUM_TYPE(Linkage);
+REGISTER_ENUM_TYPE(ThreadLocalMode);
 } // end namespace
 
 template <typename EnumTy>
 static ParseResult parseOptionalLLVMKeyword(OpAsmParser &parser,
                                             OperationState &result,
                                             StringRef name) {
   SmallVector<StringRef, 10> names;
-  for (unsigned i = 0, e = getMaxEnumValForLinkage(); i <= e; ++i)
+  for (unsigned i = 0, e = EnumTraits<EnumTy>::getMaxEnumVal(); i <= e; ++i)
     names.push_back(EnumTraits<EnumTy>::stringify(static_cast<EnumTy>(i)));
 
   int index = parseOptionalKeywordAlternative(parser, names);
@@ -993,7 +1004,7 @@ static ParseResult parseOptionalLLVMKeyword(OpAsmParser &parser,
   return success();
 }
 
-// operation ::= `llvm.mlir.global` linkage `constant`? `@` identifier
+// operation ::= `llvm.mlir.global` linkage tls-mode? `constant`? `@` identifier
 //               `(` attribute? `)` attribute-list? (`:` type)? region?
 //
 // The type can be omitted for string attributes, in which case it will be
@@ -1003,6 +1014,8 @@ static ParseResult parseGlobalOp(OpAsmParser &parser, OperationState &result) {
                                                getLinkageAttrName())))
     return parser.emitError(parser.getCurrentLocation(), "expected linkage");
 
+  parseOptionalLLVMKeyword<ThreadLocalMode>(parser, result, getThreadLocalModeAttrName());
+
   if (succeeded(parser.parseOptionalKeyword("constant")))
     result.addAttribute("constant", parser.getBuilder().getUnitAttr());
 

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -21,6 +21,8 @@
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/SetVector.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -449,14 +451,15 @@ void ModuleTranslation::convertGlobals() {
     }
 
     auto linkage = convertLinkageToLLVM(op.linkage());
+    auto tlsMode = convertThreadLocalModeToLLVM(op.tls_mode());
     bool anyExternalLinkage =
         (linkage == llvm::GlobalVariable::ExternalLinkage ||
          linkage == llvm::GlobalVariable::ExternalWeakLinkage);
     auto addrSpace = op.addr_space().getLimitedValue();
     auto *var = new llvm::GlobalVariable(
         *llvmModule, type, op.constant(), linkage,
         anyExternalLinkage ? nullptr : cst, op.sym_name(),
-        /*InsertBefore=*/nullptr, llvm::GlobalValue::NotThreadLocal, addrSpace);
+        /*InsertBefore=*/nullptr, tlsMode, addrSpace);
 
     globalsMapping.try_emplace(op, var);
   }