[MLIR][OpenMP] Added MLIR translation support for use_device clauses

[TIFitis]

Summary

Added MLIR support for translating use_device_ptr and use_device_addr clauses for LLVMIR lowering.

• use_device_ptr: The mapped variables marked with use_device_ptr are accessed through a copy of the base pointer mappers. The mapper is copied onto a new temporary pointer variable.
• use_device_addr: The mapped variables marked with use_device_addr are accessed directly through the base pointer mappers.
• If mapping information is not provided explicitly then default map_type of alloc/release is assumed and the map_size is set to 0.

Depends on D152554

[TIFitis]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

4069–4070

use_dev_ptr operands must always be pointer types so I think it is safe to always just create a ptr for them.

However if we want to support typed pointers then we would have to create something like:
Builder.CreateAlloca(UseDeviceOp.first->getType()->getPointerElementType())

This is not possible as getPointerElementType is deprecated and I couldn't find any alternative for this.

4175–4185

I have commented out some code in OpenMPToLLVMIRTranslation.cpp::1489. I think that is a better way of generating code for the use_dev_ptr, but the mapValue doesn't allow remapping.

This I am stuck with letting it use the old value inside the region and replacing it later here.

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

5017–5025

This unit test fails because of opaque pointer mismatch, The unit test runs with opaque pointers disabled however the my code only runs with it enabled.

I could not find a reasonable way to support typed pointers, I have explained this in another comment.

[TIFitis]

Split use_dev_ptr test into two.

[TIFitis]

Rebased

[TIFitis]

Need to change this patch to account for all the clang migration changes and make use of those functions.

[TIFitis]

Made changes to use updated OMPIRBuilder

[TIFitis]

Ping for review :)

[kiranchandramohan]

Could you describe the changes in the Summary? Particularly providing a short description of the code that is being emitted for use_device_ptr and use_device_addr?

[TIFitis]

Here is an example of what clang generates:

C code:

void foo(int *a, int *b, int *c){
#pragma omp target data map(from:b, c) use_device_ptr(a) use_device_addr(b){
    *a = 10;
    *b = 20;
    *c = 30;
  }
}

LLVM IR:

@.offload_sizes = private unnamed_addr constant [3 x i64] [i64 8, i64 8, i64 0]
@.offload_maptypes = private unnamed_addr constant [3 x i64] [i64 66, i64 2, i64 64]

