[flang-commits] [flang] [llvm] [Flang][OpenMP] Improve use_device_addr code generation (PR #221265)

Dominik Adamski via flang-commits flang-commits at lists.llvm.org
Thu Sep 17 08:06:32 PDT 2026


https://github.com/DominikAdamski updated https://github.com/llvm/llvm-project/pull/221265

>From e8fd8f5fcabc9a093f21728c65ec58acbea625aa Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Fri, 4 Sep 2026 10:37:23 -0500
Subject: [PATCH 1/5] [Flang][OpenMP] Improve use_device_addr code generation

Currently, the use_device_addr implementation reuses the standard
mapping mechanism, which is suboptimal for code like:

SUBROUTINE device_addr_func(x)
    INTEGER, TARGET, INTENT(IN)    :: x(:)
    !$omp target data use_device_addr (x)

For such code, Flang maps the temporary descriptor for the x array
to the GPU. This unnecessary mapping is time-consuming
and can be a large bottleneck for Fortran-to-C function wrappers
that use use_device_addr to pass a C pointer for offload code.

For use_device_addr, we only need to update the base address
in the descriptor that is used inside use_device_addr. This is cheaper
than mapping the whole descriptor to the GPU.

Scope of changes:

1) Modified mapping for the use_device_addr clause. Map only the descriptor
   and do not check whether it is present on the GPU.
2) Generated a copy of the host descriptor with an updated base address
   for the use_device_addr code region.

Signed-off-by: Dominik Adamski <dominik.adamski at amd.com>
---
 .../Optimizer/OpenMP/MapInfoFinalization.cpp  | 200 +++++++++++++++---
 .../OpenMP/use-device-addr-performance.f90    |  43 ++++
 offload/include/omptarget.h                   |   4 +
 offload/libomptarget/exports                  |   1 +
 offload/libomptarget/interface.cpp            |   7 +
 5 files changed, 222 insertions(+), 33 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/use-device-addr-performance.f90

diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index e78d194c5c7ae..738064355616a 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -888,14 +888,20 @@ class MapInfoFinalizationPass
     return false;
   }
 
