[MLIR][OpenMP] Add support for threadprivate directive

[peixin]

Summary

This supports the threadprivate directive in OpenMP dialect following
the OpenMP 5.1 [2.21.2] standard. Also lowering to LLVM IR using OpenMP
IRBduiler.

[arnamoy10]

1. Do we need a verifier for this?
2. If yes, then we should add test cases (in invalid.mlir). Also we probably should add a test case to verify the printer in ops.mlir.

[shraiysh]

In addition to Arnamoy's comments, can we have assembly format based syntax for this? It looks like the following would work -

$sym_addr `:` type($sym_addr) `->` type($res)

[peixin]

In D123350#3438111, @shraiysh wrote:

In addition to Arnamoy's comments, can we have assembly format based syntax for this? It looks like the following would work -

$sym_addr `:` type($sym_addr) `->` type($res)

Thanks for suggestions. I didn't notice the mlir format has changed. I am fixing the CI fail.

[peixin]

1. Remove let parser/printer due to update in upstream MLIR.
2. Replace syn_name() with getSymName() due to update in upstream MLIR.
3. Add test case in ops.mlir

[peixin]

In D123350#3437856, @arnamoy10 wrote:

1. Do we need a verifier for this?

This operation def is implementation defined. I am not aware of what I should verify currently. The restrictions in OpenMP standard are not suitable to verify here, which should be in semantic analysis. Please let me know if anyone finds what should to verify.

[shraiysh]

Thanks for the patch. A few comments.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

1327

nit: llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder) also works.

1333

Instead of an assert, can we return a graceful error here?

1340

nit: builder.getInt8PtrTy()

1343

[suggestion] Instead of exposing this function and changing the LLVM::ModuleTranslation class, we can use builder.GetInsertBlock()->getModule()->getDataLayout() here.

1344

For size of types, we can use the following to get the type size directly instead of using getTypeStoreSize and getFixedSize.

llvm::ConstantInt *size = Builder.getInt64(datalayout.getTypeSizeInBits(type))

[kiranchandramohan]

LGTM. A few nits. Please wait for @shraiysh and @arnamoy10.

Can you add a note to the Summary saying the following or something similar?
The current implementation uses the OpenMP runtime to provide thread-local storage. We will (in future) provide support for using the TLS feature of the LLVM IR.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

896

Nit: This operation takes in the address of a symbol that represents the original variable and returns the address of its thread local storage. All occurrences of threadprivate variables in a parallel region should use the TLS returned by this operation.

898

Nit: refers to the address of the symbol, which is a pointer to the original global variable.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

1333

Nit: When the assert is removed, please replace the auto above with a type.

[peixin]

1. Fix datalayout related comments.
2. Fix Operation description.

[peixin]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

898

Made changes: delete global.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

1333

Technically, AllocaOp works also ok here. classic-flang uses for variables in main program for fortran code. Also, OpenMP standard does not mention the variable must be global for fortran. So I add assert instead of reporting the error in verify. What do you think? @kiranchandramohan @shraiysh

1344

getTypeSizeInBits is not correct. It seems it cannot handle some special data type. I followed clang using getTypeStoreSize.

[shraiysh]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

1333

You can report an error here itself. This error will not be inside verify, instead it will be an error while translating to LLVM IR. That allows using AllocaOp in the dialect, but reports an error only while translating to LLVM IR. Example - here.

1344

I had asked about this in the discord channel and they suggested using getTypeSizeInBits to get the size of a type. Can you please specify for which types it cannot handle?

[peixin]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

1333

That makes sense. Thanks for the advice. Will fix it.

1344

https://github.com/llvm/llvm-project/blob/5562d9b3c006f3e376d98ee7b9f3c6c8608ddbee/llvm/include/llvm/IR/DataLayout.h#L467-L477

For example, returns 5 for i36 gives 40 in bits for i36. Clang uses that, so I use it. Using getTypeStoreSize works ok for all tested fortran data types, so I did not dig too much the difference between using getTypeStoreSize or getTypeSizeInBits. Fortran code doesn't have i36, not sure if c/c++ uses it.

[peixin]

Change assert into error reporting. Remove the auto. Add the test case for the error reporting

[peixin]

@arnamoy10 Is this patch OK to you now?

[shraiysh]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

1344

Okay thank you.

[arnamoy10]

LGTM

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

