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

[Kareem Ergawy via Openmp-commits]

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

>From 674f5b494f3c94b9442a79033a72a6e6cc5b5264 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 | 50 +++++++++++++++++++
 3 files changed, 57 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..fb35c5a36ab0e5
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/fortran/target_update.f90
@@ -0,0 +1,50 @@
+! Offloading test for the `target update` directive.
+
+! 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
+
+    ! Test that without a `target update` directive, the target update to x is
+    ! not yet seen by the host.
+    ! 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)
+
+    ! Test that after the `target update`, the host can see the new x value.
+    ! CHECK: After second target regions and target update: x = 84
+    print *, "After second target regions and target update: x =", x(1)
+
+    ! Make sure that offloading to the device actually happened. This way we
+    ! verify that we didn't take the fallback host execution path.
+    ! CHECK: Offloading succeeded!
+    if (host_id /= device_id(1)) then
+        print *, "Offloading succeeded!"
+    else
+        print *, "Offloading failed!"
+    end if
+end program target_update