add device_type attr to groupprivate

Author: skc7

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

@@ -2244,10 +2244,13 @@ def GroupprivateOp : OpenMP_Op<"groupprivate",
     the original variable.
   }];
 
-  let arguments = (ins OpenMP_PointerLikeType:$sym_addr);
+  let arguments = (ins
+    OpenMP_PointerLikeType:$sym_addr,
+    OptionalAttr<DeclareTargetDeviceTypeAttr>:$device_type
+  );
   let results = (outs OpenMP_PointerLikeType:$gp_addr);
   let assemblyFormat = [{
-    $sym_addr `:` type($sym_addr) `->` type($gp_addr) attr-dict
+    $sym_addr `:` type($sym_addr) ( `,` `device_type` $device_type^ )? `->` type($gp_addr) attr-dict
   }];
 }
 

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

@@ -6137,6 +6137,28 @@ convertOmpGroupprivate(Operation &opInst, llvm::IRBuilderBase &builder,
 
   if (failed(checkImplementationStatus(opInst)))
     return failure();
+  
+  bool isTargetDevice = ompBuilder->Config.isTargetDevice();
+  auto deviceType = groupprivateOp.getDeviceType();
+  
+  // skip allocation based on device_type
+  bool shouldAllocate = true;
+  if (deviceType.has_value()) {    
+    switch (*deviceType) {
+    case mlir::omp::DeclareTargetDeviceType::host:
+      // Only allocate on host
+      shouldAllocate = !isTargetDevice;
+      break;
+    case mlir::omp::DeclareTargetDeviceType::nohost:
+      // Only allocate on device
+      shouldAllocate = isTargetDevice;
+      break;
+    case mlir::omp::DeclareTargetDeviceType::any:
+      // Allocate on both
+      shouldAllocate = true;
+      break;
+    }
+  }
 
   Value symAddr = groupprivateOp.getSymAddr();
   auto *symOp = symAddr.getDefiningOp();
@@ -6151,21 +6173,32 @@ convertOmpGroupprivate(Operation &opInst, llvm::IRBuilderBase &builder,
   LLVM::GlobalOp global =
       addressOfOp.getGlobal(moduleTranslation.symbolTable());
   llvm::GlobalValue *globalValue = moduleTranslation.lookupGlobal(global);
+  llvm::Value *resultPtr;
 
-  // Get the size of the variable
-  llvm::Type *varType = globalValue->getValueType();
-  llvm::Module *llvmModule = moduleTranslation.getLLVMModule();
-  llvm::DataLayout DL = llvmModule->getDataLayout();
-  uint64_t typeSize = DL.getTypeAllocSize(varType);
-  // Call omp_alloc_shared to allocate memory for groupprivate variable.
-  llvm::FunctionCallee allocSharedFn = ompBuilder->getOrCreateRuntimeFunction(
-      *llvmModule, llvm::omp::OMPRTL___kmpc_alloc_shared);
-  // Call runtime to allocate shared memory for this group
-  llvm::Value *groupPrivatePtr =
-      builder.CreateCall(allocSharedFn, {builder.getInt64(typeSize)});
-  groupPrivatePtr =
-      builder.CreateBitCast(groupPrivatePtr, globalValue->getType());
-  moduleTranslation.mapValue(opInst.getResult(0), groupPrivatePtr);
+  if (shouldAllocate) {
+    // Get the size of the variable
+    llvm::Type *varType = globalValue->getValueType();
+    llvm::Module *llvmModule = moduleTranslation.getLLVMModule();
+    llvm::DataLayout DL = llvmModule->getDataLayout();
+    uint64_t typeSize = DL.getTypeAllocSize(varType);
+    // Call omp_alloc_shared to allocate memory for groupprivate variable.
+    llvm::FunctionCallee allocSharedFn = ompBuilder->getOrCreateRuntimeFunction(
+        *llvmModule, llvm::omp::OMPRTL___kmpc_alloc_shared);
+    // Call runtime to allocate shared memory for this group
+    llvm::Value *groupPrivatePtr =
+        builder.CreateCall(allocSharedFn, {builder.getInt64(typeSize)});
+    resultPtr =
+        builder.CreateBitCast(groupPrivatePtr, globalValue->getType());
+  }
+  else {
+    // Use original global address when not allocating group-private storage
+    resultPtr = moduleTranslation.lookupValue(symAddr);
+    if (!resultPtr) {
+      // Fallback: create address-of for the global
+      resultPtr = builder.CreateBitCast(globalValue, globalValue->getType());
+    }
+  }
+  moduleTranslation.mapValue(opInst.getResult(0), resultPtr);
   return success();
 }
 

mlir/test/Dialect/OpenMP/ops.mlir

@@ -3368,22 +3368,27 @@ func.func @omp_target_map_clause_type_test(%arg0 : memref<?xi32>) -> () {
     return
 }
 
-
-// CHECK-LABEL: func.func @omp_groupprivate
-llvm.mlir.global internal @_QFgpEx() : i32
-func.func @omp_groupprivate() {
+// CHECK-LABEL: func.func @omp_groupprivate_device_type
+func.func @omp_groupprivate_device_type() {
   %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
+  %group_private_addr = omp.groupprivate %global_addr : !llvm.ptr -> !llvm.ptr
+
+  // CHECK: {{.*}} = omp.groupprivate [[ARG0]] : !llvm.ptr, device_type (any) -> !llvm.ptr
+  %group_private_any = omp.groupprivate %global_addr : !llvm.ptr, device_type(any) -> !llvm.ptr
+  llvm.store %1, %group_private_any : i32, !llvm.ptr
+
+  // CHECK: {{.*}} = omp.groupprivate [[ARG0]] : !llvm.ptr, device_type (host) -> !llvm.ptr
+  %group_private_host = omp.groupprivate %global_addr : !llvm.ptr, device_type(host) -> !llvm.ptr
+  llvm.store %1, %group_private_host : i32, !llvm.ptr
+
+  // CHECK: {{.*}} = omp.groupprivate [[ARG0]] : !llvm.ptr, device_type (nohost) -> !llvm.ptr
+  %group_private_nohost = omp.groupprivate %global_addr : !llvm.ptr, device_type(nohost) -> !llvm.ptr
+  llvm.store %1, %group_private_nohost : i32, !llvm.ptr
+
   return
 }