[OpenMP][Offload] Fix use-after-free of the packed firstprivate transfer buffer (#212905)

`packAndTransfer()` called `submitData()`, which captured the address of
the data of
the `SmallVector` member of `PrivateArgumentManagerTy`, as the host
source of an
async transfer. The manager was later moved into an `AsyncInfo`
post-processing
lambda, changing the underlying address of the data:
`addPostProcessingFunction()`
copies the closure rather than moving it, but even if it "moved", the
address of the data owned by `SmallVector`'s inline capacity would have
changed.

The AMDGPU plugin read from the stale host source address when doing the
actual transfer.

Fix: allocate the packed buffer via
`getOrCreateSourceBufferForSubmitData()`, so
that `AsyncInfo` owns it. The non-packed
corresponding-pointer-initialization
path in this same class already built its source buffer that way.

offload/libomptarget/omptarget.cpp

@@ -61,8 +61,11 @@ int AsyncInfoTy::synchronize() {
   }
 
   // Run any pending post-processing function registered on this async object.
-  if (Result == OFFLOAD_SUCCESS && isQueueEmpty())
+  if (Result == OFFLOAD_SUCCESS && isQueueEmpty()) {
+    ODBG(ODT_DataTransfer)
+        << "Synchronization complete, running post-processing";
     Result = runPostProcessing();
+  }
 
   return Result;
 }
@@ -353,6 +356,8 @@ static char *getOrCreateSourceBufferForSubmitData(AsyncInfoTy &AsyncInfo,
   // Create a dynamic buffer for larger data and schedule its deletion.
   char *DataBuffer = new char[Size];
   AsyncInfo.addPostProcessingFunction([DataBuffer]() {
+    ODBG(ODT_DataTransfer) << "Releasing submitData source buffer at "
+                           << static_cast<void *>(DataBuffer);
     delete[] DataBuffer;
     return OFFLOAD_SUCCESS;
   });
@@ -1692,8 +1697,6 @@ class PrivateArgumentManagerTy {
 
   /// A vector of information of all first-private arguments to be packed
   SmallVector<FirstPrivateArgInfoTy> FirstPrivateArgInfo;
-  /// Host buffer for all arguments to be packed
-  SmallVector<char> FirstPrivateArgBuffer;
   /// The total size of all arguments to be packed
   int64_t FirstPrivateArgSize = 0;
 
@@ -1966,8 +1969,17 @@ public:
     if (!FirstPrivateArgInfo.empty()) {
       assert(FirstPrivateArgSize != 0 &&
              "FirstPrivateArgSize is 0 but FirstPrivateArgInfo is empty");
-      FirstPrivateArgBuffer.resize(FirstPrivateArgSize, 0);
+      // The packed buffer is the host source of the asynchronous submitData
-      auto *Itr = FirstPrivateArgBuffer.begin();
+      // below, so its lifetime must outlive this object rather than be tied to
+      // it.
+      char *FirstPrivateArgBuffer =
+          getOrCreateSourceBufferForSubmitData(AsyncInfo, FirstPrivateArgSize);
+      // Not strictly necessary: the loop below writes every entry, and the
+      // inter-entry padding gaps it skips are never read by the device. Zero
+      // them anyway to keep the transferred buffer deterministic, which makes
+      // dumps easier to read and avoids tripping memory sanitizers.
+      std::memset(FirstPrivateArgBuffer, 0, FirstPrivateArgSize);
+      auto *Itr = FirstPrivateArgBuffer;
       // Copy all host data to this buffer
       for (FirstPrivateArgInfoTy &Info : FirstPrivateArgInfo) {
         // First pad the pointer as we (have to) pad it on the device too.
@@ -1983,7 +1995,7 @@ public:
       }
       // Allocate target memory
       void *TgtPtr =
-          Device.allocData(FirstPrivateArgSize, FirstPrivateArgBuffer.data());
+          Device.allocData(FirstPrivateArgSize, FirstPrivateArgBuffer);
       if (TgtPtr == nullptr) {
         ODBG(ODT_Alloc)
             << "Failed to allocate target memory for private arguments.";
@@ -1993,7 +2005,10 @@ public:
       ODBG(ODT_Alloc) << "Allocated " << FirstPrivateArgSize
                       << " bytes of target memory at " << TgtPtr;
       // Transfer data to target device
-      int Ret = Device.submitData(TgtPtr, FirstPrivateArgBuffer.data(),
+      ODBG(ODT_DataTransfer)
+          << "Submitting packed firstprivate arguments from host buffer at "
+          << static_cast<void *>(FirstPrivateArgBuffer);
+      int Ret = Device.submitData(TgtPtr, FirstPrivateArgBuffer,
                                   FirstPrivateArgSize, AsyncInfo);
       if (Ret != OFFLOAD_SUCCESS) {
         ODBG(ODT_DataTransfer) << "Failed to submit data of private arguments.";
@@ -2008,7 +2023,8 @@ public:
         TP += Info.Padding;
         Ptr = reinterpret_cast<void *>(TP);
         TP += Info.Size;
-        ODBG(ODT_Mapping) << "Firstprivate array " << Info.HstPtrBegin
+        ODBG(ODT_Mapping) << "Firstprivate array "
+                          << static_cast<void *>(Info.HstPtrBegin)
                           << " of size " << (Info.HstPtrEnd - Info.HstPtrBegin)
                           << " mapped to " << Ptr;
       }

offload/test/offloading/firstprivate_packed_buffer_lifetime.c

@@ -0,0 +1,53 @@
+// RUN: %libomptarget-compile-generic
+//
+// RUN: env LIBOMPTARGET_DEBUG=1 LIBOMPTARGET_INFO=32 \
+// RUN: %libomptarget-run-generic 2>&1 \
+// RUN: | %fcheck-generic -check-prefix=DEBUG
+//
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic -check-prefix=CHECK
+//
+// REQUIRES: libomptarget-debug
+
+// Ensure that the packed host buffer used to send firstprivate data to the
+// device is released only after the transfer is done, i.e. that the buffer
+// submitData() reads from is still alive when the queued device operations
+// complete.
+//
+// The firstprivate payload has to be larger than a pointer for the runtime to
+// allocate a separate host buffer for it, and small enough to be packed rather
+// than transferred on its own, so use an explicitly-sized 48-byte array.
+
+#include <stdint.h>
+#include <stdio.h>
+
+#define N 6
+
+int main(void) {
+  uint64_t a[N];
+  for (int i = 0; i < N; ++i)
+    a[i] = i + 1;
+
+  int ok = 0;
+
+  // clang-format off
+  // The packed buffer is what gets submitted as the host source, ...
+  // DEBUG: Firstprivate array 0x{{.*}} of size 48 will be packed
+  // DEBUG: Submitting packed firstprivate arguments from host buffer at 0x[[#%x,BUF:]]
+  // DEBUG: Copying data from host to device, HstPtr=0x{{0*}}[[#BUF]], {{.*}}, Size=48
+  // ... and it is still alive once the queued operations have completed, which
+  // is when the release runs.
+  // DEBUG: Synchronization complete, running post-processing
+  // DEBUG: Releasing submitData source buffer at 0x{{0*}}[[#BUF]]
+  // clang-format on
+#pragma omp target firstprivate(a) map(from : ok)
+  {
+    ok = 1;
+    for (int i = 0; i < N; ++i)
+      if (a[i] != i + 1)
+        ok = 0;
+  }
+
+  // CHECK: values_intact=1
+  printf("values_intact=%d\n", ok);
+  return 0;
+}