[MLIR][OpenMP] Add OpenMPToLLVMIRTranslation support for is_device_ptr

Author: TIFitis

flang/test/Integration/OpenMP/map-types-and-sizes.f90

@@ -33,6 +33,15 @@ subroutine mapType_array
   !$omp end target
 end subroutine mapType_array
 
+!CHECK: @.offload_sizes{{.*}} = private unnamed_addr constant [1 x i64] [i64 8]
+!CHECK: @.offload_maptypes{{.*}} = private unnamed_addr constant [1 x i64] [i64 288]
+subroutine mapType_is_device_ptr
+  use iso_c_binding, only : c_ptr
+  type(c_ptr) :: p
+  !$omp target is_device_ptr(p)
+  !$omp end target
+end subroutine mapType_is_device_ptr
+
 !CHECK: @.offload_sizes{{.*}} = private unnamed_addr constant [4 x i64] [i64 0, i64 24, i64 8, i64 0]
 !CHECK: @.offload_maptypes{{.*}} = private unnamed_addr constant [4 x i64] [i64 32, i64 281474976711169, i64 281474976711171, i64 281474976711187]
 subroutine mapType_ptr

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

@@ -332,10 +332,7 @@ static LogicalResult checkImplementationStatus(Operation &op) {
         op.getInReductionSyms())
       result = todo("in_reduction");
   };
-  auto checkIsDevicePtr = [&todo](auto op, LogicalResult &result) {
-    if (!op.getIsDevicePtrVars().empty())
-      result = todo("is_device_ptr");
-  };
+  auto checkIsDevicePtr = [](auto, LogicalResult &) {};
   auto checkLinear = [&todo](auto op, LogicalResult &result) {
     if (!op.getLinearVars().empty() || !op.getLinearStepVars().empty())
       result = todo("linear");
@@ -3996,6 +3993,9 @@ static void collectMapDataFromMapOperands(
     llvm::Value *origValue = moduleTranslation.lookupValue(offloadPtr);
     auto mapType = convertClauseMapFlags(mapOp.getMapType());
     auto mapTypeAlways = llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ALWAYS;
+    bool isDevicePtr =
+        (mapOp.getMapType() & omp::ClauseMapFlags::storage) ==
+        omp::ClauseMapFlags::storage;
 
     mapData.OriginalValue.push_back(origValue);
     mapData.BasePointers.push_back(origValue);
@@ -4006,7 +4006,12 @@ static void collectMapDataFromMapOperands(
     mapData.Sizes.push_back(
         builder.getInt64(dl.getTypeSize(mapOp.getVarType())));
     mapData.MapClause.push_back(mapOp.getOperation());
-    if (llvm::to_underlying(mapType & mapTypeAlways)) {
+    if (isDevicePtr) {
+      mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TARGET_PARAM;
+      mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_LITERAL;
+      mapData.Types.push_back(mapType);
+      mapData.Mappers.push_back(nullptr);
+    } else if (llvm::to_underlying(mapType & mapTypeAlways)) {
       // Descriptors are mapped with the ALWAYS flag, since they can get
       // rematerialized, so the address of the decriptor for a given object
       // may change from one place to another.
@@ -4029,7 +4034,8 @@ static void collectMapDataFromMapOperands(
     mapData.Names.push_back(LLVM::createMappingInformation(
         mapOp.getLoc(), *moduleTranslation.getOpenMPBuilder()));
     mapData.DevicePointers.push_back(
-        llvm::OpenMPIRBuilder::DeviceInfoTy::Address);
+        isDevicePtr ? llvm::OpenMPIRBuilder::DeviceInfoTy::Pointer
+                    : llvm::OpenMPIRBuilder::DeviceInfoTy::Address);
     mapData.IsAMapping.push_back(false);
     mapData.IsAMember.push_back(checkIsAMember(hasDevAddrOperands, mapOp));
   }

mlir/test/Target/LLVMIR/openmp-todo.mlir

@@ -238,17 +238,6 @@ llvm.func @target_in_reduction(%x : !llvm.ptr) {
 
 // -----
 
-llvm.func @target_is_device_ptr(%x : !llvm.ptr) {
-  // expected-error@below {{not yet implemented: Unhandled clause is_device_ptr in omp.target operation}}
-  // expected-error@below {{LLVM Translation failed for operation: omp.target}}
-  omp.target is_device_ptr(%x : !llvm.ptr) {
-    omp.terminator
-  }
-  llvm.return
-}
-
-// -----
-
 llvm.func @target_enter_data_depend(%x: !llvm.ptr) {
   // expected-error@below {{not yet implemented: Unhandled clause depend in omp.target_enter_data operation}}
   // expected-error@below {{LLVM Translation failed for operation: omp.target_enter_data}}

offload/test/offloading/fortran/target-is-device-ptr.f90

@@ -0,0 +1,79 @@
+! Validate that a device pointer allocated via OpenMP runtime APIs can be
+! consumed by a TARGET region using the is_device_ptr clause.
+! REQUIRES: flang, amdgcn-amd-amdhsa
+! UNSUPPORTED: nvptx64-nvidia-cuda
+! UNSUPPORTED: nvptx64-nvidia-cuda-LTO
+! UNSUPPORTED: aarch64-unknown-linux-gnu
+! UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+! UNSUPPORTED: x86_64-unknown-linux-gnu
+! UNSUPPORTED: x86_64-unknown-linux-gnu-LTO
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+
+program is_device_ptr_target
+  use omp_lib
+  use iso_c_binding
+  implicit none
+
+  integer, parameter :: n = 4
+  integer, target :: host(n)
+  type(c_ptr) :: device_ptr
+  integer(c_int) :: rc
+  integer :: i
+
+  do i = 1, n
+    host(i) = i
+  end do
+
+  device_ptr = omp_target_alloc(int(n, c_size_t) * int(c_sizeof(host(1)), c_size_t), &
+                                omp_get_default_device())
+  if (.not. c_associated(device_ptr)) then
+    print *, "device alloc failed"
+    stop 1
+  end if
+
+  rc = omp_target_memcpy(device_ptr, c_loc(host), &
+                         int(n, c_size_t) * int(c_sizeof(host(1)), c_size_t), &
+                         0_c_size_t, 0_c_size_t, &
+                         omp_get_default_device(), omp_get_initial_device())
+  if (rc .ne. 0) then
+    print *, "host->device memcpy failed"
+    call omp_target_free(device_ptr, omp_get_default_device())
+    stop 1
+  end if
+
+  call fill_on_device(device_ptr)
+
+  rc = omp_target_memcpy(c_loc(host), device_ptr, &
+                         int(n, c_size_t) * int(c_sizeof(host(1)), c_size_t), &
+                         0_c_size_t, 0_c_size_t, &
+                         omp_get_initial_device(), omp_get_default_device())
+  call omp_target_free(device_ptr, omp_get_default_device())
+
+  if (rc .ne. 0) then
+    print *, "device->host memcpy failed"
+    stop 1
+  end if
+
+  if (all(host == [2, 4, 6, 8])) then
+    print *, "PASS"
+  else
+    print *, "FAIL", host
+  end if
+
+contains
+  subroutine fill_on_device(ptr)
+    type(c_ptr) :: ptr
+    integer, pointer :: p(:)
+    call c_f_pointer(ptr, p, [n])
+
+    !$omp target is_device_ptr(ptr)
+      p(1) = 2
+      p(2) = 4
+      p(3) = 6
+      p(4) = 8
+    !$omp end target
+  end subroutine fill_on_device
+end program is_device_ptr_target
+
+!CHECK: PASS