[llvm] [Offload] Introduce ATTACH map-type support for pointer attachment. (PR #149036)

[Joseph Huber via llvm-commits]

================
@@ -324,17 +325,172 @@ int targetDataMapper(ident_t *Loc, DeviceTy &Device, void *ArgBase, void *Arg,
                               MapperArgsBase.data(), MapperArgs.data(),
                               MapperArgSizes.data(), MapperArgTypes.data(),
                               MapperArgNames.data(), /*arg_mappers*/ nullptr,
-                              AsyncInfo, /*FromMapper=*/true);
+                              AsyncInfo, AttachInfo, /*FromMapper=*/true);
 
   return Rc;
 }
 
+/// Utility function to perform a pointer attachment operation.
+///
+/// For something like:
+/// ```cpp
+///  int *p;
+///  ...
+///  #pragma omp target enter data map(to:p[10:10])
+/// ```
+///
+/// for which the attachment operation gets represented using:
+/// ```
+///   &p, &p[10], sizeof(p), ATTACH
+/// ```
+///
+/// (Hst|Tgt)PtrAddr   represents &p
+/// (Hst|Tgt)PteeBase  represents &p[0]
+/// (Hst|Tgt)PteeBegin represents &p[10]
+///
+/// This function first computes the expected TgtPteeBase using:
+///   `<Select>TgtPteeBase = TgtPteeBegin - (HstPteeBegin - HstPteeBase)`
+///
+/// and then attaches TgtPteeBase to TgtPtrAddr.
+///
+/// \p HstPtrSize represents the size of the pointer p. For C/C++, this
+/// should be same as "sizeof(void*)" (say 8).
+///
+/// However, for Fortran, pointers/allocatables, which are also eligible for
+/// "pointer-attachment", may be implemented using descriptors that contain the
+/// address of the pointee in the first 8 bytes, but also contain other
+/// information such as lower-bound/upper-bound etc in their subsequent fields.
+///
+/// For example, for the following:
+/// ```fortran
+///   integer, allocatable :: x(:)
+///   integer, pointer :: p(:)
+///   ...
+///   p => x(10: 19)
+///   ...
+///   !$omp target enter data map(to:p(:))
+/// ```
+///
+/// The map should trigger a pointer-attachment (assuming the pointer-attachment
+/// conditions as noted on processAttachEntries are met) between the descriptor
+/// for p, and its pointee data.
+///
+/// Since only the first 8 bytes of the descriptor contain the address of the
+/// pointee, an attachment operation on device descriptors involves:
+/// * Setting the first 8 bytes of the device descriptor to point the device
+/// address of the pointee.
+/// * Copying the remaining information about bounds/offset etc. from the host
+/// descriptor to the device descriptor.
+///
+/// The function also handles pointer-attachment portion of PTR_AND_OBJ maps,
+/// like:
+/// ```
+///   &p, &p[10], 10 * sizeof(p[10]), PTR_AND_OBJ
+/// ```
+/// by using `sizeof(void*)` as \p HstPtrSize.
+static int performPointerAttachment(DeviceTy &Device, AsyncInfoTy &AsyncInfo,
+                                    void **HstPtrAddr, void *HstPteeBase,
+                                    void *HstPteeBegin, void **TgtPtrAddr,
+                                    void *TgtPteeBegin, int64_t HstPtrSize,
+                                    TargetPointerResultTy &PtrTPR) {
+  assert(PtrTPR.getEntry() &&
+         "Need a valid pointer entry to perform pointer-attachment");
+
+  int64_t VoidPtrSize = sizeof(void *);
+  assert(HstPtrSize >= VoidPtrSize && "PointerSize is too small");
+
+  uint64_t Delta = (uint64_t)HstPteeBegin - (uint64_t)HstPteeBase;
+  void *TgtPteeBase = (void *)((uint64_t)TgtPteeBegin - Delta);
+
+  // Add shadow pointer tracking
+  // TODO: Support shadow-tracking of larger than VoidPtrSize pointers,
+  // to support restoration of Fortran descriptors. Currently, this check
+  // would return false, even if the host Fortran descriptor had been
+  // updated since its previous map, and we should have updated its
+  // device counterpart. e.g.
+  //
+  //   !$omp target enter data map(x(1:100)) !             (1)
+  //   p => x(10: 19)
+  //   !$omp target enter data map(p, p(:)) !              (2)
+  //   p => x(5: 9)
+  //   !$omp target enter data map(attach(always): p(:)) ! (3)
+  //
+  // While PtrAddr(&desc_p) and PteeBase(&p(1)) are same for (2) and (3), the
+  // pointer attachment for (3) needs to update the bounds information
+  // in the descriptor of p on device.
+  if (!PtrTPR.getEntry()->addShadowPointer(
+          ShadowPtrInfoTy{HstPtrAddr, HstPteeBase, TgtPtrAddr, TgtPteeBase}))
+    return OFFLOAD_SUCCESS;
+
+  DP("Update pointer (" DPxMOD ") -> [" DPxMOD "]\n", DPxPTR(TgtPtrAddr),
+     DPxPTR(TgtPteeBase));
+
+  // Lambda to handle submitData result and perform final steps.
+  auto HandleSubmitResult = [&](int SubmitResult) -> int {
+    if (SubmitResult != OFFLOAD_SUCCESS) {
+      REPORT("Failed to update pointer on device.\n");
+      return OFFLOAD_FAIL;
+    }
+
+    if (PtrTPR.getEntry()->addEventIfNecessary(Device, AsyncInfo) !=
+        OFFLOAD_SUCCESS)
+      return OFFLOAD_FAIL;
+
+    return OFFLOAD_SUCCESS;
+  };
+
+  bool IsPtrAFortranDescriptor = HstPtrSize > VoidPtrSize;
+  if (!IsPtrAFortranDescriptor) {
+    // For "regular" pointers, we can use the VoidPtrLocation from AsyncInfo as
+    // the buffer space for the submission.
+    void *&BufferElement = AsyncInfo.getVoidPtrLocation();
+    BufferElement = TgtPteeBase;
+
+    // Submit the updated pointer value to device
+    return HandleSubmitResult(Device.submitData(
+        TgtPtrAddr, &BufferElement, VoidPtrSize, AsyncInfo, PtrTPR.getEntry()));
+  }
+
+  // For larger "pointers" (like Fortran's descriptors), we create a dynamic
+  // buffer, which will be eventually destroyed by AsyncInfo's post-processing
+  // callback.
+  char *DataBuffer = new char[HstPtrSize];
----------------
jhuber6 wrote:

Can probably also use unique pointer for this.

https://github.com/llvm/llvm-project/pull/149036