[OpenMP] Enable map checks under unified_shared_memory mode

[carlo.bertolli]

Summary

Enable correctness checks in runtime for map clauses under unified_shared_memory mode.
Add OPENMP_DISABLE_MAPS environment variable input to disable this behavior.

As an example, before this patch is applied, this example would not fail:

#pragma omp requires unified_shared_memory

void foo() {

int a[10];

#pragma omp target enter data map(to:a[0:5])
#pragma omp target enter data map(to:a[0:10]) // extending already mapped memory, illegal!

}

This is illegal in OpenMP, and presumably also under unified_shared_memory mode, but the current implementation bails out on map checks under this mode.
By specifications, under unified_shared_memory mode, the user is allowed not to use maps and can expect the program to work on an system that supports the mode. However, for portability to systems that do not support unified_shared_memory, map clauses should always be correct. Not applying checks in the runtime should be a conscious request by the user.

[carlo.bertolli]

Update environment variable name to reflect comments and intended name

[jdoerfert]

As discussed, we should put this under an environment variable switch and do some measurements.

openmp/libomptarget/src/rtl.h

129

This is potentially uninitialized.

[grokos]

openmp/libomptarget/src/device.cpp

229–231

We will have a problem here after D107925, D107927 and D107928 land. These patches assume that USM mappings do not have an entry in HostDataToTargetMap. That's how they figure out whether we have an allocation in shared memory or in GPU memory, i.e. the value of IsHostPtr is determined by the presence or absence of an entry in the map. If we want to move on with this change here, we'll need to introduce an extra member in struct HostDataToTargetTy, e.g. bool IsUSMAlloc = false;, and this emplacement here will have to be extended with ..., HstPtrName, /*IsUSMAlloc*/ true). In general, I think that implementing desired support involves much more work than this patch. Basically, we need to treat USM data the exact same way as device-side mappings, i.e. remove all specialized code paths for USM allocations and only check whether we have USM data when we are about to alloc memory / perform data transfers.

365–366

Rebase this patch after D107925 lands, as these lines will be removed.

389–394

This check can go away after D107925 lands because the caller of this function has already performed this check.

openmp/libomptarget/src/rtl.cpp

82

It's better to rename this env var to something like LIBOMPTARGET_DISABLE_USM_MAPS. All env vars begin either with OMP_ if they are part of the OpenMP standard or with LIBOMPTARGET_ if they are libomptarget-specific extensions. We have no env vars starting with OPENMP_, so let's keep it consistent with the existing codebase. Also, including "USM" in the name makes it more descriptive of exactly what it does.

openmp/libomptarget/src/rtl.h

129

Need to initialize, I guess we want the default value to be false?

Also, can you add a comment about what this variable does? And maybe a more meaningful name would be something like "NoUSMMapEntries"?

[carlo.bertolli]

Updated patch based on comments: add new map table entry field to track USM maps; use it to determine behavior in getTargetPointer and deallocTgtPtr; update var name and add explaining comments.

openmp/libomptarget/src/device.h

 struct HostDataToTargetTy {
   const uintptr_t HstPtrBase; // host info.
   const uintptr_t HstPtrBegin;
   const uintptr_t HstPtrEnd;       // non-inclusive.
   const map_var_info_t HstPtrName; // Optional source name of mapped variable.
 
   const uintptr_t TgtPtrBegin; // target info.
 
+  const bool IsUSMAlloc; // used to track maps under USM mode (optional)
+
 private:
   static const uint64_t INFRefCount = ~(uint64_t)0;
   static std::string refCountToStr(uint64_t RefCount) {
     return RefCount == INFRefCount ? "INF" : std::to_string(RefCount);
   }
 
   struct StatesTy {
     StatesTy(uint64_t DRC, uint64_t HRC)
private:
   };
   // When HostDataToTargetTy is used by std::set, std::set::iterator is const
   // use unique_ptr to make States mutable.
   const std::unique_ptr<StatesTy> States;
 
 public:
   HostDataToTargetTy(uintptr_t BP, uintptr_t B, uintptr_t E, uintptr_t TB,
                      bool UseHoldRefCount, map_var_info_t Name = nullptr,
-                     bool IsINF = false)
+                     bool IsINF = false, bool IsUSMAlloc = false)
       : HstPtrBase(BP), HstPtrBegin(B), HstPtrEnd(E), HstPtrName(Name),
