[flang-commits] [flang] [Flang][OpenMP] Add lowering support for is_device_ptr clause (PR #168875)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Akash Banerjee (TIFitis)

<details>
<summary>Changes</summary>

Add support for OpenMP is_device_ptr clause for target directives.

---
Full diff: https://github.com/llvm/llvm-project/pull/168875.diff


2 Files Affected:

- (modified) flang/lib/Lower/OpenMP/OpenMP.cpp (+76-3) 
- (modified) flang/test/Lower/OpenMP/target.f90 (+30) 


``````````diff
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 71067283d13f7..fa2e5281f0e1c 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2485,13 +2485,15 @@ static bool isDuplicateMappedSymbol(
     const semantics::Symbol &sym,
     const llvm::SetVector<const semantics::Symbol *> &privatizedSyms,
     const llvm::SmallVectorImpl<const semantics::Symbol *> &hasDevSyms,
-    const llvm::SmallVectorImpl<const semantics::Symbol *> &mappedSyms) {
+    const llvm::SmallVectorImpl<const semantics::Symbol *> &mappedSyms,
+    const llvm::SmallVectorImpl<const semantics::Symbol *> &isDevicePtrSyms) {
   llvm::SmallVector<const semantics::Symbol *> concatSyms;
   concatSyms.reserve(privatizedSyms.size() + hasDevSyms.size() +
-                     mappedSyms.size());
+                     mappedSyms.size() + isDevicePtrSyms.size());
   concatSyms.append(privatizedSyms.begin(), privatizedSyms.end());
   concatSyms.append(hasDevSyms.begin(), hasDevSyms.end());
   concatSyms.append(mappedSyms.begin(), mappedSyms.end());
+  concatSyms.append(isDevicePtrSyms.begin(), isDevicePtrSyms.end());
 
   auto checkSymbol = [&](const semantics::Symbol &checkSym) {
     return std::any_of(concatSyms.begin(), concatSyms.end(),
@@ -2538,6 +2540,77 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
                            /*isTargetPrivitization=*/true);
   dsp.processStep1(&clauseOps);
 
+  // Ensure is_device_ptr variables are available inside the target region by
+  // passing them as block arguments via the has_device_addr list, and record
+  // the associated var_ptr values so the clause operands reference the same
+  // MLIR SSA values used by the map info operations.
+  if (!isDevicePtrSyms.empty()) {
+    fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+    llvm::DenseMap<const semantics::Symbol *, mlir::Value> symToMapInfo;
+    for (auto [symPtr, mapVal] : llvm::zip(hasDeviceAddrSyms, clauseOps.hasDeviceAddrVars))
+      if (symPtr)
+        symToMapInfo.try_emplace(symPtr, mapVal);
+
+    llvm::SmallVector<mlir::Value> remappedIsDevicePtrVars;
+    remappedIsDevicePtrVars.reserve(isDevicePtrSyms.size());
+
+    for (const semantics::Symbol *sym : isDevicePtrSyms) {
+      if (!sym)
+        continue;
+
+      mlir::omp::MapInfoOp mapInfoOp;
+      if (auto it = symToMapInfo.find(sym); it != symToMapInfo.end())
+        mapInfoOp = mlir::dyn_cast_or_null<mlir::omp::MapInfoOp>(
+            it->second.getDefiningOp());
+
+      if (!mapInfoOp && !llvm::is_contained(mapSyms, sym)) {
+        fir::ExtendedValue dataExv = converter.getSymbolExtendedValue(*sym);
+
+        fir::factory::AddrAndBoundsInfo info =
+            Fortran::lower::getDataOperandBaseAddr(
+                converter, builder, sym->GetUltimate(),
+                converter.getCurrentLocation());
+
+        llvm::SmallVector<mlir::Value> bounds =
+            fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
+                                               mlir::omp::MapBoundsType>(
+                builder, info, dataExv,
+                semantics::IsAssumedSizeArray(sym->GetUltimate()),
+                converter.getCurrentLocation());
+
+        mlir::Value baseOp = info.rawInput;
+        if (auto refType =
+                mlir::dyn_cast<fir::ReferenceType>(baseOp.getType()))
+          (void)refType; // element type not needed for storage-only mapping
+
+        mlir::Value mapVal = createMapInfoOp(
+            builder, converter.getCurrentLocation(), baseOp,
+            /*varPtrPtr=*/mlir::Value{}, sym->name().ToString(), bounds,
+            /*members=*/{}, /*membersIndex=*/mlir::ArrayAttr{},
+            mlir::omp::ClauseMapFlags::storage,
+            mlir::omp::VariableCaptureKind::ByRef, baseOp.getType(),
+            /*partialMap=*/false, /*mapperId=*/mlir::FlatSymbolRefAttr{});
+
+        clauseOps.hasDeviceAddrVars.push_back(mapVal);
+        hasDeviceAddrSyms.push_back(sym);
+        symToMapInfo[sym] = mapVal;
+        mapInfoOp = mapVal.getDefiningOp<mlir::omp::MapInfoOp>();
+      }
+
+      if (mapInfoOp) {
+        builder.setInsertionPointAfter(mapInfoOp);
+        auto clonedOp = builder.clone(*mapInfoOp.getOperation());
+        auto clonedMapInfo = mlir::dyn_cast<mlir::omp::MapInfoOp>(clonedOp);
+        assert(clonedMapInfo && "expected cloned map info op");
+        remappedIsDevicePtrVars.push_back(clonedMapInfo.getResult());
+      }
+    }
+
+    if (!remappedIsDevicePtrVars.empty())
+      clauseOps.isDevicePtrVars.assign(remappedIsDevicePtrVars.begin(),
+                                       remappedIsDevicePtrVars.end());
+  }
+
   // 5.8.1 Implicit Data-Mapping Attribute Rules
   // The following code follows the implicit data-mapping rules to map all the
   // symbols used inside the region that do not have explicit data-environment
@@ -2570,7 +2643,7 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
       return;
 
     if (!isDuplicateMappedSymbol(sym, dsp.getAllSymbolsToPrivatize(),
-                                 hasDeviceAddrSyms, mapSyms)) {
+                                 hasDeviceAddrSyms, mapSyms, isDevicePtrSyms)) {
       if (const auto *details =
               sym.template detailsIf<semantics::HostAssocDetails>())
         converter.copySymbolBinding(details->symbol(), sym);
diff --git a/flang/test/Lower/OpenMP/target.f90 b/flang/test/Lower/OpenMP/target.f90
index 26bd62edf9d0c..1e62adc0f6f98 100644
--- a/flang/test/Lower/OpenMP/target.f90
+++ b/flang/test/Lower/OpenMP/target.f90
@@ -566,6 +566,36 @@ subroutine omp_target_device_addr
 end subroutine omp_target_device_addr
 
 
+!===============================================================================
+! Target `is_device_ptr` clause
+!===============================================================================
+
+!CHECK-LABEL: func.func @_QPomp_target_is_device_ptr() {
+subroutine omp_target_is_device_ptr
+   use iso_c_binding, only: c_associated, c_ptr
+   implicit none
+   integer :: i
+   integer :: arr(4)
+   type(c_ptr) :: p
+
+   i = 0
+   arr = 0
+
+   !CHECK: %[[P_STORAGE:.*]] = omp.map.info {{.*}}{name = "p"}
+   !CHECK: %[[P_IS:.*]] = omp.map.info {{.*}}{name = "p"}
+   !CHECK: %[[ARR_MAP:.*]] = omp.map.info {{.*}}{name = "arr"}
+   !CHECK: omp.target is_device_ptr(%[[P_IS]] :
+   !CHECK-SAME: has_device_addr(%[[P_STORAGE]] ->
+   !CHECK-SAME: map_entries({{.*}}%[[ARR_MAP]] ->
+   !$omp target is_device_ptr(p)
+      if (c_associated(p)) i = i + 1
+      arr(1) = i
+   !$omp end target
+   !CHECK: omp.terminator
+   !CHECK: }
+end subroutine omp_target_is_device_ptr
+
+
 !===============================================================================
 ! Target Data with unstructured code
 !===============================================================================

``````````

</details>


https://github.com/llvm/llvm-project/pull/168875