[Mlir-commits] [openmp] [mlir] [flang][OpenMP][Offloading][AMDGPU] Add test for `target update` (PR #76355)

[Kareem Ergawy]

https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/76355

Adds a new test for offloading `target update` directive to AMD GPUs.

>From 4c38828f0342f2fc8db8cfc378103aa16433c9f1 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Mon, 25 Dec 2023 02:17:48 -0600
Subject: [PATCH] [flang][OpenMP][Offloading][AMDGPU] Add offloading test for
 `target update`

Adds a new test for offloading `target update` directive to AMD GPUs.
---
 .../Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp  | 11 ++---
 openmp/libomptarget/test/lit.cfg              |  2 +-
 .../test/offloading/fortran/target_update.f90 | 43 +++++++++++++++++++
 3 files changed, 50 insertions(+), 6 deletions(-)
 create mode 100644 openmp/libomptarget/test/offloading/fortran/target_update.f90

diff --git a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
index cd1cfb3b7686d0..730858ffc67a71 100644
--- a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
+++ b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
@@ -239,11 +239,11 @@ void mlir::configureOpenMPToLLVMConversionLegality(
   target.addDynamicallyLegalOp<
       mlir::omp::AtomicReadOp, mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
       mlir::omp::ThreadprivateOp, mlir::omp::YieldOp, mlir::omp::EnterDataOp,
-      mlir::omp::ExitDataOp, mlir::omp::DataBoundsOp, mlir::omp::MapInfoOp>(
-      [&](Operation *op) {
-        return typeConverter.isLegal(op->getOperandTypes()) &&
-               typeConverter.isLegal(op->getResultTypes());
-      });
+      mlir::omp::ExitDataOp, mlir::omp::UpdateDataOp, mlir::omp::DataBoundsOp,
+      mlir::omp::MapInfoOp>([&](Operation *op) {
+    return typeConverter.isLegal(op->getOperandTypes()) &&
+           typeConverter.isLegal(op->getResultTypes());
+  });
   target.addDynamicallyLegalOp<mlir::omp::ReductionOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getOperandTypes());
   });
@@ -282,6 +282,7 @@ void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
       RegionLessOpConversion<omp::YieldOp>,
       RegionLessOpConversion<omp::EnterDataOp>,
       RegionLessOpConversion<omp::ExitDataOp>,
+      RegionLessOpConversion<omp::UpdateDataOp>,
       RegionLessOpWithVarOperandsConversion<omp::DataBoundsOp>>(converter);
 }
 
diff --git a/openmp/libomptarget/test/lit.cfg b/openmp/libomptarget/test/lit.cfg
index 19c5e5c4572227..dca922dcfc6bce 100644
--- a/openmp/libomptarget/test/lit.cfg
+++ b/openmp/libomptarget/test/lit.cfg
@@ -132,7 +132,7 @@ elif config.operating_system == 'Darwin':
     config.test_flags += " -Wl,-rpath," + config.library_dir
     config.test_flags += " -Wl,-rpath," + config.omp_host_rtl_directory
 else: # Unices
-    if config.libomptarget_current_target != "nvptx64-nvidia-cuda":
+    if config.libomptarget_current_target != "nvptx64-nvidia-cuda" and config.libomptarget_current_target != "amdgcn-amd-amdhsa":
         config.test_flags += " -nogpulib"
     config.test_flags += " -Wl,-rpath," + config.library_dir
     config.test_flags += " -Wl,-rpath," + config.omp_host_rtl_directory
diff --git a/openmp/libomptarget/test/offloading/fortran/target_update.f90 b/openmp/libomptarget/test/offloading/fortran/target_update.f90
new file mode 100644
index 00000000000000..87f0a46f3afe72
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/fortran/target_update.f90
@@ -0,0 +1,43 @@
+! REQUIRES: flang, amdgcn-amd-amdhsa
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+program target_update
+    implicit none
+    integer :: x(1)
+    integer :: host_id
+    integer :: device_id(1)
+
+    INTERFACE
+        FUNCTION omp_get_device_num() BIND(C)
+            USE, INTRINSIC :: iso_c_binding, ONLY: C_INT
+            integer :: omp_get_device_num
+        END FUNCTION omp_get_device_num
+    END INTERFACE
+
+    x(1) = 5
+    host_id = omp_get_device_num()
+
+!$omp target enter data map(to:x, device_id)
+!$omp target
+    x(1) = 42
+!$omp end target
+
+    ! CHECK: After first target regions and before target update: x = 5
+    print *, "After first target regions and before target update: x =", x(1)
+
+!$omp target
+    x(1) = 84
+    device_id(1) = omp_get_device_num()
+!$omp end target
+!$omp target update from(x, device_id)
+
+    ! CHECK: After second target regions and target update: x = 84
+    print *, "After second target regions and target update: x =", x(1)
+
+    ! CHECK: Offloading succeeded!
+    if (host_id /= device_id(1)) then
+        print *, "Offloading succeeded!"
+    else
+        print *, "Offloading failed!"
+    end if
+end program target_update