let extraClassDeclaration = [{
 
     /// Returns the `atomic.update` operation inside the region, if any.
     /// Otherwise, it returns nullptr.
     AtomicUpdateOp getAtomicUpdateOp();
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// [5.1] 2.21.2 threadprivate Directive
+//===----------------------------------------------------------------------===//
+
+def ThreadprivateOp : OpenMP_Op<"threadprivate"> {
+  let summary = "threadprivate directive";
+  let description = [{
+    The threadprivate directive specifies that variables are replicated, with
+    each thread having its own copy.
+
+    This model takes in the addressing of symbol and returns the addressing of
+    thread local storage (TLS) for the original variable.
+
+    The `sym_addr` refers to the addressing of the symbol, which is pointer to
+    original global variable.
+  }];
+
+  let arguments = (ins OpenMP_PointerLikeType:$sym_addr);
+  let results = (outs OpenMP_PointerLikeType);
+  let parser = [{ return parseThreadprivateOp(parser, result); }];
+  let printer = [{ return printThreadprivateOp(p, *this); }];
+}
+
+//===----------------------------------------------------------------------===//
 // 2.19.5.7 declare reduction Directive
 //===----------------------------------------------------------------------===//
 
 def ReductionDeclareOp : OpenMP_Op<"reduction.declare", [Symbol]> {
   let summary = "declares a reduction kind";
 
   let description = [{
     Declares an OpenMP reduction kind. This requires two mandatory and one

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "mlir/Target/LLVMIR/LLVMTranslationInterface.h"
 #include "mlir/Target/LLVMIR/TypeToLLVM.h"
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
+#include "llvm/IR/DataLayout.h"
 
 namespace llvm {
 class BasicBlock;
 class IRBuilderBase;
 class Function;
 class Value;
 } // namespace llvm
 
public:
   llvm::OpenMPIRBuilder *getOpenMPBuilder() {
     if (!ompBuilder) {
       ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
       ompBuilder->initialize();
     }
     return ompBuilder.get();
   }
 
+  /// Returns the DataLayout associated with the LLVM IR module being
+  /// constructed.
+  llvm::DataLayout *getDataLayout() {
+    if (!dataLayout)
+      dataLayout = std::make_unique<llvm::DataLayout>(&*llvmModule);
+    return dataLayout.get();
+  }
+
   /// Translates the given location.
   const llvm::DILocation *translateLoc(Location loc, llvm::DILocalScope *scope);
 
   /// Translates the contents of the given block to LLVM IR using this
   /// translator. The LLVM IR basic block corresponding to the given block is
   /// expected to exist in the mapping of this translator. Uses `builder` to
   /// translate the IR, leaving it at the end of the block. If `ignoreArguments`
   /// is set, does not produce PHI nodes for the block arguments. Otherwise, the
private:
   Operation *mlirModule;
   std::unique_ptr<llvm::Module> llvmModule;
   /// A converter for translating debug information.
   std::unique_ptr<detail::DebugTranslation> debugTranslation;
 
   /// Builder for LLVM IR generation of OpenMP constructs.
   std::unique_ptr<llvm::OpenMPIRBuilder> ompBuilder;
 
+  /// Provide methods for querying the target data layout string.
+  std::unique_ptr<llvm::DataLayout> dataLayout;
+
   /// Mappings between llvm.mlir.global definitions and corresponding globals.
   DenseMap<Operation *, llvm::GlobalValue *> globalsMapping;
 
   /// A stateful object used to translate types.
   TypeToLLVMIRTranslator typeTranslator;
 
   /// A dialect interface collection used for dispatching the translation to
   /// specific dialects.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

return firstReadStmt.emitError()
            << "captured variable in omp.atomic.read must be updated in "
               "second operation";
   return success();
 }
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
 
+//===----------------------------------------------------------------------===//
+// ThreadprivateOp
+//===----------------------------------------------------------------------===//
+
+/// Parser for ThreadprivateOp
+///
+/// operation ::= `omp.threadprivate` sym_addr `->` result-type
+/// address ::= operand `:` type
+static ParseResult parseThreadprivateOp(OpAsmParser &parser,
+                                        OperationState &result) {
+  OpAsmParser::OperandType sym_addr;
+  Type sym_addrType;
+  SmallVector<Type> resultType;
+
+  if (parser.parseOperand(sym_addr) || parser.parseColonType(sym_addrType) ||
+      parser.resolveOperand(sym_addr, sym_addrType, result.operands))
+    return failure();
+  if (parser.parseArrowTypeList(resultType))
+    return failure();
+  result.addTypes(resultType);
+  return success();
+}
+
+/// Printer for ThreadprivateOp
+static void printThreadprivateOp(OpAsmPrinter &p, ThreadprivateOp op) {
+  p << " " << op.sym_addr() << " : " << op.sym_addr().getType() << " -> "
+    << op.getType();
+  return;
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

if (failed(inlineConvertOmpRegions(reductionRegion, "omp.reduction.body",
                                      builder, moduleTranslation, &phis)))
     return failure();
   assert(phis.size() == 1 && "expected one value to be yielded from "
                              "the reduction body declaration region");
   builder.CreateStore(phis[0], privateReductionVar);
   return success();
 }
 
+/// Converts an OpenMP Threadprivate operation into LLVM IR using
+/// OpenMPIRBuilder.
+static LogicalResult
+convertOmpThreadprivate(Operation &opInst, llvm::IRBuilderBase &builder,
+                        LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(
+      builder.saveIP(), builder.getCurrentDebugLocation());
+  auto threadprivateOp = cast<omp::ThreadprivateOp>(opInst);
+
+  Value symAddr = threadprivateOp.sym_addr();
+  auto symOp = symAddr.getDefiningOp();
+  assert(isa<LLVM::AddressOfOp>(symOp) && "Addressing symbol not found");
+  auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(symOp);
+
+  LLVM::GlobalOp global = addressOfOp.getGlobal();
+  llvm::GlobalValue *globalValue = moduleTranslation.lookupGlobal(global);
+  llvm::LLVMContext &ctx = builder.getContext();
+  llvm::Value *data =
+      builder.CreateBitCast(globalValue, llvm::Type::getInt8PtrTy(ctx));
+  llvm::Type *type = globalValue->getValueType();
+  llvm::TypeSize typeSize =
+      moduleTranslation.getDataLayout()->getTypeStoreSize(type);
+  llvm::ConstantInt *size = llvm::ConstantInt::get(llvm::Type::getInt64Ty(ctx),
+                                                   typeSize.getFixedSize());
+  llvm::StringRef suffix = llvm::StringRef(".cache", 6);
+  llvm::Twine cacheName = Twine(global.sym_name().str()).concat(suffix);
+  // Emit runtime function and bitcast its type (i8*) to real data type.
+  llvm::Value *callInst =
+      moduleTranslation.getOpenMPBuilder()->createCachedThreadPrivate(
+          ompLoc, data, size, cacheName);
+  llvm::Value *result = builder.CreateBitCast(callInst, globalValue->getType());
+  moduleTranslation.mapValue(opInst.getResult(0), result);
+  return success();
+}
+
 namespace {
 
 /// Implementation of the dialect interface that converts operations belonging
 /// to the OpenMP dialect to LLVM IR.
 class OpenMPDialectLLVMIRTranslationInterface
     : public LLVMTranslationDialectInterface {
 public:
   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
return llvm::TypeSwitch<Operation *, LogicalResult>(op)
         // `reduction.declare` will be used by reductions and is not
         // converted directly, skip it.
         // `critical.declare` is only used to declare names of critical
         // sections which will be used by `critical` ops and hence can be
         // ignored for lowering. The OpenMP IRBuilder will create unique
         // name for critical section names.
         return success();
       })
+      .Case([&](omp::ThreadprivateOp) {
+        return convertOmpThreadprivate(*op, builder, moduleTranslation);
+      })
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();
       });
 }
 
 void mlir::registerOpenMPDialectTranslation(DialectRegistry &registry) {
   registry.insert<omp::OpenMPDialect>();

mlir/test/Target/LLVMIR/openmp-llvm.mlir

llvm.func @single_nowait(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
   }
   // CHECK: %[[t:.*]] = mul i32 %[[x]], %[[y]]
   %t = llvm.mul %x, %y : i32
   // CHECK: store i32 %[[t]], i32* %[[zaddr]]
   llvm.store %t, %zaddr : !llvm.ptr<i32>
   // CHECK: ret void
   llvm.return
 }
+
+// -----
+
+// CHECK: @_QFsubEx = internal global i32 undef
+// CHECK: @_QFsubEx.cache = common global i8** null
+
+// CHECK-LABEL: @_QPsub
+llvm.func @_QPsub() {
+// CHECK:  [[THREAD:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB:[0-9]+]])
+// CHECK:  [[TMP1:%.*]] = call i8* @__kmpc_threadprivate_cached(%struct.ident_t* @[[GLOB]], i32 [[THREAD]], i8* bitcast (i32* @_QFsubEx to i8*), i64 4, i8*** @_QFsubEx.cache)
+// CHECK:  [[TMP2:%.*]] = bitcast i8* [[TMP1]] to i32*
+// CHECK:  store i32 1, i32* [[TMP2]], align 4
+// CHECK:  store i32 3, i32* [[TMP2]], align 4
+
+// CHECK-LABEL: omp.par.region{{.*}}
+// CHECK:  [[THREAD2:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB2:[0-9]+]])
+// CHECK:  [[TMP3:%.*]] = call i8* @__kmpc_threadprivate_cached(%struct.ident_t* @[[GLOB2]], i32 [[THREAD2]], i8* bitcast (i32* @_QFsubEx to i8*), i64 4, i8*** @_QFsubEx.cache)
+// CHECK:  [[TMP4:%.*]] = bitcast i8* [[TMP3]] to i32*
+// CHECK:  store i32 2, i32* [[TMP4]], align 4
+
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %1 = llvm.mlir.constant(2 : i32) : i32
+  %2 = llvm.mlir.constant(3 : i32) : i32
+
+  %3 = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>
+  %4 = omp.threadprivate %3 : !llvm.ptr<i32> -> !llvm.ptr<i32>
+
+  llvm.store %0, %4 : !llvm.ptr<i32>
+
+  omp.parallel  {
+    %5 = omp.threadprivate %3 : !llvm.ptr<i32> -> !llvm.ptr<i32>
+    llvm.store %1, %5 : !llvm.ptr<i32>
+    omp.terminator
+  }
+
+  llvm.store %2, %4 : !llvm.ptr<i32>
+  llvm.return
+}
+
+llvm.mlir.global internal @_QFsubEx() : i32 {
+  %0 = llvm.mlir.undef : i32
+  llvm.return %0 : i32
+}