[OpenMP][FIX] Properly align for firstprivate variables

[jdoerfert]

Summary

The old code didn't actually align the values, and it added padding even
when none was necessary. This approach will pad entries if necessary
and, similar to the struct case, use the host pointer as guidance.

Fixes: https://github.com/llvm/llvm-project/issues/61034

[jdoerfert]

add context

[tianshilei1992]

LG

[jdoerfert]

committed as 89a8077f3d68f0f431f3657a8805f7751f5eac69

openmp/libomptarget/src/omptarget.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "omptarget.h"
 #include "device.h"
 #include "private.h"
 #include "rtl.h"
 
+#include "llvm/ADT/bit.h"
+
 #include <cassert>
 #include <cstdint>
 #include <vector>
 
 using llvm::SmallVector;
 
 int AsyncInfoTy::synchronize() {
   int Result = OFFLOAD_SUCCESS;

  * extending the size of the allocated chuck accordingly, the chuck on the
  * device will start at 0x200 with the padding (4 bytes), then &s1.b=0x204 and
  * &s1.p=0x208, as they should be to satisfy the alignment requirements.
  */
 static const int64_t MaxAlignment = 16;
 
 /// Return the alignment requirement of partially mapped structs, see
 /// MaxAlignment above.
-static int64_t getPartialStructRequiredAlignment(void *HstPtrBase) {
+static uint64_t getPartialStructRequiredAlignment(void *HstPtrBase) {
   auto BaseAlignment = reinterpret_cast<uintptr_t>(HstPtrBase) % MaxAlignment;
   return BaseAlignment == 0 ? MaxAlignment : BaseAlignment;
 }
 
 /// Map global data and execute pending ctors
 static int initLibrary(DeviceTy &Device) {
   /*
    * Map global data

 }
 
 /// A class manages private arguments in a target region.
 class PrivateArgumentManagerTy {
   /// A data structure for the information of first-private arguments. We can
   /// use this information to optimize data transfer by packing all
   /// first-private arguments and transfer them all at once.
   struct FirstPrivateArgInfoTy {
-    /// The index of the element in \p TgtArgs corresponding to the argument
-    int Index;
     /// Host pointer begin
     char *HstPtrBegin;
     /// Host pointer end
     char *HstPtrEnd;
-    /// Aligned size
-    int64_t AlignedSize;
+    /// The index of the element in \p TgtArgs corresponding to the argument
+    int Index;
+    /// Alignment of the entry (base of the entry, not after the entry).
+    uint32_t Alignment;
+    /// Size (without alignment, see padding)
+    uint32_t Size;
+    /// Padding used to align the next argument.
+    uint32_t Padding;
     /// Host pointer name
     map_var_info_t HstPtrName = nullptr;
 
-    FirstPrivateArgInfoTy(int Index, void *HstPtr, int64_t Size,
+    FirstPrivateArgInfoTy(int Index, void *HstPtr, uint32_t Size,
+                          uint32_t Alignment, uint32_t Padding,
                           const map_var_info_t HstPtrName = nullptr)
-        : Index(Index), HstPtrBegin(reinterpret_cast<char *>(HstPtr)),
-          HstPtrEnd(HstPtrBegin + Size),
-          AlignedSize(Size + Size % MaxAlignment), HstPtrName(HstPtrName) {}
+        : HstPtrBegin(reinterpret_cast<char *>(HstPtr)),
+          HstPtrEnd(HstPtrBegin + Size), Index(Index), Alignment(Alignment),
+          Size(Size), Padding(Padding), HstPtrName(HstPtrName) {}
   };
 
   /// A vector of target pointers for all private arguments
   SmallVector<void *> TgtPtrs;
 
   /// A vector of information of all first-private arguments to be packed
   SmallVector<FirstPrivateArgInfoTy> FirstPrivateArgInfo;
   /// Host buffer for all arguments to be packed
#endif
       // We will pack all this kind of arguments to transfer them all at once
       // to reduce the number of data transfer. We will not take
       // non-first-private arguments, aka. private arguments that doesn't need
       // to be mapped to target device, into account because data allocation
       // can be very efficient with memory manager.
 
       // Placeholder value
       TgtPtr = nullptr;
+      auto *LastFPArgInfo =
+          FirstPrivateArgInfo.empty() ? nullptr : &FirstPrivateArgInfo.back();
+
+      // Compute the start alignment of this entry, add padding if necessary.
+      // TODO: Consider sorting instead.
+      uint32_t Padding = 0;
+      uint32_t StartAlignment =
+          LastFPArgInfo ? LastFPArgInfo->Alignment : MaxAlignment;
+      if (LastFPArgInfo) {
+        // Check if we keep the start alignment or if it is shrunk due to the
+        // size of the last element.
+        uint32_t Offset = LastFPArgInfo->Size % StartAlignment;
+        if (Offset)
+          StartAlignment = Offset;
+        // We only need as much alignment as the host pointer had (since we
+        // don't know the alignment information from the source we might end up
+        // overaligning accesses but not too much).
+        uint32_t RequiredAlignment =
+            llvm::bit_floor(getPartialStructRequiredAlignment(HstPtr));
+        if (RequiredAlignment > StartAlignment) {
+          Padding = RequiredAlignment - StartAlignment;
+          StartAlignment = RequiredAlignment;
+        }
+      }
+
       FirstPrivateArgInfo.emplace_back(TgtArgsIndex, HstPtr, ArgSize,
-                                       HstPtrName);
-      FirstPrivateArgSize += FirstPrivateArgInfo.back().AlignedSize;
+                                       StartAlignment, Padding, HstPtrName);
+      FirstPrivateArgSize += Padding + ArgSize;
     }
 
     return OFFLOAD_SUCCESS;
   }
 
   /// Pack first-private arguments, replace place holder pointers in \p TgtArgs,
   /// and start the transfer.
   int packAndTransfer(SmallVector<void *> &TgtArgs) {
     if (!FirstPrivateArgInfo.empty()) {
       assert(FirstPrivateArgSize != 0 &&
              "FirstPrivateArgSize is 0 but FirstPrivateArgInfo is empty");
       FirstPrivateArgBuffer.resize(FirstPrivateArgSize, 0);
       auto Itr = FirstPrivateArgBuffer.begin();
       // 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.
+        Itr = std::next(Itr, Info.Padding);
         std::copy(Info.HstPtrBegin, Info.HstPtrEnd, Itr);
-        Itr = std::next(Itr, Info.AlignedSize);
+        Itr = std::next(Itr, Info.Size);
       }
       // Allocate target memory
       void *TgtPtr =
           Device.allocData(FirstPrivateArgSize, FirstPrivateArgBuffer.data());
       if (TgtPtr == nullptr) {
         DP("Failed to allocate target memory for private arguments.\n");
         return OFFLOAD_FAIL;
       }
       TgtPtrs.push_back(TgtPtr);
       DP("Allocated %" PRId64 " bytes of target memory at " DPxMOD "\n",
          FirstPrivateArgSize, DPxPTR(TgtPtr));
       // Transfer data to target device
       int Ret = Device.submitData(TgtPtr, FirstPrivateArgBuffer.data(),
                                   FirstPrivateArgSize, AsyncInfo);
       if (Ret != OFFLOAD_SUCCESS) {
         DP("Failed to submit data of private arguments.\n");
         return OFFLOAD_FAIL;
       }
       // Fill in all placeholder pointers
       auto TP = reinterpret_cast<uintptr_t>(TgtPtr);
       for (FirstPrivateArgInfoTy &Info : FirstPrivateArgInfo) {
         void *&Ptr = TgtArgs[Info.Index];
         assert(Ptr == nullptr && "Target pointer is already set by mistaken");
+        // Pad the device pointer to get the right alignment.
+        TP += Info.Padding;
         Ptr = reinterpret_cast<void *>(TP);
-        TP += Info.AlignedSize;
+        TP += Info.Size;
         DP("Firstprivate array " DPxMOD " of size %" PRId64 " mapped to " DPxMOD
            "\n",
            DPxPTR(Info.HstPtrBegin), Info.HstPtrEnd - Info.HstPtrBegin,
            DPxPTR(Ptr));
       }
     }
 
     return OFFLOAD_SUCCESS;

openmp/libomptarget/test/mapping/firstprivate_aligned.cpp

+// RUN: %libomptarget-compilexx-generic -O3 && %libomptarget-run-generic
+
+#include <stdio.h>
+
+// CHECK: rx: 16, ry: 16;
+// CHECK: rx: 16, ry: 16;
+// CHECK: rx: 16, ry: 16;
+// CHECK: rx: 16, ry: 16;
+
+template <bool Aligned> void test() {
+  printf("Test %saligned firstprivate\n", Aligned ? "" : "non-");
+  char z1[3 + Aligned], z2[3 + Aligned];
+  int x[4];
+  int y[4];
+  y[0] = y[1] = y[2] = y[3] = 4;
+  x[0] = x[1] = x[2] = x[3] = 4;
+  int rx = -1, ry = -1;
+#pragma omp target firstprivate(z1, y, z2) map(from : ry, rx) map(to : x)
+  {
+    ry = (y[0] + y[1] + y[2] + y[3]);
+    rx = (x[0] + x[1] + x[2] + x[3]);
+  }
+  printf(" rx:%i, ry:%i\n", rx, ry);
+#pragma omp target firstprivate(z1, y, z2) map(from : ry, rx) map(to : x)
+  {
+    z1[2] += 5;
+    ry = (y[0] + y[1] + y[2] + y[3]);
+    rx = (x[0] + x[1] + x[2] + x[3]);
+    z2[2] += 7;
+  }
+  printf(" rx:%i, ry:%i\n", rx, ry);
+}
+
+int main() {
+  test<true>();
+  test<false>();
+}