-        TgtPtrBegin(TB), States(std::make_unique<StatesTy>(UseHoldRefCount ? 0
-                                                           : IsINF ? INFRefCount
-                                                                   : 1,
-                                                           !UseHoldRefCount ? 0
-                                                           : IsINF ? INFRefCount
-                                                                   : 1)) {}
+        TgtPtrBegin(TB), IsUSMAlloc(IsUSMAlloc),
+        States(std::make_unique<StatesTy>(UseHoldRefCount ? 0
+                                          : IsINF         ? INFRefCount
+                                          : IsINF          ? INFRefCount
 
   /// Get the total reference count.  This is smarter than just getDynRefCount()
   /// + getHoldRefCount() because it handles the case where at least one is
   /// infinity and the other is non-zero.
   uint64_t getTotalRefCount() const {
     if (States->DynRefCount == INFRefCount ||
         States->HoldRefCount == INFRefCount)
       return INFRefCount;

openmp/libomptarget/src/device.cpp

if (LR.Flags.IsContained ||
     // In addition to the mapping rules above, the close map modifier forces the
     // mapping of the variable to the device.
     if (Size) {
       DP("Return HstPtrBegin " DPxMOD " Size=%" PRId64 " for unified shared "
          "memory\n",
          DPxPTR((uintptr_t)HstPtrBegin), Size);
       IsHostPtr = true;
       TargetPointer = HstPtrBegin;
+
+      // If user requested map checks under USM mode, add map information
+      // to the table and mark as usm allocation.
+      // Infinite ref count to disable memcopies to/from device memory
+      if (PM->RTLs.USMMapChecks)
+        Entry = HostDataToTargetMap
+                    .emplace((uintptr_t)HstPtrBase, (uintptr_t)HstPtrBegin,
+                             (uintptr_t)HstPtrBegin + Size,
+                             (uintptr_t)HstPtrBegin, HasHoldModifier,
+                             HstPtrName, /*IsInf=*/true, /*IsUSMAlloc=*/true)
+                    .first;
     }
   } else if (HasPresentModifier) {
     DP("Mapping required by 'present' map type modifier does not exist for "
        "HstPtrBegin=" DPxMOD ", Size=%" PRId64 "\n",
        DPxPTR(HstPtrBegin), Size);
     MESSAGE("device mapping required by 'present' map type modifier does not "
             "exist for host address " DPxMOD " (%" PRId64 " bytes)",
             DPxPTR(HstPtrBegin), Size);
void *DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
                                bool &IsHostPtr, bool MustContain,
                                bool ForceDelete) {
   void *rc = NULL;
   IsHostPtr = false;
   IsLast = false;
   DataMapMtx.lock();
   LookupResult lr = lookupMapping(HstPtrBegin, Size);
 
-  if (lr.Flags.IsContained ||
-      (!MustContain && (lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter))) {
+  // When map checks are enabled under USM mode, mapped host pointers are
+  // tracked in the map table but should be treated as in the USM case
+  if ((lr.Flags.IsContained ||
+       (!MustContain && (lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter))) &&
+      !lr.Entry->IsUSMAlloc) {
     auto &HT = *lr.Entry;
     // We do not zero the total reference count here.  deallocTgtPtr does that
     // atomically with removing the mapping.  Otherwise, before this thread
     // removed the mapping in deallocTgtPtr, another thread could retrieve the
     // mapping, increment and decrement back to zero, and then both threads
     // would try to remove the mapping, resulting in a double free.
     IsLast = HT.decShouldRemove(UseHoldRefCount, ForceDelete);
     const char *RefCountAction;
if ((lr.Flags.IsContained ||
     uintptr_t tp = HT.TgtPtrBegin + ((uintptr_t)HstPtrBegin - HT.HstPtrBegin);
     INFO(OMP_INFOTYPE_MAPPING_EXISTS, DeviceID,
          "Mapping exists with HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD ", "
          "Size=%" PRId64 ", DynRefCount=%s%s, HoldRefCount=%s%s\n",
          DPxPTR(HstPtrBegin), DPxPTR(tp), Size, HT.dynRefCountToStr().c_str(),
          DynRefCountAction, HT.holdRefCountToStr().c_str(), HoldRefCountAction);
     rc = (void *)tp;
   } else if (PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) {
-    // If the value isn't found in the mapping and unified shared memory
-    // is on then it means we have stumbled upon a value which we need to
-    // use directly from the host.
+    // If the value isn't found in the mapping or if the value is found but it
+    // is related to a USM mapping and unified shared memory is on then it means
+    // we have stumbled upon a value which we need to use directly from the
+    // host.
     DP("Get HstPtrBegin " DPxMOD " Size=%" PRId64 " for unified shared "
        "memory\n",
        DPxPTR((uintptr_t)HstPtrBegin), Size);
     IsHostPtr = true;
     rc = HstPtrBegin;
   }
 
   DataMapMtx.unlock();
if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {
     auto &HT = *lr.Entry;
     uintptr_t tp = HT.TgtPtrBegin + (hp - HT.HstPtrBegin);
     return (void *)tp;
   }
 
   return NULL;
 }
 
 int DeviceTy::deallocTgtPtr(void *HstPtrBegin, int64_t Size,
                             bool HasHoldModifier) {
   // Check if the pointer is contained in any sub-nodes.
   int rc;
   DataMapMtx.lock();
   LookupResult lr = lookupMapping(HstPtrBegin, Size);
   if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {
     auto &HT = *lr.Entry;
     if (HT.decRefCount(HasHoldModifier) == 0) {
+      if (!HT.IsUSMAlloc) {
         DP("Deleting tgt data " DPxMOD " of size %" PRId64 "\n",
            DPxPTR(HT.TgtPtrBegin), Size);
         deleteData((void *)HT.TgtPtrBegin);
+      }
       INFO(OMP_INFOTYPE_MAPPING_CHANGED, DeviceID,
            "Removing map entry with HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD
            ", Size=%" PRId64 ", Name=%s\n",
            DPxPTR(HT.HstPtrBegin), DPxPTR(HT.TgtPtrBegin), Size,
            (HT.HstPtrName) ? getNameFromMapping(HT.HstPtrName).c_str()
                            : "unknown");
+      if (PM->RTLs.USMMapChecks)
         HostDataToTargetMap.erase(lr.Entry);
     }
     rc = OFFLOAD_SUCCESS;
   } else {
     REPORT("Section to delete (hst addr " DPxMOD ") does not exist in the"
            " allocated memory\n",
            DPxPTR(HstPtrBegin));
     rc = OFFLOAD_FAIL;
   }

openmp/libomptarget/src/rtl.h

struct RTLsTy {
   std::list<RTLInfoTy> AllRTLs;
 
   // Array of pointers to the detected runtime libraries that have compatible
   // binaries.
   std::vector<RTLInfoTy *> UsedRTLs;
 
   int64_t RequiresFlags = OMP_REQ_UNDEFINED;
 
+  // When true, perform map checks under USM mode
+  bool USMMapChecks = false;
+
   explicit RTLsTy() = default;
 
   // Register the clauses of the requires directive.
   void RegisterRequires(int64_t flags);
 
   // Initialize RTL if it has not been initialized
   void initRTLonce(RTLInfoTy &RTL);
 

openmp/libomptarget/src/rtl.cpp

 void RTLsTy::LoadRTLs() {
   // Parse environment variable OMP_TARGET_OFFLOAD (if set)
   PM->TargetOffloadPolicy =
       (kmp_target_offload_kind_t)__kmpc_get_target_offload();
   if (PM->TargetOffloadPolicy == tgt_disabled) {
     return;
   }
 
+  // Parse environment variable LIBOMPTARGET_USM_MAP_CHECKS to enable map checks
+  // under USM mode Expected values are:
+  // * "check" : enable map checks under usm mode
+  // * "none"  : disable map checks under usm mode
+  // * not set : disable map checks under usm mode
+  // (no device allocations or host-to-device and device-to-host transfers)
+  if (const char *USMMapChecksStr = getenv("LIBOMPTARGET_USM_MAP_CHECKS")) {
+    const char *checkStr = "check";
+    if (strncmp(USMMapChecksStr, checkStr, strlen(checkStr)) == 0)
+      USMMapChecks = true;
+  }
+
   DP("Loading RTLs...\n");
 
   // Attempt to open all the plugins and, if they exist, check if the interface
   // is correct and if they are supporting any devices.
   for (auto *Name : RTLNames) {
     DP("Loading library '%s'...\n", Name);
     void *dynlib_handle = dlopen(Name, RTLD_NOW);
 

openmp/libomptarget/test/unified_shared_memory/check_maps.c

+// RUN: %libomptarget-compile-generic && env LIBOMPTARGET_USM_MAP_CHECKS=check
+// %libomptarget-run-generic 2>&1 | %fcheck-generic -allow-empty
+// -check-prefix=CHECK-USM-CHECK RUN: %libomptarget-compile-generic && env
+// LIBOMPTARGET_USM_MAP_CHECKS=none %libomptarget-run-generic 2>&1 |
+// %fcheck-generic -allow-empty -check-prefix=CHECK-USM-NONE RUN:
+// %libomptarget-compile-generic && %libomptarget-run-generic 2>&1 |
+// %fcheck-generic -allow-empty -check-prefix=CHECK-USM-UNSPEC
+
+// CHECK-USM-CHECK: Libomptarget message: explicit extension not allowed: host
+// address specified is 0x{{.*}}, but device allocation maps to host at 0x{{.*}}
+// CHECK-USM-CHECK: Libomptarget error: Call to getTargetPointer returned null
+// pointer (device failure or illegal mapping). CHECK-USM-CHECK: Libomptarget
+// fatal error 1: failure of target construct while offloading is mandatory
+
+// CHECK-USM-NONE: No errors
+
+// CHECK-USM-UNSPEC: No errors
+
+// REQUIRES: unified_shared_memory
+
+#include <omp.h>
+#include <stdio.h>
+
+// ---------------------------------------------------------------------------
+// Various definitions copied from OpenMP RTL
+
+extern void __tgt_register_requires(int64_t);
+
+// End of definitions copied from OpenMP RTL.
+// ---------------------------------------------------------------------------
+
+#pragma omp requires unified_shared_memory
+
+#define N 1024
+
+int main(int argc, char *argv[]) {
+  int data[N];
+
+  // Manual registration of requires flags for Clang versions
+  // that do not support requires.
+  __tgt_register_requires(8);
+
+#pragma omp target enter data map(to : data [0:N / 2])
+
+// extending a map under unified_shared_memory mode is prohibited
+// when map checks are on
+#pragma omp target enter data map(to : data [0:N])
+
+  printf("No errors\n");
+
+#pragma omp target
+  {}
+
+  return 0;
+}