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

[Dominik Adamski via flang-commits]

================
@@ -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 =
----------------
DominikAdamski wrote:

Done

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