define dso_local void @foo(ptr noundef %a, ptr noundef %b, ptr noundef %c) #0 {
entry:
  %a.addr = alloca ptr, align 8
  %b.addr = alloca ptr, align 8
  %c.addr = alloca ptr, align 8
  %.offload_baseptrs = alloca [3 x ptr], align 8
  %.offload_ptrs = alloca [3 x ptr], align 8
  %.offload_mappers = alloca [3 x ptr], align 8
  %0 = alloca ptr, align 8
  store ptr %a, ptr %a.addr, align 8
  store ptr %b, ptr %b.addr, align 8
  store ptr %c, ptr %c.addr, align 8
  %1 = load ptr, ptr %a.addr, align 8
  %2 = getelementptr inbounds [3 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
  store ptr %b.addr, ptr %2, align 8
  %3 = getelementptr inbounds [3 x ptr], ptr %.offload_ptrs, i32 0, i32 0
  store ptr %b.addr, ptr %3, align 8
  %4 = getelementptr inbounds [3 x ptr], ptr %.offload_mappers, i64 0, i64 0
  store ptr null, ptr %4, align 8
  %5 = getelementptr inbounds [3 x ptr], ptr %.offload_baseptrs, i32 0, i32 1
  store ptr %c.addr, ptr %5, align 8
  %6 = getelementptr inbounds [3 x ptr], ptr %.offload_ptrs, i32 0, i32 1
  store ptr %c.addr, ptr %6, align 8
  %7 = getelementptr inbounds [3 x ptr], ptr %.offload_mappers, i64 0, i64 1
  store ptr null, ptr %7, align 8
  %8 = getelementptr inbounds [3 x ptr], ptr %.offload_baseptrs, i32 0, i32 2
  store ptr %1, ptr %8, align 8
  %9 = getelementptr inbounds [3 x ptr], ptr %.offload_ptrs, i32 0, i32 2
  store ptr %1, ptr %9, align 8
  %10 = getelementptr inbounds [3 x ptr], ptr %.offload_mappers, i64 0, i64 2
  store ptr null, ptr %10, align 8
  %11 = getelementptr inbounds [3 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
  %12 = getelementptr inbounds [3 x ptr], ptr %.offload_ptrs, i32 0, i32 0
  call void @__tgt_target_data_begin_mapper(ptr @1, i64 -1, i32 3, ptr %11, ptr %12, ptr @.offload_sizes, ptr @.offload_maptypes, ptr null, ptr null)
  %13 = load ptr, ptr %8, align 8
  store ptr %13, ptr %0, align 8
  %14 = load ptr, ptr %2, align 8
  %15 = load ptr, ptr %0, align 8
  store i32 10, ptr %15, align 4
  %16 = load ptr, ptr %14, align 8
  store i32 20, ptr %16, align 4
  %17 = load ptr, ptr %c.addr, align 8
  store i32 30, ptr %17, align 4
  %18 = getelementptr inbounds [3 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
  %19 = getelementptr inbounds [3 x ptr], ptr %.offload_ptrs, i32 0, i32 0
  call void @__tgt_target_data_end_mapper(ptr @1, i64 -1, i32 3, ptr %18, ptr %19, ptr @.offload_sizes, ptr @.offload_maptypes, ptr null, ptr null)
  ret void
}

[TIFitis]

Ping for review

[kiranchandramohan]

I am starting to go through the patch.

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

1422–1438

This block seems to be similar to devPtrOperands. Could you share code here?

mlir/test/Target/LLVMIR/omptarget-llvm.mlir

148–238

Please add a testcase with both use_device_ptr and use_device_addr.

[TIFitis]

Rebased. Addressed reviewer comments. Fixed bug in generating map info for device_addr - device info type was incorrectly being set to pointer.

[TIFitis]

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

1422–1438

I've moved them into a functor.

[kiranchandramohan]

There is a build failure currently.

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

1382

failure()?

1390

failure()?

[TIFitis]

Renamed failure() to fail()

[TIFitis]

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

1382

I was calling the failed() functor defined above. I've renamed it to fail to avoid confusion.

[TIFitis]

Attempt at fixing build error

[kiranchandramohan]

While looking into this patch, I noticed that the current implementation in Semantics and MLIR requires the MAP clause to always be present. But the standard only says the following,

At least one map, use_device_addr or use_device_ptr clause must appear on the directive.

I will fix this. But this made having a look at the code generated only with use_device_addr, use_device_ptr and without map difficult.

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

1375–1376

Nit:This only needs to take in as reference combinedInfo.

1516

Nit: expand the type here. Put an assert before this line for the op being a DataOp.

1516

Should this be an assert instead of an if? If not, Please add a comment on when it can and cannot be empty.

1516

Should this be an assert instead of an if? If not, Please add a comment on when it can and cannot be empty.

1518–1531

Might be worth creating a lambda and callling it separately with useDevPtrOperands, useDevAddrOperands.

mlir/test/Target/LLVMIR/omptarget-llvm.mlir

193

Is the default usage of dev_addr also with the pointer type?
If not using a single reference example might be easier to follow.

llvm.func @_QPopenmp_target_use_dev_addr() {
  %0 = llvm.mlir.constant(1 : i64) : i64
  %a = llvm.alloca %0 x !llvm.ptr<i32> : (i64) -> !llvm.ptr<i32>
  omp.target_data  map((from -> %a : !llvm.ptr<i32>)) use_device_addr(%a : !llvm.ptr<i32>)  {
  ^bb0(%arg0: !llvm.ptr<i32>): 
    %1 = llvm.mlir.constant(10 : i32) : i32
    llvm.store %1, %arg0 : !llvm.ptr<i32>
    omp.terminator
  }
  llvm.return
}

[TIFitis]

Addressed reviewer comments. Added missing LoadInst for use_device_addr inside region block.

[TIFitis]

In D146648#4527573, @kiranchandramohan wrote:

While looking into this patch, I noticed that the current implementation in Semantics and MLIR requires the MAP clause to always be present. But the standard only says the following,

At least one map, use_device_addr or use_device_ptr clause must appear on the directive.

I will fix this. But this made having a look at the code generated only with use_device_addr, use_device_ptr and without map difficult.

The test openmp_target_use_dev_addr_nomap does not provide any map clause for %a, so for %a marked with use_device_addr it has the same effect as no map clause being present. Key change when use_device_addr or use_device_ptr is used without providing any map information is that the map_size should be set to 0, and only the OMP_MAP_RETURN_PARAM flag should be set for the map_type.

[TIFitis]

BTW I am not sure how to resolve the windows build error that I am getting repeatedly. It's giving an unknown identifier error for the fail functor being called at OpenMPToLLVMIRTranslation.cpp:1429. Any insight would be much appreciated.

mlir/test/Target/LLVMIR/omptarget-llvm.mlir

193

I've added a openmp_target_use_dev_addr_no_ptr test

[kiranchandramohan]

BTW I am not sure how to resolve the windows build error that I am getting repeatedly. It's giving an unknown identifier error for the fail functor being called at OpenMPToLLVMIRTranslation.cpp:1429. Any insight would be much appreciated.

Might be a bug in the compiler in windows. You are trying to call a lambda function inside a lambda function here. Try moving the second lambda into a static function or try capturing the second lambda explicitly in the first lambda.

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

1374–1376

Can you move this to just before the loop at 1406?

[TIFitis]

Addressed reviewer comments. Try to fix build error.

[TIFitis]

In D146648#4531472, @kiranchandramohan wrote:

BTW I am not sure how to resolve the windows build error that I am getting repeatedly. It's giving an unknown identifier error for the fail functor being called at OpenMPToLLVMIRTranslation.cpp:1429. Any insight would be much appreciated.

Might be a bug in the compiler in windows. You are trying to call a lambda function inside a lambda function here. Try moving the second lambda into a static function or try capturing the second lambda explicitly in the first lambda.

Thanks for the suggestion, explicitly capturing the lambda worked. The builds are all clean now.

[TIFitis]

Ping for review :)

[kiranchandramohan]

use_device_ptr: The mapped variables marked with use_device_ptr are accessed through new pointer allocas inside the region.
use_device_addr: The mapped variables marked with use_device_addr are accessed through the base pointer mappers.

This explanation in the summary talks of mapped variables. Mapping is always not necessary. Could you also expand the explanation a little more? And also talk a bit about the difference between use_device_ptr and use_device_addr.

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

1517

If you do not need any codegen in this case for other switch options as well, then you can consider moving the if outside the switch.

1531–1532

Why is the load required here and not for devPtr?

[TIFitis]

In D146648#4554063, @kiranchandramohan wrote:

use_device_ptr: The mapped variables marked with use_device_ptr are accessed through new pointer allocas inside the region.
use_device_addr: The mapped variables marked with use_device_addr are accessed through the base pointer mappers.

This explanation in the summary talks of mapped variables. Mapping is always not necessary. Could you also expand the explanation a little more? And also talk a bit about the difference between use_device_ptr and use_device_addr.

The regular behaviour is for mapping to be present. In the event that mapping information is absent we simply don't have the map_type and size information available. As such according to the specification the size gets set to 0 and the map_type bits are omitted, i.e, default alloc/release is used as appropriate.

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

1517

No, we need the if condition here in combination with the switch case to generate code for device privatisation.

1531–1532

In case of dev_ptr the load inst is added by OMPIRBuilder itself along with a store Inst to store it into the new temp ptr var added.

For dev_addr we don't add the LoadInst in the OMPIRBuilder only because of compatibility issues with Clang.

[kiranchandramohan]

Looks OK. It will be good to have one more review.

[raghavendhra]

LGTM

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

public:
   struct MapInfo {
     Constant *Name = nullptr;
     uint64_t Type = 0x00;
     uint64_t Size = 0;
     Value *BasePtr = nullptr;
     Value *Ptr = nullptr;
   };
 
+  /// This structure contains information generated for use_deviceptr/addr
+  /// clause operands including the alloca pointers, base pointers and their
+  /// use sets.
+  struct UseDeviceInfo {
+    Value *AllocaPtr = nullptr;
+    Value *BasePtr = nullptr;
+    SmallDenseSet<Use *> Uses;
+  };
+
   /// Generator for '#omp target data'
   ///
   /// \param Loc The location where the target data construct was encountered.
   /// \param AllocaIP The insertion points to be used for alloca instructions.
   /// \param CodeGenIP The insertion point at which the target directive code
   /// should be placed.
   /// \param IsBegin If true then emits begin mapper call otherwise emits
   /// end mapper call.
   /// \param DeviceID Stores the DeviceID from the device clause.
   /// \param IfCond Value which corresponds to the if clause condition.
   /// \param MapInfos Map from a 'map' Value* to its corresponding MapInfo.
+  /// \param UseDeviceInfos Map from a use device pointer Value* to its
+  /// corresponding UseDeviceInfo.
   /// \param BodyGenCB Callback that will generate the region body code.
-  InsertPointTy createTargetData(const LocationDescription &Loc,
-                                 InsertPointTy AllocaIP,
-                                 InsertPointTy CodeGenIP, bool IsBegin,
-                                 int64_t DeviceID, Value *IfCond,
-                                 SmallMapVector<Value *, MapInfo, 4> &MapInfos,
-                                 BodyGenCallbackTy BodyGenCB = {});
+  InsertPointTy
+  createTargetData(const LocationDescription &Loc, InsertPointTy AllocaIP,
+                   InsertPointTy CodeGenIP, bool IsBegin, int64_t DeviceID,
+                   Value *IfCond, SmallMapVector<Value *, MapInfo, 4> &MapInfos,
+                   SmallMapVector<Value *, UseDeviceInfo, 4> &UseDeviceInfos,
 
   /// Declarations for LLVM-IR types (simple, array, function and structure) are
   /// generated below. Their names are defined and used in OpenMPKinds.def. Here
   /// we provide the declarations, the initializeTypes function will provide the
   /// values.
   ///
   ///{
 #define OMP_TYPE(VarName, InitValue) Type *VarName = nullptr;
private:
 
   /// Generator for map clauses
   ///
   /// \param Loc The location where the target data construct was encountered.
   /// \param AllocaIP The insertion points to be used for alloca instructions.
   /// \param CodeGenIP The insertion point at which the target directive code
   /// should be placed.
   /// \param MapInfos Map from a 'map' Value* to its corresponding MapInfo.
+  /// \param UseDeviceInfos Map from a use device pointer Value* to its
+  /// corresponding UseDeviceInfo.
   /// \param BodyGenCB Callback that will generate the region body code.
   /// \param MapperAllocas Return param for the Alloca pointers from
   /// BasePtr, Ptr and Size ptr.
   /// \param MapNamesGV Return param for the MapNames Global.
   /// \param MapTypesGV Return param for the MapTypes Global.
   /// \param MapSizesGV Return param for the MapSizes Global.
   void ProcessMapInfo(const LocationDescription &Loc, InsertPointTy AllocaIP,
                       InsertPointTy CodeGenIP,
                       SmallMapVector<Value *, MapInfo, 4> &MapInfos,
+                      SmallMapVector<Value *, UseDeviceInfo, 4> &UseDeviceInfos,
                       struct MapperAllocas &MapperAllocas, Value **MapNamesGV,
                       Value **MapTypesGV, Value **MapSizesGV);
 
+  /// Generator for use device ptr & addr clauses
+  ///
+  /// \param UseDeviceInfos Map from a use device pointer Value* to its
+  /// corresponding UseDeviceInfo.
+  void ProcessUseDeviceInfo(
+      SmallMapVector<Value *, UseDeviceInfo, 4> &UseDeviceInfos);
+
 private:
   enum AtomicKind { Read, Write, Update, Capture, Compare };
 
   /// Determine whether to emit flush or not
   ///
   /// \param Loc    The insert and source location description.
   /// \param AO     The required atomic ordering
   /// \param AK     The OpenMP atomic operation kind used.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Constant *OpenMPIRBuilder::registerTargetRegionFunction(
   auto OutlinedFnID = createOutlinedFunctionID(OutlinedFn, EntryFnIDName);
   auto EntryAddr = createTargetRegionEntryAddr(OutlinedFn, EntryFnName);
   OffloadInfoManager.registerTargetRegionEntryInfo(
       EntryInfo, EntryAddr, OutlinedFnID,
       OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion);
   return OutlinedFnID;
 }
 
+void OpenMPIRBuilder::ProcessUseDeviceInfo(
+    SmallMapVector<Value *, UseDeviceInfo, 4> &UseDeviceInfos) {
+  for (auto &UseDeviceOp : UseDeviceInfos) {
+    auto *UseDeviceVal = UseDeviceOp.first;
+    auto &UseDeviceInfo = UseDeviceOp.second;
+    auto *LoadInst =
+        Builder.CreateLoad(Builder.getPtrTy(), UseDeviceInfo.BasePtr);
+    Builder.CreateStore(LoadInst, UseDeviceInfo.AllocaPtr);
+
+    for (auto &Use : UseDeviceVal->uses())
+      UseDeviceInfo.Uses.insert(&Use);
+  }
+}
+
 void OpenMPIRBuilder::ProcessMapInfo(
     const LocationDescription &Loc, InsertPointTy AllocaIP,
     InsertPointTy CodeGenIP, SmallMapVector<Value *, MapInfo, 4> &MapInfos,
+    SmallMapVector<Value *, UseDeviceInfo, 4> &UseDeviceInfos,
     struct MapperAllocas &MapperAllocas, Value **MapNamesArg,
     Value **MapTypesArg, Value **MapSizesArg) {
   auto NumMapOperands = MapInfos.size();
 
   createMapperAllocas(Loc, AllocaIP, NumMapOperands, MapperAllocas, false);
 
+  Builder.restoreIP(AllocaIP);
+
+  for (auto &UseDeviceOp : UseDeviceInfos)
+    UseDeviceOp.second.AllocaPtr = Builder.CreateAlloca(Builder.getPtrTy());
+
   Builder.restoreIP(CodeGenIP);
 
   ArrayType *ArrPtrTy = ArrayType::get(Builder.getPtrTy(), NumMapOperands);
 
   unsigned Index = 0;
   for (auto &MapOp : MapInfos) {
     auto &MapInfo = MapOp.second;
     Value *MapOpPtrBase, *MapOpPtr;
Value *GEPPtr = Builder.CreateInBoundsGEP(
         {Builder.getInt32(0), Builder.getInt32(Index)});
     Builder.CreateStore(MapOpPtr, GEPPtr);
 
     MapInfo.BasePtr = BaseGEPPtr;
     MapInfo.Ptr = GEPPtr;
     Index++;
   }
 
+  for (auto &UseDeviceOp : UseDeviceInfos)
+    UseDeviceOp.second.BasePtr = MapInfos[UseDeviceOp.first].BasePtr;
+
   SmallVector<Constant *, 4> MapNames;
   SmallVector<uint64_t, 4> MapTypes;
   SmallVector<uint64_t, 4> MapSizes;
 
   for (const auto &MapInfoIter : MapInfos) {
     const auto &MapInfo = MapInfoIter.second;
     MapNames.emplace_back(MapInfo.Name);
     MapTypes.emplace_back(MapInfo.Type);
*MapSizesArg = Builder.CreateConstInBoundsGEP2_32(
       ArrayType::get(Builder.getInt64Ty(), NumMapOperands), MapSizesGV,
       /*Idx0=*/0, /*Idx1=*/0);
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetData(
     const LocationDescription &Loc, InsertPointTy AllocaIP,
     InsertPointTy CodeGenIP, bool IsBegin, int64_t DeviceID, Value *IfCond,
     SmallMapVector<Value *, MapInfo, 4> &MapInfos,
+    SmallMapVector<Value *, UseDeviceInfo, 4> &UseDeviceInfos,
     BodyGenCallbackTy BodyGenCB) {
   if (!updateToLocation(Loc))
     return InsertPointTy();
 
   Builder.restoreIP(CodeGenIP);
 
   // LLVM utilities like blocks with terminators.
   // The UI acts as a resume point for code insertion after the BodyGen
   auto *UI = Builder.CreateUnreachable();
   if (IfCond) {
     auto *ThenTI =
         SplitBlockAndInsertIfThen(IfCond, UI, /* Unreachable */ false);
     ThenTI->getParent()->setName("omp_if.then");
     Builder.SetInsertPoint(ThenTI);
   } else {
     Builder.SetInsertPoint(UI);
   }
 
   struct MapperAllocas MapperAllocas;
   Value *MapNamesArg = nullptr, *MapTypesArg = nullptr, *MapSizesArg = nullptr;
-  ProcessMapInfo(Loc, AllocaIP, Builder.saveIP(), MapInfos, MapperAllocas,
-                 &MapNamesArg, &MapTypesArg, &MapSizesArg);
+  ProcessMapInfo(Loc, AllocaIP, Builder.saveIP(), MapInfos, UseDeviceInfos,
+                 MapperAllocas, &MapNamesArg, &MapTypesArg, &MapSizesArg);
 
   uint32_t SrcLocStrSize;
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
   Value *srcLocInfo = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
 
   Function *beginMapperFunc =
       getOrCreateRuntimeFunctionPtr(omp::OMPRTL___tgt_target_data_begin_mapper);
   Function *endMapperFunc =
       getOrCreateRuntimeFunctionPtr(omp::OMPRTL___tgt_target_data_end_mapper);
 
   if (BodyGenCB) {
     // Create call to start the data region.
     emitMapperCall(Builder.saveIP(), beginMapperFunc, srcLocInfo, MapTypesArg,
                    MapNamesArg, MapperAllocas, DeviceID, MapInfos.size(),
                    MapSizesArg);
 
+    ProcessUseDeviceInfo(UseDeviceInfos);
+
     BodyGenCB(Builder.saveIP(), Builder.saveIP());
 
+    for (auto &UseDeviceOp : UseDeviceInfos) {
+      auto *UseDeviceVal = UseDeviceOp.first;
+      auto &UseDeviceInfo = UseDeviceOp.second;
+      SmallDenseSet<Use *> ReplaceUses;
+      for (auto &Use : UseDeviceVal->uses()) {
+        if (UseDeviceInfo.Uses.find(&Use) == UseDeviceInfo.Uses.end())
+          ReplaceUses.insert(&Use);
+      }
+      for (auto &Use : ReplaceUses)
+        Use->set(UseDeviceInfo.AllocaPtr);
+    }
+
     Builder.SetInsertPoint(UI->getParent());
     // Create call to end the data region.
     emitMapperCall(Builder.saveIP(), endMapperFunc, srcLocInfo, MapTypesArg,
                    MapNamesArg, MapperAllocas, DeviceID, MapInfos.size(),
                    MapSizesArg);
   } else {
     emitMapperCall(Builder.saveIP(), IsBegin ? beginMapperFunc : endMapperFunc,
                    srcLocInfo, MapTypesArg, MapNamesArg, MapperAllocas,

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

AllocaInst *Val1 =
       Builder.CreateAlloca(Builder.getInt32Ty(), Builder.getInt64(1));
   ASSERT_NE(Val1, nullptr);
 
   IRBuilder<>::InsertPoint AllocaIP(&F->getEntryBlock(),
                                     F->getEntryBlock().getFirstInsertionPt());
 
   llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::MapInfo, 4>
       MapInfos;
+  llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::UseDeviceInfo, 4>
+      UseDeviceInfos;
 
   uint32_t temp;
   MapInfos[Val1].Name = OMPBuilder.getOrCreateSrcLocStr("unknown", temp);
   MapInfos[Val1].Size = 4;
   MapInfos[Val1].Type = 1;
 
   Builder.restoreIP(OMPBuilder.createTargetData(
       Loc, AllocaIP, Builder.saveIP(), /* IsBegin= */ true, DeviceID,
-      /* IfCond= */ nullptr, MapInfos));
+      /* IfCond= */ nullptr, MapInfos, UseDeviceInfos));
 
   CallInst *TargetDataCall = dyn_cast<CallInst>(&BB->back());
   EXPECT_NE(TargetDataCall, nullptr);
   EXPECT_EQ(TargetDataCall->arg_size(), 9U);
   EXPECT_EQ(TargetDataCall->getCalledFunction()->getName(),
             "__tgt_target_data_begin_mapper");
   EXPECT_TRUE(TargetDataCall->getOperand(1)->getType()->isIntegerTy(64));
   EXPECT_TRUE(TargetDataCall->getOperand(2)->getType()->isIntegerTy(32));
AllocaInst *Val1 =
       Builder.CreateAlloca(Builder.getInt32Ty(), Builder.getInt64(1));
   ASSERT_NE(Val1, nullptr);
 
   IRBuilder<>::InsertPoint AllocaIP(&F->getEntryBlock(),
                                     F->getEntryBlock().getFirstInsertionPt());
 
   llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::MapInfo, 4>
       MapInfos;
+  llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::UseDeviceInfo, 4>
+      UseDeviceInfos;
 
   uint32_t temp;
   MapInfos[Val1].Name = OMPBuilder.getOrCreateSrcLocStr("unknown", temp);
   MapInfos[Val1].Size = 4;
   MapInfos[Val1].Type = 2;
 
   Builder.restoreIP(OMPBuilder.createTargetData(
       Loc, AllocaIP, Builder.saveIP(), /* IsBegin= */ false, DeviceID,
-      /* IfCond= */ nullptr, MapInfos));
+      /* IfCond= */ nullptr, MapInfos, UseDeviceInfos));
 
   CallInst *TargetDataCall = dyn_cast<CallInst>(&BB->back());
   EXPECT_NE(TargetDataCall, nullptr);
   EXPECT_EQ(TargetDataCall->arg_size(), 9U);
   EXPECT_EQ(TargetDataCall->getCalledFunction()->getName(),
             "__tgt_target_data_end_mapper");
   EXPECT_TRUE(TargetDataCall->getOperand(1)->getType()->isIntegerTy(64));
   EXPECT_TRUE(TargetDataCall->getOperand(2)->getType()->isIntegerTy(32));
TEST_F(OpenMPIRBuilderTest, TargetDataRegion) {
   IRBuilder<> Builder(BB);
   OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
   int64_t DeviceID = 2;
 
   AllocaInst *Val1 =
       Builder.CreateAlloca(Builder.getInt32Ty(), Builder.getInt64(1));
   ASSERT_NE(Val1, nullptr);
 
+  AllocaInst *Val2 = Builder.CreateAlloca(Builder.getInt32Ty()->getPointerTo());
+  ASSERT_NE(Val2, nullptr);
+
   IRBuilder<>::InsertPoint AllocaIP(&F->getEntryBlock(),
                                     F->getEntryBlock().getFirstInsertionPt());
 
   llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::MapInfo, 4>
       MapInfos;
+  llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::UseDeviceInfo, 4>
+      UseDeviceInfos;
 
   uint32_t temp;
   MapInfos[Val1].Name = OMPBuilder.getOrCreateSrcLocStr("unknown", temp);
   MapInfos[Val1].Size = 4;
   MapInfos[Val1].Type = 3;
+  MapInfos[Val2].Name = OMPBuilder.getOrCreateSrcLocStr("unknown", temp);
+  MapInfos[Val2].Size = 8;
+  MapInfos[Val2].Type = 67;
+
+  UseDeviceInfos[Val2];
 
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
   auto BodyCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP) {
-    CallInst *TargetDataCall = dyn_cast<CallInst>(BB->back().getPrevNode());
+    CallInst *TargetDataCall = dyn_cast<CallInst>(
+        BB->back().getPrevNode()->getPrevNode()->getPrevNode());
     EXPECT_NE(TargetDataCall, nullptr);
     EXPECT_EQ(TargetDataCall->arg_size(), 9U);
     EXPECT_EQ(TargetDataCall->getCalledFunction()->getName(),
               "__tgt_target_data_begin_mapper");
     EXPECT_TRUE(TargetDataCall->getOperand(1)->getType()->isIntegerTy(64));
     EXPECT_TRUE(TargetDataCall->getOperand(2)->getType()->isIntegerTy(32));
     EXPECT_TRUE(TargetDataCall->getOperand(8)->getType()->isPointerTy());
     Builder.restoreIP(codeGenIP);
-    Builder.CreateStore(Builder.getInt32(99), Val1);
+    Builder.CreateStore(Builder.getInt32(10), Val1);
+    auto *LI = Builder.CreateLoad(Builder.getInt32Ty()->getPointerTo(), Val2);
+    Builder.CreateStore(Builder.getInt32(20), LI);
   };
 
   Builder.restoreIP(OMPBuilder.createTargetData(
       Loc, AllocaIP, Builder.saveIP(), /* IsBegin= */ false, DeviceID,
-      /* IfCond= */ nullptr, MapInfos, BodyCB));
+      /* IfCond= */ nullptr, MapInfos, UseDeviceInfos, BodyCB));
+
+  LoadInst *UseDeviceLoadInst = dyn_cast<LoadInst>(
+      BB->back().getPrevNode()->getPrevNode()->getPrevNode()->getPrevNode());
+  EXPECT_NE(UseDeviceLoadInst->getPointerOperand(), Val2);
+  EXPECT_EQ(UseDeviceLoadInst->getPointerOperand(),
+            UseDeviceInfos[Val2].AllocaPtr);
 
   CallInst *TargetDataCall = dyn_cast<CallInst>(&BB->back());
   EXPECT_NE(TargetDataCall, nullptr);
   EXPECT_EQ(TargetDataCall->arg_size(), 9U);
   EXPECT_EQ(TargetDataCall->getCalledFunction()->getName(),
             "__tgt_target_data_end_mapper");
   EXPECT_TRUE(TargetDataCall->getOperand(1)->getType()->isIntegerTy(64));
   EXPECT_TRUE(TargetDataCall->getOperand(2)->getType()->isIntegerTy(32));

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

