[OpenMP][mlir] Add Groupprivate op in omp dialect.

Author: skc7

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

@@ -2224,4 +2224,31 @@ def WorkdistributeOp : OpenMP_Op<"workdistribute"> {
   let assemblyFormat = "$region attr-dict";
 }
 
+//===----------------------------------------------------------------------===//
+// [6.0] groupprivate Directive
+//===----------------------------------------------------------------------===//
+
+def GroupprivateOp : OpenMP_Op<"groupprivate",
+                      [AllTypesMatch<["sym_addr", "gp_addr"]>]> {
+  let summary = "groupprivate directive";
+  let description = [{
+    The groupprivate directive specifies that variables are replicated, with
+    each group having its own copy.
+
+    This operation takes in the address of a symbol that represents the original
+    variable and returns the address of its groupprivate copy. All occurrences of
+    groupprivate variables in a parallel region should use the groupprivate copy
+    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:$gp_addr);
+  let assemblyFormat = [{
+    $sym_addr `:` type($sym_addr) `->` type($gp_addr) attr-dict
+  }];
+}
+
 #endif // OPENMP_OPS

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

@@ -6092,6 +6092,40 @@ convertTargetFreeMemOp(Operation &opInst, llvm::IRBuilderBase &builder,
   return success();
 }
 
+/// Converts an OpenMP Groupprivate operation into LLVM IR.
+static LogicalResult
+convertOmpGroupprivate(Operation &opInst, llvm::IRBuilderBase &builder,
+                        LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+  auto groupprivateOp = cast<omp::GroupprivateOp>(opInst);
+
+  if (failed(checkImplementationStatus(opInst)))
+    return failure();
+
+  Value symAddr = groupprivateOp.getSymAddr();
+  auto *symOp = symAddr.getDefiningOp();
+
+  if (auto asCast = dyn_cast<LLVM::AddrSpaceCastOp>(symOp))
+    symOp = asCast.getOperand().getDefiningOp();
+
+  if (!isa<LLVM::AddressOfOp>(symOp))
+    return opInst.emitError("Addressing symbol not found");
+  LLVM::AddressOfOp addressOfOp = dyn_cast<LLVM::AddressOfOp>(symOp);
+
+  LLVM::GlobalOp global =
+      addressOfOp.getGlobal(moduleTranslation.symbolTable());
+  llvm::GlobalValue *globalValue = moduleTranslation.lookupGlobal(global);
+
+  if (!ompBuilder->Config.isTargetDevice()) {
+    llvm_unreachable("NYI");
+  } else {
+    moduleTranslation.mapValue(opInst.getResult(0), globalValue);
+  }
+
+  return success();
+}
+
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR (including
 /// OpenMP runtime calls).
 static LogicalResult
@@ -6275,6 +6309,9 @@ convertHostOrTargetOperation(Operation *op, llvm::IRBuilderBase &builder,
           .Case([&](omp::TargetFreeMemOp) {
             return convertTargetFreeMemOp(*op, builder, moduleTranslation);
           })
+          .Case([&](omp::GroupprivateOp) {
+            return convertOmpGroupprivate(*op, builder, moduleTranslation);
+          })
           .Default([&](Operation *inst) {
             return inst->emitError()
                    << "not yet implemented: " << inst->getName();

mlir/test/Dialect/OpenMP/ops.mlir

@@ -3321,3 +3321,22 @@ func.func @omp_workdistribute() {
   }
   return
 }
+
+llvm.mlir.global internal @_QFgpEx() : i32
+
+func.func @omp_groupprivate() {
+  %0 = arith.constant 1 : i32
+  %1 = arith.constant 2 : i32
+  // CHECK: [[ARG0:%.*]] = llvm.mlir.addressof @_QFgpEx : !llvm.ptr
+  %global_addr = llvm.mlir.addressof @_QFgpEx : !llvm.ptr
+  omp.teams {
+    // CHECK: {{.*}} = omp.groupprivate [[ARG0]] : !llvm.ptr -> !llvm.ptr
+    %group_private_addr_in_teams = omp.groupprivate %global_addr : !llvm.ptr -> !llvm.ptr
+    llvm.store %0, %group_private_addr_in_teams : i32, !llvm.ptr
+    omp.terminator
+  }
+  // CHECK: {{.*}} = omp.groupprivate [[ARG0]] : !llvm.ptr -> !llvm.ptr
+  %group_private_addr_after_teams = omp.groupprivate %global_addr : !llvm.ptr -> !llvm.ptr
+  llvm.store %1, %group_private_addr_after_teams : i32, !llvm.ptr
+  return
+}