[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.
+
+    The current implementation uses the OpenMP runtime to provide thread-local
+    storage (TLS). Using the TLS feature of the LLVM IR will be supported in
+    future.
+
+    This operation takes in the address of a symbol that represents the original
+    variable and returns the address of its TLS. All occurrences of
+    threadprivate variables in a parallel region should use the TLS returned by
+    this operation.
+
+    The `sym_addr` refers to the address of the symbol, which is a pointer to
+    the original variable.
+  }];
+
+  let arguments = (ins OpenMP_PointerLikeType:$sym_addr);
+  let results = (outs OpenMP_PointerLikeType:$tls_addr);
+  let assemblyFormat = [{
+    $sym_addr `:` type($sym_addr) `->` type($tls_addr) attr-dict
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // 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/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);
+  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::Value *data =
+      builder.CreateBitCast(globalValue, builder.getInt8PtrTy());
+  llvm::Type *type = globalValue->getValueType();
+  llvm::TypeSize typeSize =
+      builder.GetInsertBlock()->getModule()->getDataLayout().getTypeStoreSize(
+          type);
+  llvm::ConstantInt *size = builder.getInt64(typeSize.getFixedSize());
+  llvm::StringRef suffix = llvm::StringRef(".cache", 6);
+  llvm::Twine cacheName = Twine(global.getSymName()).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/Dialect/OpenMP/ops.mlir

-// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// RUN: mlir-opt -split-input-file %s | mlir-opt | FileCheck %s
 
 func @omp_barrier() -> () {
   // CHECK: omp.barrier
   omp.barrier
   return
 }
 
 func @omp_master() -> () {
omp.single allocate(%data_var : memref<i32> -> %data_var : memref<i32>) nowait {
       // CHECK: omp.terminator
       omp.terminator
     }
     // CHECK: omp.terminator
     omp.terminator
   }
   return
 }
+
+// -----
+
+func @omp_threadprivate() {
+  %0 = arith.constant 1 : i32
+  %1 = arith.constant 2 : i32
+  %2 = arith.constant 3 : i32
+
+  // CHECK: [[ARG0:%.*]] = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>
+  // CHECK: {{.*}} = omp.threadprivate [[ARG0]] : !llvm.ptr<i32> -> !llvm.ptr<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>
+
+  // CHECK:  omp.parallel
+  // CHECK:    {{.*}} = omp.threadprivate [[ARG0]] : !llvm.ptr<i32> -> !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>
+  return
+}
+
+llvm.mlir.global internal @_QFsubEx() : i32

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: @omp_threadprivate
+llvm.func @omp_threadprivate() {
+// 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