From 44a91e839ed7656146647c16613d9ee86f1d9b31 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler@amd.com>
Date: Fri, 24 Oct 2025 13:52:28 -0500
Subject: [PATCH] [OpenMP][offload] Fix map-type order

Reorder mappers such that
- on target entry: "from" mappers are handled at the end, where they
  only act as decayed alloc/storage mappers.
- on target exit: "to" mappers are handled first so that they can act as
  decayed release/storage mappers and decrement the reference count.

This avoids that map to+from or from+to result in different outcomes
than mapping with a tofrom mapper. The previous behavior:
- `target map(to: ptr[0:size]) map(from: ptr[0:size])`: the from-mapper
  was hindered from copying the data back to the host because it was
handled first on target exit which means that the to-mapper hadn't been
able to decrement the reference count first.
- `target map(from: ptr[0:size]) map(to: ptr[0:size])`: the to-mapper
  was hindered from copying the data to the device because it was
handled second on target entry which means that the reference count had
already been incremented by the decayed alloc/storage operation of the
from-mapper.
---
 offload/libomptarget/omptarget.cpp      | 28 ++++++++++++-
 offload/test/mapping/map_type_order.cpp | 56 +++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 2 deletions(-)
 create mode 100644 offload/test/mapping/map_type_order.cpp

diff --git a/offload/libomptarget/omptarget.cpp b/offload/libomptarget/omptarget.cpp
index 1753917667fb4..99f31d39668fb 100644
--- a/offload/libomptarget/omptarget.cpp
+++ b/offload/libomptarget/omptarget.cpp
@@ -483,6 +483,26 @@ static int performPointerAttachment(DeviceTy &Device, AsyncInfoTy &AsyncInfo,
   return HandleSubmitResult(SubmitResult);
 }
 
+/// Reorder the mapper indices such that the mappers that are of the specified
+/// map type occur last.
+/// Important: aside from this reordering, the order of the mappers must be
+/// maintained to not violate other ordering requirements.
+static SmallVector<int32_t> reorderMapType(int32_t ArgNum, int64_t *ArgTypes,
+                                           tgt_map_type MapType) {
+  SmallVector<int32_t> Mappers(ArgNum);
+  auto *Front = Mappers.begin();
+  auto *Back = Mappers.end() - 1;
+
+  for (int32_t I = 0; I < ArgNum; ++I) {
+    if (ArgTypes[I] & MapType)
+      *Back-- = I;
+    else
+      *Front++ = I;
+  }
+
+  return Mappers;
+}
+
 /// Internal function to do the mapping and transfer the data to the device
 int targetDataBegin(ident_t *Loc, DeviceTy &Device, int32_t ArgNum,
                     void **ArgsBase, void **Args, int64_t *ArgSizes,
@@ -492,7 +512,9 @@ int targetDataBegin(ident_t *Loc, DeviceTy &Device, int32_t ArgNum,
   assert(AttachInfo && "AttachInfo must be available for targetDataBegin for "
                        "handling ATTACH map-types.");
   // process each input.
-  for (int32_t I = 0; I < ArgNum; ++I) {
+  SmallVector<int32_t> Mappers =
+      reorderMapType(ArgNum, ArgTypes, OMP_TGT_MAPTYPE_FROM);
+  for (int32_t I : Mappers) {
     // Ignore private variables and arrays - there is no mapping for them.
     if ((ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) ||
         (ArgTypes[I] & OMP_TGT_MAPTYPE_PRIVATE))
@@ -1008,7 +1030,9 @@ int targetDataEnd(ident_t *Loc, DeviceTy &Device, int32_t ArgNum,
   int Ret = OFFLOAD_SUCCESS;
   auto *PostProcessingPtrs = new SmallVector<PostProcessingInfo>();
   // process each input.
-  for (int32_t I = ArgNum - 1; I >= 0; --I) {
+  SmallVector<int32_t> Mappers =
+      reorderMapType(ArgNum, ArgTypes, OMP_TGT_MAPTYPE_TO);
+  for (int32_t I : llvm::reverse(Mappers)) {
     // Ignore private variables and arrays - there is no mapping for them.
     // Also, ignore the use_device_ptr directive, it has no effect here.
     if ((ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) ||
diff --git a/offload/test/mapping/map_type_order.cpp b/offload/test/mapping/map_type_order.cpp
new file mode 100644
index 0000000000000..1d5cd93254594
--- /dev/null
+++ b/offload/test/mapping/map_type_order.cpp
@@ -0,0 +1,56 @@
+// RUN: %libomptarget-compilexx-run-and-check-generic
+
+#include <cstdio>
+#include <omp.h>
+
+int main() {
+  int i;
+
+  i = -1;
+#pragma omp target map(to : i) map(from : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(from : i) map(to : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(alloc : i) map(to : i) map(from : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(alloc : i) map(from : i) map(to : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(to : i) map(alloc : i) map(from : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(from : i) map(alloc : i) map(to : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(to : i) map(from : i) map(alloc : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+
+  i = -1;
+#pragma omp target map(from : i) map(to : i) map(alloc : i)
+  i += 2;
+  printf("%d\n", i);
+  // CHECK: 1
+}
\ No newline at end of file