-  bool isUseDeviceAddr(mlir::omp::MapInfoOp mapOp, mlir::Operation &userOp) {
+  mlir::BlockArgument getUseDeviceAddrBlockArg(mlir::omp::MapInfoOp mapOp,
+                                               mlir::Operation &userOp) {
     if (auto targetDataOp = llvm::dyn_cast<mlir::omp::TargetDataOp>(userOp)) {
-      for (mlir::Value uda : targetDataOp.getUseDeviceAddrVars()) {
-        if (uda.getDefiningOp() == mapOp)
-          return true;
+      auto iface = mlir::cast<mlir::omp::BlockArgOpenMPOpInterface>(
+          targetDataOp.getOperation());
+      auto useDeviceAddrArgs = iface.getUseDeviceAddrBlockArgs();
+      auto useDeviceAddrVars = targetDataOp.getUseDeviceAddrVars();
+      for (auto [useDeviceAddrVar, useDeviceAddrArg] :
+           llvm::zip_equal(useDeviceAddrVars, useDeviceAddrArgs)) {
+        if (useDeviceAddrVar.getDefiningOp() == mapOp)
+          return mlir::cast<mlir::BlockArgument>(useDeviceAddrArg);
       }
     }
-    return false;
+    return nullptr;
   }
 
   bool isUseDevicePtr(mlir::omp::MapInfoOp mapOp, mlir::Operation &userOp) {
@@ -1033,10 +1039,11 @@ class MapInfoFinalizationPass
   /// additional attach map which indicates to the runtime to try and attach
   /// the base address to the descriptor if it's available and it's the first
   /// time the ref_ptr has been allocated on the device.
-  void genRefPtrMap(mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
-                    mlir::Operation *target, mlir::Value descriptor,
-                    llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
-                    bool isAttachNever, bool isAttachAlways) {
+  mlir::omp::MapInfoOp
+  genRefPtrMap(mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
+               mlir::Operation *target, mlir::Value descriptor,
+               llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
+               bool isAttachNever, bool isAttachAlways) {
     auto newMapInfoOp = mlir::omp::MapInfoOp::create(
         builder, op->getLoc(), op.getResult().getType(), descriptor,
         mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
@@ -1053,6 +1060,7 @@ class MapInfoFinalizationPass
                            mlir::omp::ClauseMapFlags::ref_ptr, isAttachAlways);
     op.replaceAllUsesWith(newMapInfoOp.getResult());
     op->erase();
+    return newMapInfoOp;
   }
 
   /// Helper function to generate a ref_ptee map. This handles the case where
@@ -1064,11 +1072,12 @@ class MapInfoFinalizationPass
   /// additional attach map which indicates to the runtime to try and attach
   /// the base address to the descriptor if it's available and it's the first
   /// time the ref_ptee has been allocated on the device.
-  void genRefPteeMap(mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
-                     mlir::Operation *target, mlir::Value descriptor,
-                     llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
-                     bool isAttachNever, bool isAttachAlways,
-                     mlir::FlatSymbolRefAttr mapperId) {
+  mlir::omp::MapInfoOp
+  genRefPteeMap(mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
+                mlir::Operation *target, mlir::Value descriptor,
+                llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
+                bool isAttachNever, bool isAttachAlways,
+                mlir::FlatSymbolRefAttr mapperId) {
     // NOTE: We replace the descriptor map with the base address map. This
     // effectively replaces the descriptor's index position in any complex
     // structure mapping. This is a little different to the
@@ -1086,6 +1095,7 @@ class MapInfoFinalizationPass
                            newMapInfoOp.getVarPtrPtr());
     op.replaceAllUsesWith(newMapInfoOp.getResult());
     op->erase();
+    return newMapInfoOp;
   }
 
   /// Helper function to generate a ref_ptr_ptee or default descriptor map.
@@ -1096,13 +1106,13 @@ class MapInfoFinalizationPass
   /// a map is generated for the descriptor and its base address,
   /// similarly in the default auto attach case, we generate an additional
   /// attach map.
-  void genRefPtrPteeOrDefaultMap(
+  mlir::omp::MapInfoOp genRefPtrPteeOrDefaultMap(
       mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
       mlir::Operation *target, mlir::Value descriptor,
       llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
       bool isAttachNever, bool isAttachAlways, bool isHasDeviceAddrFlag,
       bool descCanBeDeferred, bool canOptimizeDescViaPrivatization,
-      mlir::FlatSymbolRefAttr mapperId) {
+      mlir::FlatSymbolRefAttr mapperId, bool mapOnlyDescriptor) {
     bool isRefPtrPtee =
         bitEnumContainsAll(op.getMapType(),
                            mlir::omp::ClauseMapFlags::ref_ptr) &&
@@ -1158,12 +1168,13 @@ class MapInfoFinalizationPass
         /*partial_map=*/builder.getBoolAttr(false));
 
     mlir::Operation *attachMap = nullptr;
-    if (!isAttachNever && !isHasDeviceAddrFlag)
+    if (!isAttachNever && !isHasDeviceAddrFlag && !mapOnlyDescriptor) {
       attachMap =
           genImplicitAttachMap(op, descriptor, mapMemberUsers, target, builder,
                                mlir::omp::ClauseMapFlags::ref_ptr |
                                    mlir::omp::ClauseMapFlags::ref_ptee,
                                isAttachAlways, baseAddr.getVarPtrPtr());
+    }
 
     op.replaceAllUsesWith(newMapInfoOp.getResult());
     op->erase();
@@ -1171,8 +1182,10 @@ class MapInfoFinalizationPass
     // The deferral only applies to cases where we map both the descriptor and
     // base address at once, and when provided ref_ptr_ptee by a user we
     // assume they know what they're asking for and don't intervene.
-    if (descCanBeDeferred && !isRefPtrPtee)
+    if (descCanBeDeferred && !isRefPtrPtee &&
+        !getUseDeviceAddrBlockArg(op, *target))
       deferrableDesc.push_back(std::make_pair(newMapInfoOp, attachMap));
+    return newMapInfoOp;
   }
 
   // This function handles the splitting of allocatable/pointer maps in
@@ -1185,15 +1198,17 @@ class MapInfoFinalizationPass
   // - genRefPteeMap: for ref_ptee mappings
   // - genRefPtrPteeOrDefaultMap: for ref_ptr_ptee or default descriptor
   // mappings
-  void genDescriptorMaps(mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
-                         mlir::Operation *target) {
+  mlir::omp::MapInfoOp genDescriptorMaps(mlir::omp::MapInfoOp op,
+                                         fir::FirOpBuilder &builder,
+                                         mlir::Operation *target,
+                                         bool &canOptimizeUseDeviceAddr) {
     bool descCanBeDeferred = false;
     bool canOptimizeDescViaPrivatization = false;
     llvm::SmallVector<ParentAndPlacement> mapMemberUsers;
     getMemberUserList(op, mapMemberUsers);
-
     // TODO: map the addendum segment of the descriptor, similarly to the
     // base address/data pointer member.
+    bool mapOnlyDescriptor = false;
     bool isHasDeviceAddrFlag = isHasDeviceAddr(op, *target);
     bool isAttachNever = bitEnumContainsAll(
         op.getMapType(), mlir::omp::ClauseMapFlags::attach_never);
@@ -1210,11 +1225,24 @@ class MapInfoFinalizationPass
 
     mlir::Value descriptor = getDescriptorFromBoxMap(
         op, builder, descCanBeDeferred, canOptimizeDescViaPrivatization);
+    bool isNewDescriptor = mlir::isa<fir::AllocaOp>(descriptor.getDefiningOp());
+    bool isUseDeviceAddrItem =
+        (getUseDeviceAddrBlockArg(op, *target) != nullptr);
+    bool isArray = false;
+    bool knownRanks = false;
+    fir::BaseBoxType bt = mlir::dyn_cast<fir::BaseBoxType>(
+        fir::unwrapRefType(descriptor.getType()));
+    if (bt) {
+      isArray = bt.isArray();
+      knownRanks = !bt.isAssumedRank();
+    }
+    canOptimizeUseDeviceAddr =
+        (isNewDescriptor && isUseDeviceAddrItem && isArray && knownRanks);
+    mapOnlyDescriptor = isHasDeviceAddrFlag | canOptimizeUseDeviceAddr;
     mlir::FlatSymbolRefAttr mapperId = op.getMapperIdAttr();
-
     // Exclude irregular maps from optimization via privatization; at least for
     // the moment.
-    if (isHasDeviceAddrFlag || isUseDeviceAddr(op, *target) ||
+    if (isHasDeviceAddrFlag || getUseDeviceAddrBlockArg(op, *target) ||
         isUseDevicePtr(op, *target))
       canOptimizeDescViaPrivatization = false;
 
@@ -1234,18 +1262,21 @@ class MapInfoFinalizationPass
     // TODO: This currently only works for the first level of a
     // derived-type descriptor chain and will likely need to be extended for the
     // case where we do a similar style of mapping for deeper nestings.
+    mlir::omp::MapInfoOp newMapInfo;
     if (isRefPtr && op.getMembers().empty()) {
-      genRefPtrMap(op, builder, target, descriptor, mapMemberUsers,
-                   isAttachNever, isAttachAlways);
+      newMapInfo = genRefPtrMap(op, builder, target, descriptor, mapMemberUsers,
+                                isAttachNever, isAttachAlways);
     } else if (isRefPtee) {
-      genRefPteeMap(op, builder, target, descriptor, mapMemberUsers,
-                    isAttachNever, isAttachAlways, mapperId);
+      newMapInfo =
+          genRefPteeMap(op, builder, target, descriptor, mapMemberUsers,
+                        isAttachNever, isAttachAlways, mapperId);
     } else {
-      genRefPtrPteeOrDefaultMap(op, builder, target, descriptor, mapMemberUsers,
-                                isAttachNever, isAttachAlways,
-                                isHasDeviceAddrFlag, descCanBeDeferred,
-                                canOptimizeDescViaPrivatization, mapperId);
+      newMapInfo = genRefPtrPteeOrDefaultMap(
+          op, builder, target, descriptor, mapMemberUsers, isAttachNever,
+          isAttachAlways, mapOnlyDescriptor, descCanBeDeferred,
+          canOptimizeDescViaPrivatization, mapperId, canOptimizeUseDeviceAddr);
     }
+    return newMapInfo;
   }
 
   void addImplicitDescriptorMapToTargetDataOp(mlir::omp::MapInfoOp op,
@@ -1283,7 +1314,7 @@ class MapInfoFinalizationPass
     if (!llvm::isa<mlir::omp::TargetDataOp>(target) || op.getMembers().empty())
       return;
 
-    if (!isUseDeviceAddr(op, target) && !isUseDevicePtr(op, target))
+    if (!getUseDeviceAddrBlockArg(op, target) && !isUseDevicePtr(op, target))
       return;
 
     auto targetDataOp = llvm::cast<mlir::omp::TargetDataOp>(target);
@@ -1404,6 +1435,101 @@ class MapInfoFinalizationPass
     return false;
   }
 
+  mlir::Value genTgtGetMappedPtrCall(fir::FirOpBuilder &builder,
+                                     mlir::Location loc, mlir::Value deviceNum,
+                                     mlir::Value hostPtr,
+                                     mlir::ModuleOp module) {
+    auto *context = builder.getContext();
+    auto voidPtrType = fir::LLVMPointerType::get(context, builder.getI8Type());
+    auto i32Type = builder.getI32Type();
+    auto i64Type = builder.getI64Type();
+    auto funcName = "__tgt_get_mapped_ptr";
+    auto funcOp = module.lookupSymbol<mlir::func::FuncOp>(funcName);
+
+    if (!funcOp) {
+      auto funcType = mlir::FunctionType::get(context, {i64Type, voidPtrType},
+                                              {voidPtrType});
+
+      mlir::OpBuilder::InsertionGuard guard(builder);
+      builder.setInsertionPointToStart(module.getBody());
+
+      funcOp = mlir::func::FuncOp::create(builder, loc, funcName, funcType);
+      funcOp.setPrivate();
+    }
+    if (!deviceNum) {
+      auto funcGetDefaultDeviceName = "omp_get_default_device";
+      auto funcGetDefaultDeviceOp =
+          module.lookupSymbol<mlir::func::FuncOp>(funcGetDefaultDeviceName);
+      if (!funcGetDefaultDeviceOp) {
+        auto funcType = mlir::FunctionType::get(context, {}, {i32Type});
+
+        mlir::OpBuilder::InsertionGuard guard(builder);
+        builder.setInsertionPointToStart(module.getBody());
+
+        funcGetDefaultDeviceOp = mlir::func::FuncOp::create(
+            builder, loc, funcGetDefaultDeviceName, funcType);
+        funcGetDefaultDeviceOp.setPrivate();
+      }
+      auto callGetDefaultDeviceOp =
+          fir::CallOp::create(builder, loc, funcGetDefaultDeviceOp, {});
+      deviceNum = callGetDefaultDeviceOp.getResult(0);
+    }
+    llvm::SmallVector<mlir::Value> args;
+    args.push_back(fir::ConvertOp::create(builder, loc, i64Type, deviceNum));
+    args.push_back(fir::ConvertOp::create(builder, loc, voidPtrType, hostPtr));
+    auto callOp = fir::CallOp::create(builder, loc, funcOp, args);
+    return callOp.getResult(0);
+  }
+
+  void genOptimizedUseDeviceAddr(fir::FirOpBuilder &builder,
+                                 mlir::omp::TargetDataOp targetDataOp,
+                                 mlir::omp::MapInfoOp mapOp,
+                                 mlir::ModuleOp module) {
+    mlir::Location loc = targetDataOp.getLoc();
+    mapOp.setMapType(mapOp.getMapType() | mlir::omp::ClauseMapFlags::literal);
+    auto arg = getUseDeviceAddrBlockArg(mapOp, *targetDataOp.getOperation());
+    auto insertionPoint = builder.saveInsertionPoint();
+    builder.setInsertionPoint(&targetDataOp->getRegion(0).front().front());
+    // We need to create a temporary copy of the host descriptor, which will
+    // be used inside the use_device_addr code region. The copy will be updated
+    // with the target pointer of the mapped array. The lifetime of the
+    // temporary copy is equal to the scope of the use_device_addr.
+    // The additional copy eliminates the need to synchronize the host
+    // descriptor if we want to update the host descriptor inside the
+    // use_device_addr.
+    auto allocaTgtDescriptor =
+        fir::AllocaOp::create(builder, loc, arg.getType());
+    auto allocaHostDescriptor =
+        fir::AllocaOp::create(builder, loc, arg.getType());
+    fir::StoreOp::create(builder, loc, arg, allocaHostDescriptor);
+    auto hostDescriptor =
+        fir::LoadOp::create(builder, loc, allocaHostDescriptor);
+    auto hostAddrPtr = fir::BoxAddrOp::create(builder, loc, hostDescriptor);
+    auto convertedAddr = fir::ConvertOp::create(
+        builder, loc,
+        fir::LLVMPointerType::get(builder.getContext(), builder.getI8Type()),
+        hostAddrPtr);
+    auto newAddr = genTgtGetMappedPtrCall(
+        builder, loc, targetDataOp.getDevice(), convertedAddr, module);
+    auto convertedGPUAddr =
+        fir::ConvertOp::create(builder, loc, hostAddrPtr.getType(), newAddr);
+    llvm::SmallVector<mlir::Value> lbounds;
+    llvm::SmallVector<mlir::Value> extents;
+    llvm::SmallVector<mlir::Value> strides;
+    fir::factory::genDimInfoFromBox(builder, loc, hostDescriptor, &lbounds,
+                                    &extents, &strides);
+    auto newDescriptor =
+        fir::CreateBoxOp::create(builder, loc, hostDescriptor.getType(),
+                                 convertedGPUAddr, lbounds, extents, strides);
+    fir::StoreOp::create(builder, loc, newDescriptor, allocaTgtDescriptor);
+    auto res = fir::LoadOp::create(builder, loc, allocaTgtDescriptor);
+    arg.replaceUsesWithIf(res, [&](mlir::OpOperand &use) {
+      mlir::Operation *user = use.getOwner();
+      return res->isBeforeInBlock(user);
+    });
+    builder.restoreInsertionPoint(insertionPoint);
+  }
+
   // This pass executes on omp::MapInfoOp's containing descriptor based types
   // (allocatables, pointers, assumed shape etc.) and expanding them into
   // multiple omp::MapInfoOp's for each pointer member contained within the
@@ -1659,7 +1785,15 @@ class MapInfoFinalizationPass
           builder.setInsertionPoint(op);
           mlir::Operation *targetUser = getFirstTargetUser(op);
           assert(targetUser && "expected user of map operation was not found");
-          genDescriptorMaps(op, builder, targetUser);
+          auto targetDataOp =
+              llvm::dyn_cast<mlir::omp::TargetDataOp>(*targetUser);
+          bool canOptimizeUseDeviceAddr = false;
+          mlir::omp::MapInfoOp newMapInfo = genDescriptorMaps(
+              op, builder, targetUser, canOptimizeUseDeviceAddr);
+          if (canOptimizeUseDeviceAddr && targetDataOp) {
+            genOptimizedUseDeviceAddr(builder, targetDataOp, newMapInfo,
+                                      module);
+          }
         }
       });
 
diff --git a/flang/test/Lower/OpenMP/use-device-addr-performance.f90 b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
new file mode 100644
index 0000000000000..c4785e7e8a097
--- /dev/null
+++ b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
@@ -0,0 +1,43 @@
+! The "use_device_addr" was added to the "target data" directive in OpenMP 5.0.
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 %s -o - | FileCheck %s
+! This test primary goal is to check that we update only base addr for
+! arrays used in used_device_addr clause.
+
+!CHECK: func.func @{{.*}}device_addr_default(
+!CHECK: %[[MAP:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, to, literal) capture(ByRef) name("x") -> !fir.ref<!fir.array<?xi32>>
+!CHECK: omp.target_data use_device_addr(%[[MAP]] -> %[[ARG:.*]] : !fir.ref<!fir.array<?xi32>>) {
+!CHECK:   %[[ALLOCA_TGT_DESC:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
+!CHECK:   %[[ALLOCA_HOST_DESC:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
+!CHECK:   fir.store %[[ARG]] to %[[ALLOCA_HOST_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
+!CHECK:   %[[LOADED_HOST_DESC:.*]] = fir.load %[[ALLOCA_HOST_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
+!CHECK:   %[[HOST_ARR_ADDR:.*]] = fir.box_addr %[[LOADED_HOST_DESC]] : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+!CHECK:   %[[HOST_ARR_PTR:.*]] = fir.convert %[[HOST_ARR_ADDR]] : (!fir.ref<!fir.array<?xi32>>) -> !fir.llvm_ptr<i8>
+!CHECK:   %[[DEVICE_ID:.*]] = fir.call @omp_get_default_device() : () -> i32
+!CHECK:   %[[DEVICE_ID_CONV:.*]] = fir.convert %[[DEVICE_ID]] : (i32) -> i64
+!CHECK:   %[[PTR_ARG:.*]] = fir.convert %[[HOST_ARR_PTR]] : (!fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
+!CHECK:   %[[TGT_PTR:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONV]], %[[PTR_ARG]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
+!CHECK:   %[[TGT_PTR_CONV:.*]] = fir.convert %[[TGT_PTR]] : (!fir.llvm_ptr<i8>) -> !fir.ref<!fir.array<?xi32>>
+!CHECK:   %[[C0:.*]] = arith.constant 0 : index
+!CHECK:   %[[ARR_DIMS:.*]]:3 = fir.box_dims %[[LOADED_HOST_DESC]], %[[C0]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
+!CHECK:   %[[TGT_DESC:.*]] = fir.create_box %[[TGT_PTR_CONV]] lbs(%[[ARR_DIMS]]#0) extents(%[[ARR_DIMS]]#1) strides(%[[ARR_DIMS]]#2) : (!fir.ref<!fir.array<?xi32>>, index, index, index) -> !fir.box<!fir.array<?xi32>>
+!CHECK:   fir.store %[[TGT_DESC]] to %[[ALLOCA_TGT_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
+!CHECK:   %[[RES_TGT_DESC:.*]] = fir.load %[[ALLOCA_TGT_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
+!CHECK:   %[[DECL:.*]] = hlfir.declare %[[RES_TGT_DESC]] {fortran_attrs = #fir.var_attrs<intent_in, target>, uniq_name = "_QFdevice_addr_defaultEx"} : (!fir.box<!fir.array<?xi32>>) -> (!fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>)
+  SUBROUTINE device_addr_default(x)
+    INTEGER, TARGET, INTENT(IN)    :: x(:)
+    !$omp target data use_device_addr (x)
+    !$omp end target data
+  END SUBROUTINE
+
+! Goal: check if we take into account device clause
+!CHECK: func.func @{{.*}}device_addr_device_2(
+!CHECK: omp.target_data device(%[[DEVICE_ID_CONST:.*]] : i32) use_device_addr(%{{.*}} -> %{{.*}} : !fir.ref<!fir.array<?xi32>>)
+!CHECK: %[[DEVICE_ID_CONST_CONV:.*]] = fir.convert %c2_i32 : (i32) -> i64
+!CHECK: %[[TGT_PTR1:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONST_CONV]], %[[BASE_PTR:.*]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
+  SUBROUTINE device_addr_device_2(x)
+    INTEGER, TARGET, INTENT(IN)    :: x(:)
+    !$omp target data use_device_addr (x) device(2)
+    !$omp end target data
+  END SUBROUTINE
+
diff --git a/offload/include/omptarget.h b/offload/include/omptarget.h
index db9590844b2fd..6bd61036cc4b3 100644
--- a/offload/include/omptarget.h
+++ b/offload/include/omptarget.h
@@ -291,6 +291,7 @@ const char *omp_get_uid_from_device(int DeviceNum);
 int omp_get_initial_device(void);
 size_t omp_get_gprivate_limit(int DeviceNum,
                               omp_access_t AccessGroup = omp_access_cgroup);
+void *omp_get_mapped_ptr(const void *Ptr, int DeviceNum);
 void *omp_target_alloc(size_t Size, int DeviceNum);
 void omp_target_free(void *DevicePtr, int DeviceNum);
 int omp_target_is_present(const void *Ptr, int DeviceNum);
@@ -442,6 +443,9 @@ int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize,
                                  void *VAddr, bool IsRecord, bool SaveOutput,
                                  bool EmitReport, const char *OutputDirPath);
 
+// Gets mapped device pointer. If device pointer is not found, returns
+// host pointer
+void *__tgt_get_mapped_ptr(int64_t DeviceId, const void *HostPtr);
 // Registers a callback for the RPC server. Expects this function type.
 // unsigned callback(rpc::Server::Port *Port, unsigned NumLanes). See the RPC
 // code for details.
diff --git a/offload/libomptarget/exports b/offload/libomptarget/exports
index 1831c43cc5f29..91d29fe41679c 100644
--- a/offload/libomptarget/exports
+++ b/offload/libomptarget/exports
@@ -5,6 +5,7 @@ VERS1.0 {
     __tgt_register_requires;
     __tgt_register_lib;
     __tgt_unregister_lib;
+    __tgt_get_mapped_ptr;
     __tgt_init_all_rtls;
     __tgt_target_data_begin;
     __tgt_target_data_end;
diff --git a/offload/libomptarget/interface.cpp b/offload/libomptarget/interface.cpp
index 5d7d948711b99..5833925209fee 100644
--- a/offload/libomptarget/interface.cpp
+++ b/offload/libomptarget/interface.cpp
@@ -653,3 +653,10 @@ EXTERN void __tgt_register_rpc_callback(unsigned (*Callback)(void *,
     if (Plugin.is_initialized() && Plugin.getNumDevices() > 0)
       Plugin.getRPCServer().registerCallback(Callback);
 }
+
+EXTERN void *__tgt_get_mapped_ptr(int64_t DeviceId, const void *HostPtr) {
+  void *TargetPtr = omp_get_mapped_ptr(HostPtr, DeviceId);
+  if (!TargetPtr)
+    return const_cast<void *>(HostPtr);
+  return TargetPtr;
+}

>From c832877ebad604b07da3597c802bdd673a0bff63 Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Thu, 10 Sep 2026 14:09:18 -0500
Subject: [PATCH 2/5] Applied_remarks

---
 .../Optimizer/OpenMP/MapInfoFinalization.cpp  | 34 ++++++++++---------
 .../OpenMP/use-device-addr-performance.f90    |  7 ++--
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index 738064355616a..8523c57387f79 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -1110,9 +1110,9 @@ class MapInfoFinalizationPass
       mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
       mlir::Operation *target, mlir::Value descriptor,
       llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
-      bool isAttachNever, bool isAttachAlways, bool isHasDeviceAddrFlag,
+      bool isAttachNever, bool isAttachAlways, bool mapOnlyDescriptor,
       bool descCanBeDeferred, bool canOptimizeDescViaPrivatization,
-      mlir::FlatSymbolRefAttr mapperId, bool mapOnlyDescriptor) {
+      mlir::FlatSymbolRefAttr mapperId) {
     bool isRefPtrPtee =
         bitEnumContainsAll(op.getMapType(),
                            mlir::omp::ClauseMapFlags::ref_ptr) &&
@@ -1131,7 +1131,7 @@ class MapInfoFinalizationPass
     // For has_device_address we currently do not emit the base address
     // or an attach map.
     mlir::omp::MapInfoOp baseAddr;
-    if (!isHasDeviceAddrFlag) {
+    if (!mapOnlyDescriptor) {
       baseAddr =
           genBaseAddrMap(op.getLoc(), descriptor, op, op.getMapType(), builder,
                          /*IsRefPtee=*/false, mapperId);
@@ -1168,7 +1168,7 @@ class MapInfoFinalizationPass
         /*partial_map=*/builder.getBoolAttr(false));
 
     mlir::Operation *attachMap = nullptr;
-    if (!isAttachNever && !isHasDeviceAddrFlag && !mapOnlyDescriptor) {
+    if (!isAttachNever && !mapOnlyDescriptor) {
       attachMap =
           genImplicitAttachMap(op, descriptor, mapMemberUsers, target, builder,
                                mlir::omp::ClauseMapFlags::ref_ptr |
@@ -1188,6 +1188,12 @@ class MapInfoFinalizationPass
     return newMapInfoOp;
   }
 
+  /// This function checks if the given value is the result of fir::alloca
+  /// operation
+  bool isAllocaOp(mlir::Value &val) {
+    return mlir::isa<fir::AllocaOp>(val.getDefiningOp());
+  }
+
   // This function handles the splitting of allocatable/pointer maps in
   // Fortran into descriptor, pointer and attach map components, as
   // well as the handling of ref_ptr, ref_ptee, ref_ptr_ptee and attach
@@ -1206,6 +1212,7 @@ class MapInfoFinalizationPass
     bool canOptimizeDescViaPrivatization = false;
     llvm::SmallVector<ParentAndPlacement> mapMemberUsers;
     getMemberUserList(op, mapMemberUsers);
+
     // TODO: map the addendum segment of the descriptor, similarly to the
     // base address/data pointer member.
     bool mapOnlyDescriptor = false;
@@ -1225,7 +1232,7 @@ class MapInfoFinalizationPass
 
     mlir::Value descriptor = getDescriptorFromBoxMap(
         op, builder, descCanBeDeferred, canOptimizeDescViaPrivatization);
-    bool isNewDescriptor = mlir::isa<fir::AllocaOp>(descriptor.getDefiningOp());
+    bool isNewDescriptor = isAllocaOp(descriptor);
     bool isUseDeviceAddrItem =
         (getUseDeviceAddrBlockArg(op, *target) != nullptr);
     bool isArray = false;
@@ -1274,7 +1281,7 @@ class MapInfoFinalizationPass
       newMapInfo = genRefPtrPteeOrDefaultMap(
           op, builder, target, descriptor, mapMemberUsers, isAttachNever,
           isAttachAlways, mapOnlyDescriptor, descCanBeDeferred,
-          canOptimizeDescViaPrivatization, mapperId, canOptimizeUseDeviceAddr);
+          canOptimizeDescViaPrivatization, mapperId);
     }
     return newMapInfo;
   }
@@ -1499,12 +1506,7 @@ class MapInfoFinalizationPass
     // use_device_addr.
     auto allocaTgtDescriptor =
         fir::AllocaOp::create(builder, loc, arg.getType());
-    auto allocaHostDescriptor =
-        fir::AllocaOp::create(builder, loc, arg.getType());
-    fir::StoreOp::create(builder, loc, arg, allocaHostDescriptor);
-    auto hostDescriptor =
-        fir::LoadOp::create(builder, loc, allocaHostDescriptor);
-    auto hostAddrPtr = fir::BoxAddrOp::create(builder, loc, hostDescriptor);
+    auto hostAddrPtr = fir::BoxAddrOp::create(builder, loc, arg);
     auto convertedAddr = fir::ConvertOp::create(
         builder, loc,
         fir::LLVMPointerType::get(builder.getContext(), builder.getI8Type()),
@@ -1516,11 +1518,11 @@ class MapInfoFinalizationPass
     llvm::SmallVector<mlir::Value> lbounds;
     llvm::SmallVector<mlir::Value> extents;
     llvm::SmallVector<mlir::Value> strides;
-    fir::factory::genDimInfoFromBox(builder, loc, hostDescriptor, &lbounds,
-                                    &extents, &strides);
+    fir::factory::genDimInfoFromBox(builder, loc, arg, &lbounds, &extents,
+                                    &strides);
     auto newDescriptor =
-        fir::CreateBoxOp::create(builder, loc, hostDescriptor.getType(),
-                                 convertedGPUAddr, lbounds, extents, strides);
+        fir::CreateBoxOp::create(builder, loc, arg.getType(), convertedGPUAddr,
+                                 lbounds, extents, strides);
     fir::StoreOp::create(builder, loc, newDescriptor, allocaTgtDescriptor);
     auto res = fir::LoadOp::create(builder, loc, allocaTgtDescriptor);
     arg.replaceUsesWithIf(res, [&](mlir::OpOperand &use) {
diff --git a/flang/test/Lower/OpenMP/use-device-addr-performance.f90 b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
index c4785e7e8a097..406e34dcfcd5c 100644
--- a/flang/test/Lower/OpenMP/use-device-addr-performance.f90
+++ b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
@@ -8,10 +8,7 @@
 !CHECK: %[[MAP:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, to, literal) capture(ByRef) name("x") -> !fir.ref<!fir.array<?xi32>>
 !CHECK: omp.target_data use_device_addr(%[[MAP]] -> %[[ARG:.*]] : !fir.ref<!fir.array<?xi32>>) {
 !CHECK:   %[[ALLOCA_TGT_DESC:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
-!CHECK:   %[[ALLOCA_HOST_DESC:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
-!CHECK:   fir.store %[[ARG]] to %[[ALLOCA_HOST_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
-!CHECK:   %[[LOADED_HOST_DESC:.*]] = fir.load %[[ALLOCA_HOST_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
-!CHECK:   %[[HOST_ARR_ADDR:.*]] = fir.box_addr %[[LOADED_HOST_DESC]] : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+!CHECK:   %[[HOST_ARR_ADDR:.*]] = fir.box_addr %[[ARG]] : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
 !CHECK:   %[[HOST_ARR_PTR:.*]] = fir.convert %[[HOST_ARR_ADDR]] : (!fir.ref<!fir.array<?xi32>>) -> !fir.llvm_ptr<i8>
 !CHECK:   %[[DEVICE_ID:.*]] = fir.call @omp_get_default_device() : () -> i32
 !CHECK:   %[[DEVICE_ID_CONV:.*]] = fir.convert %[[DEVICE_ID]] : (i32) -> i64
@@ -19,7 +16,7 @@
 !CHECK:   %[[TGT_PTR:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONV]], %[[PTR_ARG]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
 !CHECK:   %[[TGT_PTR_CONV:.*]] = fir.convert %[[TGT_PTR]] : (!fir.llvm_ptr<i8>) -> !fir.ref<!fir.array<?xi32>>
 !CHECK:   %[[C0:.*]] = arith.constant 0 : index
-!CHECK:   %[[ARR_DIMS:.*]]:3 = fir.box_dims %[[LOADED_HOST_DESC]], %[[C0]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
+!CHECK:   %[[ARR_DIMS:.*]]:3 = fir.box_dims %[[ARG]], %[[C0]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
 !CHECK:   %[[TGT_DESC:.*]] = fir.create_box %[[TGT_PTR_CONV]] lbs(%[[ARR_DIMS]]#0) extents(%[[ARR_DIMS]]#1) strides(%[[ARR_DIMS]]#2) : (!fir.ref<!fir.array<?xi32>>, index, index, index) -> !fir.box<!fir.array<?xi32>>
 !CHECK:   fir.store %[[TGT_DESC]] to %[[ALLOCA_TGT_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
 !CHECK:   %[[RES_TGT_DESC:.*]] = fir.load %[[ALLOCA_TGT_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>

>From 8edf3c77761fdf5b3b927bebf240985329374930 Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Mon, 14 Sep 2026 10:16:04 -0500
Subject: [PATCH 3/5] Applied remarks

---
 .../Optimizer/OpenMP/MapInfoFinalization.cpp  | 35 ++++++++++++----
 .../OpenMP/use-device-addr-performance.f90    | 11 ++---
 .../fortran/target-use-device-addr-opt.f90    | 41 +++++++++++++++++++
 3 files changed, 73 insertions(+), 14 deletions(-)
 create mode 100644 offload/test/offloading/fortran/target-use-device-addr-opt.f90

diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index f22340810b2b7..db6921c146478 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -1530,11 +1530,28 @@ class MapInfoFinalizationPass
                                  mlir::omp::TargetDataOp targetDataOp,
                                  mlir::omp::MapInfoOp mapOp,
                                  mlir::ModuleOp module) {
-    mlir::Location loc = targetDataOp.getLoc();
+    // Disable mapping of the descriptor to the GPU by setting literal map type
     mapOp.setMapType(mapOp.getMapType() | mlir::omp::ClauseMapFlags::literal);
+    auto loc = targetDataOp.getLoc();
+    // Make sure that updateUseDeviceDescriptorArgs was launched earlier
     auto arg = getUseDeviceAddrBlockArg(mapOp, *targetDataOp.getOperation());
+    bool isArgBoxType = mlir::isa<fir::BaseBoxType>(arg.getType());
+    bool useArgLoad =
+        arg.hasOneUse() && mlir::isa<fir::LoadOp>(*arg.use_begin()->getOwner());
+    assert((isArgBoxType || useArgLoad) &&
+           "Expected either BaseBox item or Load operation");
+
+    auto argUserOp = arg.use_begin()->getOwner();
     auto insertionPoint = builder.saveInsertionPoint();
-    builder.setInsertionPoint(&targetDataOp->getRegion(0).front().front());
+    mlir::Value hostDescriptor;
+    if (isArgBoxType) {
+      hostDescriptor = arg;
+      builder.setInsertionPoint(&targetDataOp->getRegion(0).front().front());
+    } else {
+      auto loadOp = arg.use_begin()->getOwner();
+      hostDescriptor = loadOp->getResult(0);
+      builder.setInsertionPointAfter(loadOp);
+    }
     // We need to create a temporary copy of the host descriptor, which will
     // be used inside the use_device_addr code region. The copy will be updated
     // with the target pointer of the mapped array. The lifetime of the
@@ -1543,8 +1560,8 @@ class MapInfoFinalizationPass
     // descriptor if we want to update the host descriptor inside the
     // use_device_addr.
     auto allocaTgtDescriptor =
-        fir::AllocaOp::create(builder, loc, arg.getType());
-    auto hostAddrPtr = fir::BoxAddrOp::create(builder, loc, arg);
+        fir::AllocaOp::create(builder, loc, hostDescriptor.getType());
+    auto hostAddrPtr = fir::BoxAddrOp::create(builder, loc, hostDescriptor);
     auto convertedAddr = fir::ConvertOp::create(
         builder, loc,
         fir::LLVMPointerType::get(builder.getContext(), builder.getI8Type()),
@@ -1556,14 +1573,14 @@ class MapInfoFinalizationPass
     llvm::SmallVector<mlir::Value> lbounds;
     llvm::SmallVector<mlir::Value> extents;
     llvm::SmallVector<mlir::Value> strides;
-    fir::factory::genDimInfoFromBox(builder, loc, arg, &lbounds, &extents,
-                                    &strides);
+    fir::factory::genDimInfoFromBox(builder, loc, hostDescriptor, &lbounds,
+                                    &extents, &strides);
     auto newDescriptor =
-        fir::CreateBoxOp::create(builder, loc, arg.getType(), convertedGPUAddr,
-                                 lbounds, extents, strides);
+        fir::CreateBoxOp::create(builder, loc, hostDescriptor.getType(),
+                                 convertedGPUAddr, lbounds, extents, strides);
     fir::StoreOp::create(builder, loc, newDescriptor, allocaTgtDescriptor);
     auto res = fir::LoadOp::create(builder, loc, allocaTgtDescriptor);
-    arg.replaceUsesWithIf(res, [&](mlir::OpOperand &use) {
+    hostDescriptor.replaceUsesWithIf(res, [&](mlir::OpOperand &use) {
       mlir::Operation *user = use.getOwner();
       return res->isBeforeInBlock(user);
     });
diff --git a/flang/test/Lower/OpenMP/use-device-addr-performance.f90 b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
index 406e34dcfcd5c..93321c644e6ff 100644
--- a/flang/test/Lower/OpenMP/use-device-addr-performance.f90
+++ b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
@@ -5,10 +5,11 @@
 ! arrays used in used_device_addr clause.
 
 !CHECK: func.func @{{.*}}device_addr_default(
-!CHECK: %[[MAP:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, to, literal) capture(ByRef) name("x") -> !fir.ref<!fir.array<?xi32>>
-!CHECK: omp.target_data use_device_addr(%[[MAP]] -> %[[ARG:.*]] : !fir.ref<!fir.array<?xi32>>) {
+!CHECK: %[[MAP:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, to, literal) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.array<?xi32>>>
+!CHECK: omp.target_data use_device_addr(%[[MAP]] -> %[[ARG:.*]] : !fir.ref<!fir.box<!fir.array<?xi32>>>) {
+!CHECK:   %[[HOST_DESCRIPTOR:.*]] = fir.load %[[ARG]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
 !CHECK:   %[[ALLOCA_TGT_DESC:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
-!CHECK:   %[[HOST_ARR_ADDR:.*]] = fir.box_addr %[[ARG]] : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+!CHECK:   %[[HOST_ARR_ADDR:.*]] = fir.box_addr %[[HOST_DESCRIPTOR]] : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
 !CHECK:   %[[HOST_ARR_PTR:.*]] = fir.convert %[[HOST_ARR_ADDR]] : (!fir.ref<!fir.array<?xi32>>) -> !fir.llvm_ptr<i8>
 !CHECK:   %[[DEVICE_ID:.*]] = fir.call @omp_get_default_device() : () -> i32
 !CHECK:   %[[DEVICE_ID_CONV:.*]] = fir.convert %[[DEVICE_ID]] : (i32) -> i64
@@ -16,7 +17,7 @@
 !CHECK:   %[[TGT_PTR:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONV]], %[[PTR_ARG]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
 !CHECK:   %[[TGT_PTR_CONV:.*]] = fir.convert %[[TGT_PTR]] : (!fir.llvm_ptr<i8>) -> !fir.ref<!fir.array<?xi32>>
 !CHECK:   %[[C0:.*]] = arith.constant 0 : index
-!CHECK:   %[[ARR_DIMS:.*]]:3 = fir.box_dims %[[ARG]], %[[C0]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
+!CHECK:   %[[ARR_DIMS:.*]]:3 = fir.box_dims %[[HOST_DESCRIPTOR]], %[[C0]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
 !CHECK:   %[[TGT_DESC:.*]] = fir.create_box %[[TGT_PTR_CONV]] lbs(%[[ARR_DIMS]]#0) extents(%[[ARR_DIMS]]#1) strides(%[[ARR_DIMS]]#2) : (!fir.ref<!fir.array<?xi32>>, index, index, index) -> !fir.box<!fir.array<?xi32>>
 !CHECK:   fir.store %[[TGT_DESC]] to %[[ALLOCA_TGT_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
 !CHECK:   %[[RES_TGT_DESC:.*]] = fir.load %[[ALLOCA_TGT_DESC]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
@@ -29,7 +30,7 @@ SUBROUTINE device_addr_default(x)
 
 ! Goal: check if we take into account device clause
 !CHECK: func.func @{{.*}}device_addr_device_2(
-!CHECK: omp.target_data device(%[[DEVICE_ID_CONST:.*]] : i32) use_device_addr(%{{.*}} -> %{{.*}} : !fir.ref<!fir.array<?xi32>>)
+!CHECK: omp.target_data device(%[[DEVICE_ID_CONST:.*]] : i32) use_device_addr(%{{.*}} -> %{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>)
 !CHECK: %[[DEVICE_ID_CONST_CONV:.*]] = fir.convert %c2_i32 : (i32) -> i64
 !CHECK: %[[TGT_PTR1:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONST_CONV]], %[[BASE_PTR:.*]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
   SUBROUTINE device_addr_device_2(x)
diff --git a/offload/test/offloading/fortran/target-use-device-addr-opt.f90 b/offload/test/offloading/fortran/target-use-device-addr-opt.f90
new file mode 100644
index 0000000000000..bb4a49973e2d3
--- /dev/null
+++ b/offload/test/offloading/fortran/target-use-device-addr-opt.f90
@@ -0,0 +1,41 @@
+! REQUIRES: flang
+! REQUIRES: gpu, amdgpu
+
+! RUN: %libomptarget-compile-fortran-generic
+! RUN: env LIBOMPTARGET_INFO=8  %libomptarget-run-generic 2>&1 | %fcheck-generic
+MODULE foo
+  IMPLICIT NONE
+  PRIVATE
+  PUBLIC :: bar_device_addr
+
+CONTAINS
+
+  SUBROUTINE bar_device_addr(x)
+    INTEGER, TARGET, INTENT(IN)    :: x(:)
+    !$omp target data use_device_addr (x)
+    !$omp end target data
+  END SUBROUTINE
+
+END MODULE foo
+
+PROGRAM test_ptr
+  USE, intrinsic :: iso_fortran_env, only: error_unit
+  USE foo
+  IMPLICIT NONE
+
+  INTEGER, ALLOCATABLE, TARGET :: x(:)
+  ALLOCATE(x(10))
+  !$omp target enter data map(to: x)
+    CALL bar_device_addr(x)
+  !$omp target exit data map(from: x)
+  DEALLOCATE(x)
+  write(error_unit, *) 'Success'
+END PROGRAM test_ptr
+
+! CHECK: Creating new map entry
+! CHECK: Creating new map entry
+! CHECK-NOT: Creating new map entry
+! CHECK: Removing map entry
+! CHECK: Removing map entry
+! CHECK-NOT: Removing map entry
+! CHECK: Success

>From 3b0381112dd3e6c693b0874cbc65b081f23b2f80 Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Wed, 16 Sep 2026 01:51:45 -0500
Subject: [PATCH 4/5] Addressed Akash's comments

---
 .../Optimizer/OpenMP/MapInfoFinalization.cpp  | 142 +++++++++++++-----
 .../OpenMP/use-device-addr-performance.f90    |  47 +++++-
 2 files changed, 146 insertions(+), 43 deletions(-)

diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index db6921c146478..5e4b9b956f634 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -1191,7 +1191,7 @@ class MapInfoFinalizationPass
   /// This function checks if the given value is the result of fir::alloca
   /// operation
   bool isAllocaOp(mlir::Value &val) {
-    return mlir::isa<fir::AllocaOp>(val.getDefiningOp());
+    return mlir::isa_and_present<fir::AllocaOp>(val.getDefiningOp());
   }
 
   /// Use the materialized descriptor's address for target data block arguments.
@@ -1231,6 +1231,57 @@ class MapInfoFinalizationPass
                argIface.getUseDevicePtrBlockArgs());
   }
 
+  /// Check if we can optimize the descriptor mappings for use_device_addr
+  bool canOptimizeUseDeviceAddrMapping(fir::FirOpBuilder &builder,
+                                       mlir::Value descriptor,
+                                       mlir::omp::MapInfoOp op,
+                                       mlir::Operation *target) {
+    // optimize only temporary descriptors (i.e. allocated on function stack)
+    if (!isAllocaOp(descriptor))
+      return false;
+    // check if given descriptor is mapped as use_device_addr argument
+    if (getUseDeviceAddrBlockArg(op, *target) == nullptr)
+      return false;
+    auto module = builder.getModule();
+    // check if the OpenMP offload target device is specified
+    auto iface =
+        llvm::cast<mlir::omp::OffloadModuleInterface>(module.getOperation());
+    if (iface.getTargetTriples().empty())
+      return false;
+
+    // Only optimize descriptors whose element size is known at compile time.
+    // Excluded:
+    //   - polymorphic entities (!fir.class): elem_len is a runtime property,
+    //     since the dynamic type may extend the declared type.
+    //   - deferred-length characters (!fir.char<k,?>) and parameterized
+    //     derived types: LEN parameters are runtime values.
+    //   - assumed-rank arrays (!fir.array<*:T>): descriptor layout depends on
+    //     a rank that is not known here.
+    //
+    // TODO: The restrictions can be lifted if fir.create_box operation supports
+    // creation of box with dynamical element size.
+    bool isArray = false;
+    bool knownRanks = false;
+    fir::BaseBoxType baseBoxTy = mlir::dyn_cast<fir::BaseBoxType>(
+        fir::unwrapRefType(descriptor.getType()));
+    if (baseBoxTy) {
+      isArray = baseBoxTy.isArray();
+      knownRanks = !baseBoxTy.isAssumedRank();
+    }
+    if (!isArray)
+      return false;
+    if (!knownRanks)
+      return false;
+    auto eleTy = baseBoxTy.unwrapInnerType();
+    if (fir::hasDynamicSize(eleTy))
+      return false;
+    if (fir::isPolymorphicType(eleTy))
+      return false;
+    if (fir::isAssumedType(eleTy))
+      return false;
+    return true;
+  }
+
   // This function handles the splitting of allocatable/pointer maps in
   // Fortran into descriptor, pointer and attach map components, as
   // well as the handling of ref_ptr, ref_ptee, ref_ptr_ptee and attach
@@ -1270,19 +1321,8 @@ class MapInfoFinalizationPass
     mlir::Value descriptor = getDescriptorFromBoxMap(
         op, builder, descCanBeDeferred, canOptimizeDescViaPrivatization);
     updateUseDeviceDescriptorArgs(op, descriptor, target, builder);
-    bool isNewDescriptor = isAllocaOp(descriptor);
-    bool isUseDeviceAddrItem =
-        (getUseDeviceAddrBlockArg(op, *target) != nullptr);
-    bool isArray = false;
-    bool knownRanks = false;
-    fir::BaseBoxType bt = mlir::dyn_cast<fir::BaseBoxType>(
-        fir::unwrapRefType(descriptor.getType()));
-    if (bt) {
-      isArray = bt.isArray();
-      knownRanks = !bt.isAssumedRank();
-    }
     canOptimizeUseDeviceAddr =
-        (isNewDescriptor && isUseDeviceAddrItem && isArray && knownRanks);
+        canOptimizeUseDeviceAddrMapping(builder, descriptor, op, target);
     mapOnlyDescriptor = isHasDeviceAddrFlag | canOptimizeUseDeviceAddr;
     mlir::FlatSymbolRefAttr mapperId = op.getMapperIdAttr();
     // Exclude irregular maps from optimization via privatization; at least for
@@ -1482,12 +1522,59 @@ class MapInfoFinalizationPass
 
   mlir::Value genTgtGetMappedPtrCall(fir::FirOpBuilder &builder,
                                      mlir::Location loc, mlir::Value deviceNum,
-                                     mlir::Value hostPtr,
+                                     mlir::Value ifCond, mlir::Value hostPtr,
                                      mlir::ModuleOp module) {
     auto *context = builder.getContext();
     auto voidPtrType = fir::LLVMPointerType::get(context, builder.getI8Type());
     auto i32Type = builder.getI32Type();
     auto i64Type = builder.getI64Type();
+
+    // Helper funtion which creates calls to omp_get_default_device()
+    // or omp_get_initial_device()
+    auto createOmpGetFunction = [&](llvm::StringRef funcName) -> mlir::Value {
+      auto funcOmpGetDeviceOp =
+          module.lookupSymbol<mlir::func::FuncOp>(funcName);
+      if (!funcOmpGetDeviceOp) {
+        auto funcType = mlir::FunctionType::get(context, {}, {i32Type});
+
+        mlir::OpBuilder::InsertionGuard guard(builder);
+        builder.setInsertionPointToStart(module.getBody());
+
+        funcOmpGetDeviceOp =
+            mlir::func::FuncOp::create(builder, loc, funcName, funcType);
+        funcOmpGetDeviceOp.setPrivate();
+      }
+      auto callOmpGetFuncOp =
+          fir::CallOp::create(builder, loc, funcOmpGetDeviceOp, {});
+      return callOmpGetFuncOp.getResult(0);
+      ;
+    };
+
+    if (!deviceNum) {
+      deviceNum = createOmpGetFunction("omp_get_default_device");
+    }
+
+    if (ifCond) {
+      auto allocaIfDevice = fir::AllocaOp::create(builder, loc, i32Type);
+      mlir::Value initialDeviceNum = nullptr;
+      mlir::Value boolIfCond =
+          builder.createConvert(loc, builder.getI1Type(), ifCond);
+
+      builder.genIfThenElse(loc, boolIfCond)
+          .genThen([&]() {
+            fir::StoreOp::create(builder, loc, deviceNum, allocaIfDevice);
+          })
+          .genElse([&]() {
+            // omp_get_initial_device returns host id
+            initialDeviceNum = createOmpGetFunction("omp_get_initial_device");
+            fir::StoreOp::create(builder, loc, initialDeviceNum,
+                                 allocaIfDevice);
+          })
+          .end();
+      auto loadIfDevice = fir::LoadOp::create(builder, loc, allocaIfDevice);
+      deviceNum = loadIfDevice.getResult();
+    }
+
     auto funcName = "__tgt_get_mapped_ptr";
     auto funcOp = module.lookupSymbol<mlir::func::FuncOp>(funcName);
 
@@ -1501,24 +1588,6 @@ class MapInfoFinalizationPass
       funcOp = mlir::func::FuncOp::create(builder, loc, funcName, funcType);
       funcOp.setPrivate();
     }
-    if (!deviceNum) {
-      auto funcGetDefaultDeviceName = "omp_get_default_device";
-      auto funcGetDefaultDeviceOp =
-          module.lookupSymbol<mlir::func::FuncOp>(funcGetDefaultDeviceName);
-      if (!funcGetDefaultDeviceOp) {
-        auto funcType = mlir::FunctionType::get(context, {}, {i32Type});
-
-        mlir::OpBuilder::InsertionGuard guard(builder);
-        builder.setInsertionPointToStart(module.getBody());
-
-        funcGetDefaultDeviceOp = mlir::func::FuncOp::create(
-            builder, loc, funcGetDefaultDeviceName, funcType);
-        funcGetDefaultDeviceOp.setPrivate();
-      }
-      auto callGetDefaultDeviceOp =
-          fir::CallOp::create(builder, loc, funcGetDefaultDeviceOp, {});
-      deviceNum = callGetDefaultDeviceOp.getResult(0);
-    }
     llvm::SmallVector<mlir::Value> args;
     args.push_back(fir::ConvertOp::create(builder, loc, i64Type, deviceNum));
     args.push_back(fir::ConvertOp::create(builder, loc, voidPtrType, hostPtr));
@@ -1541,8 +1610,7 @@ class MapInfoFinalizationPass
     assert((isArgBoxType || useArgLoad) &&
            "Expected either BaseBox item or Load operation");
 
-    auto argUserOp = arg.use_begin()->getOwner();
-    auto insertionPoint = builder.saveInsertionPoint();
+    mlir::OpBuilder::InsertionGuard guard(builder);
     mlir::Value hostDescriptor;
     if (isArgBoxType) {
       hostDescriptor = arg;
@@ -1566,8 +1634,9 @@ class MapInfoFinalizationPass
         builder, loc,
         fir::LLVMPointerType::get(builder.getContext(), builder.getI8Type()),
         hostAddrPtr);
-    auto newAddr = genTgtGetMappedPtrCall(
-        builder, loc, targetDataOp.getDevice(), convertedAddr, module);
+    auto newAddr =
+        genTgtGetMappedPtrCall(builder, loc, targetDataOp.getDevice(),
+                               targetDataOp.getIfExpr(), convertedAddr, module);
     auto convertedGPUAddr =
         fir::ConvertOp::create(builder, loc, hostAddrPtr.getType(), newAddr);
     llvm::SmallVector<mlir::Value> lbounds;
@@ -1584,7 +1653,6 @@ class MapInfoFinalizationPass
       mlir::Operation *user = use.getOwner();
       return res->isBeforeInBlock(user);
     });
-    builder.restoreInsertionPoint(insertionPoint);
   }
 
   // This pass executes on omp::MapInfoOp's containing descriptor based types
diff --git a/flang/test/Lower/OpenMP/use-device-addr-performance.f90 b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
index 93321c644e6ff..6bedd346a31e9 100644
--- a/flang/test/Lower/OpenMP/use-device-addr-performance.f90
+++ b/flang/test/Lower/OpenMP/use-device-addr-performance.f90
@@ -1,6 +1,7 @@
 ! The "use_device_addr" was added to the "target data" directive in OpenMP 5.0.
-! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 %s -o - | FileCheck %s
-! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 -fopenmp-targets=amdgcn-amd-amdhsa %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 -fopenmp-targets=amdgcn-amd-amdhsa %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 %s -o - | FileCheck %s --check-prefix=HOSTONLY
 ! This test primary goal is to check that we update only base addr for
 ! arrays used in used_device_addr clause.
 
@@ -29,13 +30,47 @@ SUBROUTINE device_addr_default(x)
   END SUBROUTINE
 
 ! Goal: check if we take into account device clause
-!CHECK: func.func @{{.*}}device_addr_device_2(
-!CHECK: omp.target_data device(%[[DEVICE_ID_CONST:.*]] : i32) use_device_addr(%{{.*}} -> %{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>)
-!CHECK: %[[DEVICE_ID_CONST_CONV:.*]] = fir.convert %c2_i32 : (i32) -> i64
-!CHECK: %[[TGT_PTR1:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONST_CONV]], %[[BASE_PTR:.*]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
+!CHECK:   func.func @{{.*}}device_addr_device_2(
+!CHECK:   omp.target_data device(%[[DEVICE_ID_CONST:.*]] : i32) use_device_addr(%{{.*}} -> %{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>)
+!CHECK:   %[[DEVICE_ID_CONST_CONV:.*]] = fir.convert %c2_i32 : (i32) -> i64
+!CHECK:   %[[TGT_PTR1:.*]] = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONST_CONV]], %[[BASE_PTR:.*]]) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
   SUBROUTINE device_addr_device_2(x)
     INTEGER, TARGET, INTENT(IN)    :: x(:)
     !$omp target data use_device_addr (x) device(2)
     !$omp end target data
   END SUBROUTINE
 
+! Goal: check if we take into account if clause
+!CHECK:   func.func @{{.*}}device_addr_device_if(
+!CHECK:   omp.target_data if(%[[COND:.*]]) use_device_addr(%{{.*}} -> %{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>) {
+!CHECK:   %[[GPU_ID:.*]] = fir.call @omp_get_default_device() : () -> i32
+!CHECK:   %[[DEVICE_ID_ALLOCA:.*]] = fir.alloca i32
+!CHECK:   fir.if %[[COND]] {
+!CHECK:     fir.store %[[GPU_ID]] to %[[DEVICE_ID_ALLOCA]] : !fir.ref<i32>
+!CHECK:   } else {
+!CHECK:     %[[HOST_ID:.*]] = fir.call @omp_get_initial_device() : () -> i32
+!CHECK:      fir.store %[[HOST_ID]] to %[[DEVICE_ID_ALLOCA]] : !fir.ref<i32>
+!CHECK:   }
+!CHECK:   %[[DEVICE_ID:.*]] = fir.load %[[DEVICE_ID_ALLOCA]] : !fir.ref<i32>
+!CHECK:   %[[DEVICE_ID_CONV:.*]] = fir.convert %[[DEVICE_ID]] : (i32) -> i64
+!CHECK:  %{{.*}} = fir.call @__tgt_get_mapped_ptr(%[[DEVICE_ID_CONV]], %{{.*}}) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
+
+  SUBROUTINE device_addr_device_if(x, n)
+    INTEGER, TARGET, INTENT(IN)    :: x(:)
+    INTEGER, INTENT(IN)            :: n
+    !$omp target data use_device_addr (x) if(n > 2)
+    !$omp end target data
+  END SUBROUTINE
+
+! Goal: check if we don't optimize array with dynamically sized elements
+!CHECK:   func.func @{{.*}}device_addr_device_char(
+!CHECK-NOT:  %{{.*}} = fir.call @__tgt_get_mapped_ptr(%{{.*}}, %{{.*}}) : (i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>
+  SUBROUTINE device_addr_device_char(x, n)
+    CHARACTER(*), TARGET :: x(:)
+    !$omp target data use_device_addr(x)
+    !$omp end target data
+  END SUBROUTINE
+
+! Goal: check if we skip optimization for host only code (i.e. we don't use
+! __tgt_get_mapped_ptr).
+! HOSTONLY-NOT: func.func private @__tgt_get_mapped_ptr(i64, !fir.llvm_ptr<i8>) -> !fir.llvm_ptr<i8>

>From 502edb07aadd0527a11ef317c932961f53f4afec Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Thu, 17 Sep 2026 10:06:01 -0500
Subject: [PATCH 5/5] Applied remark

---
 flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index 5e4b9b956f634..b3b235e07aece 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -1275,7 +1275,7 @@ class MapInfoFinalizationPass
     auto eleTy = baseBoxTy.unwrapInnerType();
     if (fir::hasDynamicSize(eleTy))
       return false;
-    if (fir::isPolymorphicType(eleTy))
+    if (fir::isPolymorphicType(baseBoxTy))
       return false;
     if (fir::isAssumedType(eleTy))
       return false;



More information about the flang-commits mailing list