[MLIR][OpenMP] Add OpenMPToLLVMIRTranslation support for is_device_ptr

This PR adds support for the OpenMP is_device_ptr clause in the MLIR to LLVM IR translation for target regions. The is_device_ptr clause allows device pointers (allocated via OpenMP runtime APIs) to be used directly in target regions without implicit mapping.

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,6 @@ 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 checkLinear = [&todo](auto op, LogicalResult &result) {
     if (!op.getLinearVars().empty() || !op.getLinearStepVars().empty())
       result = todo("linear");
@@ -444,7 +440,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
         checkBare(op, result);
         checkDevice(op, result);
         checkInReduction(op, result);
-        checkIsDevicePtr(op, result);
       })
       .Default([](Operation &) {
         // Assume all clauses for an operation can be translated unless they are
@@ -3875,6 +3870,11 @@ convertClauseMapFlags(omp::ClauseMapFlags mlirFlags) {
   if (mapTypeToBool(omp::ClauseMapFlags::attach))
     mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ATTACH;
 
+  if (mapTypeToBool(omp::ClauseMapFlags::is_device_ptr)) {
+    mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TARGET_PARAM;
+    mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_LITERAL;
+  }
+
   return mapType;
 }
 
@@ -3996,6 +3996,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::is_device_ptr) !=
+        omp::ClauseMapFlags::none;
 
     mapData.OriginalValue.push_back(origValue);
     mapData.BasePointers.push_back(origValue);
@@ -4029,7 +4032,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/omptarget-llvm.mlir

@@ -622,3 +622,20 @@ module attributes {omp.target_triples = ["amdgcn-amd-amdhsa"]} {
 // CHECK:         br label %[[VAL_40]]
 // CHECK:       omp.done:                                         ; preds = %[[VAL_68]], %[[VAL_63]], %[[VAL_32]]
 // CHECK:         ret void
+
+// -----
+
+module attributes {omp.target_triples = ["amdgcn-amd-amdhsa"]} {
+  llvm.func @_QPomp_target_is_device_ptr(%arg0 : !llvm.ptr) {
+    %map = omp.map.info var_ptr(%arg0 : !llvm.ptr, !llvm.ptr)
+        map_clauses(is_device_ptr) capture(ByRef) -> !llvm.ptr {name = ""}
+    omp.target map_entries(%map -> %ptr_arg : !llvm.ptr) {
+      omp.terminator
+    }
+    llvm.return
+  }
+}
+
+// CHECK: @.offload_sizes = private unnamed_addr constant [1 x i64] [i64 8]
+// CHECK: @.offload_maptypes = private unnamed_addr constant [1 x i64] [i64 288]
+// CHECK-LABEL: define void @_QPomp_target_is_device_ptr

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,46 @@
+! Validate that a device pointer obtained via omp_get_mapped_ptr can be used
+! inside a TARGET region with the is_device_ptr clause.
+! REQUIRES: flang, amdgcn-amd-amdhsa
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+
+program is_device_ptr_target
+  use iso_c_binding, only : c_ptr, c_loc
+  implicit none
+
+  interface
+    function omp_get_mapped_ptr(host_ptr, device_num)                       &
+        bind(C, name="omp_get_mapped_ptr")
+      use iso_c_binding, only : c_ptr, c_int
+      type(c_ptr) :: omp_get_mapped_ptr
+      type(c_ptr), value :: host_ptr
+      integer(c_int), value :: device_num
+    end function omp_get_mapped_ptr
+  end interface
+
+  integer, parameter :: n = 4
+  integer, parameter :: dev = 0
+  integer, target :: a(n)
+  type(c_ptr) :: dptr
+  integer :: flag
+
+  a = [2, 4, 6, 8]
+  flag = 0
+
+  !$omp target data map(tofrom: a, flag)
+    dptr = omp_get_mapped_ptr(c_loc(a), dev)
+
+    !$omp target is_device_ptr(dptr) map(tofrom: flag)
+      flag = flag + 1
+    !$omp end target
+  !$omp end target data
+
+  if (flag .eq. 1 .and. all(a == [2, 4, 6, 8])) then
+    print *, "PASS"
+  else
+    print *, "FAIL", a
+  end if
+
+end program is_device_ptr_target
+
+!CHECK: PASS