From 794bf300693e7826a7e2b7d50e6823d9dd421d4e Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy@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 ++--
 .../test/offloading/fortran/target_update.f90 | 50 +++++++++++++++++++
 2 files changed, 56 insertions(+), 5 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 cd1cfb3b7686d..730858ffc67a7 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/offloading/fortran/target_update.f90 b/openmp/libomptarget/test/offloading/fortran/target_update.f90
new file mode 100644
index 0000000000000..fb35c5a36ab0e
--- /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