int64_t getSizeInBytes(DataLayout &DL, const mlir::Type &type) {
 
   return -1;
 }
 
 /// Gather Type and Name and other information for the map operands.
 static LogicalResult gatherMapInfo(
     LLVM::ModuleTranslation &moduleTranslation, DataLayout &DL,
     const SmallVectorImpl<Value> &mapOperands, const ArrayAttr &mapTypes,
+    const SmallVectorImpl<Value> &useDevicePtrOperands,
     llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::MapInfo, 4>
-        &mapInfos) {
+        &mapInfos,
+    llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::UseDeviceInfo, 4>
+        &useDeviceInfos) {
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
   unsigned index = 0;
   for (const auto &mapOp : mapOperands) {
     if (!mapOp.getType().isa<LLVM::LLVMPointerType>())
       return failure();
     llvm::Value *mapOpValue = moduleTranslation.lookupValue(mapOp);
     mapInfos[mapOpValue].Name =
         mlir::LLVM::createMappingInformation(mapOp.getLoc(), *ompBuilder);
     mapInfos[mapOpValue].Type =
         mapTypes[index].dyn_cast<mlir::IntegerAttr>().getInt();
     mapInfos[mapOpValue].Size = getSizeInBytes(DL, mapOp.getType());
     index++;
   }
+
+  // Look at the use_device_ptr and use_device_addr clauses information and
+  // mark the existing map entries as such. If there is no map information for
+  // an entry in the use_device_ptr and use_device_addr list, we create one
+  // with map type 'alloc' and zero size section.
+  for (const auto &ptrOp : useDevicePtrOperands) {
+    llvm::Value *ptrOpValue = moduleTranslation.lookupValue(ptrOp);
+    if (mapInfos.find(ptrOpValue) == mapInfos.end()) {
+      mapInfos[ptrOpValue].Name =
+          mlir::LLVM::createMappingInformation(ptrOp.getLoc(), *ompBuilder);
+      mapInfos[ptrOpValue].Type = 0x00;
+      mapInfos[ptrOpValue].Size = 0;
+    }
+    mapInfos[ptrOpValue].Type |= static_cast<
+        std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+        llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_RETURN_PARAM);
+    useDeviceInfos[ptrOpValue];
+  }
+
   return success();
 }
 
 static LogicalResult
 convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
                      LLVM::ModuleTranslation &moduleTranslation) {
   unsigned numMapOperands;
   bool isBegin;
   llvm::Value *ifCond = nullptr;
   int64_t deviceID = llvm::omp::OMP_DEVICEID_UNDEF;
-  SmallVector<Value> mapOperands;
+  SmallVector<Value> mapOperands, useDevicePtrOperands;
   ArrayAttr mapTypes;
   llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::MapInfo, 4>
       mapInfos;
+  llvm::SmallMapVector<llvm::Value *, llvm::OpenMPIRBuilder::UseDeviceInfo, 4>
+      useDeviceInfos;
   DataLayout DL = DataLayout(op->getParentOfType<ModuleOp>());
 
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
 
   LogicalResult result =
       llvm::TypeSwitch<Operation *, LogicalResult>(op)
           .Case([&](omp::DataOp dataOp) {
-            if (!dataOp.getUseDeviceAddr().empty() ||
-                !dataOp.getUseDevicePtr().empty())
+            if (!dataOp.getUseDeviceAddr().empty())
               return failure();
 
             if (auto ifExprVar = dataOp.getIfExpr())
               ifCond = moduleTranslation.lookupValue(ifExprVar);
 
             if (auto devId = dataOp.getDevice())
               if (auto constOp = mlir::dyn_cast<mlir::LLVM::ConstantOp>(
                       devId.getDefiningOp()))
                 if (auto intAttr =
                         constOp.getValue().dyn_cast<mlir::IntegerAttr>())
                   deviceID = intAttr.getInt();
 
             numMapOperands = dataOp.getMapOperands().size();
             mapOperands = dataOp.getMapOperands();
             mapTypes = dataOp.getMapTypes();
+            useDevicePtrOperands = dataOp.getUseDevicePtr();
             return success();
           })
           .Case([&](omp::EnterDataOp enterDataOp) {
             if (enterDataOp.getNowait())
               return failure();
 
             if (auto ifExprVar = enterDataOp.getIfExpr())
               ifCond = moduleTranslation.lookupValue(ifExprVar);
LogicalResult result =
           });
 
   if (failed(result))
     return failure();
 
   LogicalResult bodyGenStatus = success();
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
   auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP) {
+    // llvm::SmallDenseMap<Value, llvm::Value*> saveValueMap;
+    // for (const auto &ptrOp : useDevicePtrOperands) {
+    //   llvm::Value *ptrOpValue = moduleTranslation.lookupValue(ptrOp);
+    //   saveValueMap[ptrOp] = ptrOpValue;
+    //   moduleTranslation.mapValue(ptrOp,
+    //   useDeviceInfos[ptrOpValue].AllocaPtr);
+    // }
+
     // DataOp has only one region associated with it.
     auto &region = cast<omp::DataOp>(op).getRegion();
     builder.restoreIP(codeGenIP);
     bodyGenStatus = inlineConvertOmpRegions(region, "omp.data.region", builder,
                                             moduleTranslation);
+
+    // for (const auto &ptrOp : useDevicePtrOperands)
+    //   moduleTranslation.mapValue(ptrOp, saveValueMap[ptrOp]);
   };
 
   if (failed(gatherMapInfo(moduleTranslation, DL, mapOperands, mapTypes,
-                           mapInfos)))
+                           useDevicePtrOperands, mapInfos, useDeviceInfos)))
     return failure();
 
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
       findAllocaInsertPoint(builder, moduleTranslation);
 
   if (isa<omp::DataOp>(op)) {
     builder.restoreIP(ompBuilder->createTargetData(
         ompLoc, allocaIP, builder.saveIP(), isBegin, deviceID, ifCond, mapInfos,
-        bodyCB));
+        useDeviceInfos, bodyCB));
   } else {
-    builder.restoreIP(ompBuilder->createTargetData(ompLoc, allocaIP,
-                                                   builder.saveIP(), isBegin,
-                                                   deviceID, ifCond, mapInfos));
+    builder.restoreIP(ompBuilder->createTargetData(
+        ompLoc, allocaIP, builder.saveIP(), isBegin, deviceID, ifCond, mapInfos,
+        useDeviceInfos));
   }
 
   return bodyGenStatus;
 }
 
 namespace {
 
 /// Implementation of the dialect interface that converts operations belonging
 /// to the OpenMP dialect to LLVM IR.

mlir/test/Target/LLVMIR/omptarget-llvm.mlir

 // CHECK:         store ptr %[[ARG_1]], ptr %[[VAL_29]], align 8
 // CHECK:         %[[VAL_30:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_0]], i32 0, i32 0
 // CHECK:         %[[VAL_31:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_1]], i32 0, i32 0
 // CHECK:         call void @__tgt_target_data_end_mapper(ptr @3, i64 -1, i32 2, ptr %[[VAL_30]], ptr %[[VAL_31]], ptr @.offload_sizes.3, ptr @.offload_maptypes.2, ptr @.offload_mapnames.1, ptr null)
 // CHECK:         br label %[[VAL_25]]
 // CHECK:       omp_if.end4:                                      ; preds = %[[VAL_12]], %[[VAL_24]]
 // CHECK:         ret void
 
 // -----
+
+llvm.func @_QPopenmp_target_use_dev_ptr() {
+  %0 = llvm.mlir.constant(1 : i64) : i64
+  %a = llvm.alloca %0 x !llvm.ptr<!llvm.ptr<i32>> : (i64) -> !llvm.ptr<!llvm.ptr<i32>>
+  %1 = llvm.mlir.constant(1 : i64) : i64
+  %b = llvm.alloca %0 x !llvm.ptr<!llvm.ptr<i32>> : (i64) -> !llvm.ptr<!llvm.ptr<i32>>
+  omp.target_data   use_device_ptr(%a : !llvm.ptr<!llvm.ptr<i32>>) map((from -> %a : !llvm.ptr<!llvm.ptr<i32>>)) {
+    %3 = llvm.mlir.constant(10 : i32) : i32
+    %4 = llvm.load %a : !llvm.ptr<!llvm.ptr<i32>>
+    llvm.store %3, %4 : !llvm.ptr<i32>
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK:         @.offload_maptypes = private unnamed_addr constant [1 x i64] [i64 66]
+// CHECK:         @.offload_sizes = private unnamed_addr constant [1 x i64]  [i64 8]
+// CHECK-LABEL: define void @_QPopenmp_target_use_dev_ptr
+// CHECK:         %[[VAL_0:.*]] = alloca [1 x ptr], align 8
+// CHECK:         %[[VAL_1:.*]] = alloca [1 x ptr], align 8
+// CHECK:         %[[VAL_2:.*]] = alloca ptr, align 8
+// CHECK:         %[[VAL_3:.*]] = alloca ptr, i64 1, align 8
+// CHECK:         %[[VAL_4:.*]] = alloca ptr, i64 1, align 8
+// CHECK:         br label %[[VAL_5:.*]]
+// CHECK:       entry:                                            ; preds = %[[VAL_6:.*]]
+// CHECK:         %[[VAL_7:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_0]], i32 0, i32 0
+// CHECK:         store ptr %[[VAL_3]], ptr %[[VAL_7]], align 8
+// CHECK:         %[[VAL_8:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_1]], i32 0, i32 0
+// CHECK:         store ptr %[[VAL_3]], ptr %[[VAL_8]], align 8
+// CHECK:         %[[VAL_9:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_0]], i32 0, i32 0
+// CHECK:         %[[VAL_10:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_1]], i32 0, i32 0
+// CHECK:         call void @__tgt_target_data_begin_mapper(ptr @2, i64 -1, i32 1, ptr %[[VAL_9]], ptr %[[VAL_10]], ptr @.offload_sizes, ptr @.offload_maptypes, ptr @.offload_mapnames, ptr null)
+// CHECK:         %[[VAL_11:.*]] = load ptr, ptr %[[VAL_7]], align 8
+// CHECK:         store ptr %[[VAL_11]], ptr %[[VAL_2]], align 8
+// CHECK:         %[[VAL_12:.*]] = load ptr, ptr %[[VAL_2]], align 8
+// CHECK:         store i32 10, ptr %[[VAL_12]], align 4
+// CHECK:         %[[VAL_13:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_0]], i32 0, i32 0
+// CHECK:         %[[VAL_14:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_1]], i32 0, i32 0
+// CHECK:         call void @__tgt_target_data_end_mapper(ptr @2, i64 -1, i32 1, ptr %[[VAL_13]], ptr %[[VAL_14]], ptr @.offload_sizes, ptr @.offload_maptypes, ptr @.offload_mapnames, ptr null)
+// CHECK:         ret void
+
+// -----
+
+
+
+llvm.func @_QPopenmp_target_use_dev_ptr_nomap() {
+  %0 = llvm.mlir.constant(1 : i64) : i64
+  %a = llvm.alloca %0 x !llvm.ptr<!llvm.ptr<i32>> : (i64) -> !llvm.ptr<!llvm.ptr<i32>>
+  %1 = llvm.mlir.constant(1 : i64) : i64
+  %b = llvm.alloca %0 x !llvm.ptr<!llvm.ptr<i32>> : (i64) -> !llvm.ptr<!llvm.ptr<i32>>
+  omp.target_data   use_device_ptr(%a : !llvm.ptr<!llvm.ptr<i32>>) map((from -> %b : !llvm.ptr<!llvm.ptr<i32>>)) {
+    %3 = llvm.mlir.constant(10 : i32) : i32
+    %4 = llvm.load %a : !llvm.ptr<!llvm.ptr<i32>>
+    llvm.store %3, %4 : !llvm.ptr<i32>
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK:         @.offload_maptypes = private unnamed_addr constant [2 x i64] [i64 2, i64 64]
+// CHECK:         @.offload_sizes = private unnamed_addr constant [2 x i64]  [i64 8, i64 0]
+// CHECK-LABEL: define void @_QPopenmp_target_use_dev_ptr
+// CHECK:         %[[VAL_0:.*]] = alloca [2 x ptr], align 8
+// CHECK:         %[[VAL_1:.*]] = alloca [2 x ptr], align 8
+// CHECK:         %[[VAL_2:.*]] = alloca ptr, align 8
+// CHECK:         %[[VAL_3:.*]] = alloca ptr, i64 1, align 8
+// CHECK:         %[[VAL_4:.*]] = alloca ptr, i64 1, align 8
+// CHECK:         br label %[[VAL_5:.*]]
+// CHECK:       entry:                                            ; preds = %[[VAL_6:.*]]
+// CHECK:         %[[VAL_7:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_0]], i32 0, i32 0
+// CHECK:         store ptr %[[VAL_4]], ptr %[[VAL_7]], align 8
+// CHECK:         %[[VAL_8:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_1]], i32 0, i32 0
+// CHECK:         store ptr %[[VAL_4]], ptr %[[VAL_8]], align 8
+// CHECK:         %[[VAL_9:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_0]], i32 0, i32 1
+// CHECK:         store ptr %[[VAL_3]], ptr %[[VAL_9]], align 8
+// CHECK:         %[[VAL_10:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_1]], i32 0, i32 1
+// CHECK:         store ptr %[[VAL_3]], ptr %[[VAL_10]], align 8
+// CHECK:         %[[VAL_11:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_0]], i32 0, i32 0
+// CHECK:         %[[VAL_12:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_1]], i32 0, i32 0
+// CHECK:         call void @__tgt_target_data_begin_mapper(ptr @3, i64 -1, i32 2, ptr %[[VAL_11]], ptr %[[VAL_12]], ptr @.offload_sizes, ptr @.offload_maptypes, ptr @.offload_mapnames, ptr null)
+// CHECK:         %[[VAL_13:.*]] = load ptr, ptr %[[VAL_9]], align 8
+// CHECK:         store ptr %[[VAL_13]], ptr %[[VAL_2]], align 8
+// CHECK:         %[[VAL_14:.*]] = load ptr, ptr %[[VAL_2]], align 8
+// CHECK:         store i32 10, ptr %[[VAL_14]], align 4
+// CHECK:         %[[VAL_15:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_0]], i32 0, i32 0
+// CHECK:         %[[VAL_16:.*]] = getelementptr inbounds [2 x ptr], ptr %[[VAL_1]], i32 0, i32 0
+// CHECK:         call void @__tgt_target_data_end_mapper(ptr @3, i64 -1, i32 2, ptr %[[VAL_15]], ptr %[[VAL_16]], ptr @.offload_sizes, ptr @.offload_maptypes, ptr @.offload_mapnames, ptr null)
+// CHECK:         ret void
+
+// -----