[OpenMP][OMPT] Fix device identifier collision during callbacks (#65595)

Fixes: #65104
When a user assigns devices to target regions it may happen that
different identifiers will map onto the same id within different
plugins. This will lead to situations where callbacks will become much
harder to read, as ambiguous identifiers are reported.

We fix this by collecting the index-offset upon general RTL
initialization. Which in turn, allows to calculate the unique,
user-observable device id.

Author: mhalk

openmp/libomptarget/include/omptargetplugin.h

@@ -218,6 +218,10 @@ int32_t __tgt_rtl_data_notify_mapped(int32_t ID, void *HstPtr, int64_t Size);
 // host address \p HstPtr and \p Size bytes.
 int32_t __tgt_rtl_data_notify_unmapped(int32_t ID, void *HstPtr);
 
+// Set the global device identifier offset, such that the plugin may determine a
+// unique device number.
+int32_t __tgt_rtl_set_device_offset(int32_t DeviceIdOffset);
+
 #ifdef __cplusplus
 }
 #endif

openmp/libomptarget/include/rtl.h

@@ -72,6 +72,7 @@ struct RTLInfoTy {
   typedef int32_t(data_unlock_ty)(int32_t, void *);
   typedef int32_t(data_notify_mapped_ty)(int32_t, void *, int64_t);
   typedef int32_t(data_notify_unmapped_ty)(int32_t, void *);
+  typedef int32_t(set_device_offset_ty)(int32_t);
   typedef int32_t(activate_record_replay_ty)(int32_t, uint64_t, bool, bool);
 
   int32_t Idx = -1;             // RTL index, index is the number of devices
@@ -125,6 +126,7 @@ struct RTLInfoTy {
   data_unlock_ty *data_unlock = nullptr;
   data_notify_mapped_ty *data_notify_mapped = nullptr;
   data_notify_unmapped_ty *data_notify_unmapped = nullptr;
+  set_device_offset_ty *set_device_offset = nullptr;
   activate_record_replay_ty *activate_record_replay = nullptr;
 
   // Are there images associated with this RTL.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

@@ -542,7 +542,8 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
     bool ExpectedStatus = false;
     if (OmptInitialized.compare_exchange_strong(ExpectedStatus, true))
       performOmptCallback(device_initialize,
-                          /* device_num */ DeviceId,
+                          /* device_num */ DeviceId +
+                              Plugin.getDeviceIdStartIndex(),
                           /* type */ getComputeUnitKind().c_str(),
                           /* device */ reinterpret_cast<ompt_device_t *>(this),
                           /* lookup */ ompt::lookupCallbackByName,
@@ -587,7 +588,7 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
   return Plugin::success();
 }
 
-Error GenericDeviceTy::deinit() {
+Error GenericDeviceTy::deinit(GenericPluginTy &Plugin) {
   // Delete the memory manager before deinitializing the device. Otherwise,
   // we may delete device allocations after the device is deinitialized.
   if (MemoryManager)
@@ -605,7 +606,9 @@ Error GenericDeviceTy::deinit() {
   if (ompt::Initialized) {
     bool ExpectedStatus = true;
     if (OmptInitialized.compare_exchange_strong(ExpectedStatus, false))
-      performOmptCallback(device_finalize, /* device_num */ DeviceId);
+      performOmptCallback(device_finalize,
+                          /* device_num */ DeviceId +
+                              Plugin.getDeviceIdStartIndex());
   }
 #endif
 
@@ -656,7 +659,8 @@ GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
     size_t Bytes =
         getPtrDiff(InputTgtImage->ImageEnd, InputTgtImage->ImageStart);
     performOmptCallback(device_load,
-                        /* device_num */ DeviceId,
+                        /* device_num */ DeviceId +
+                            Plugin.getDeviceIdStartIndex(),
                         /* FileName */ nullptr,
                         /* File Offset */ 0,
                         /* VmaInFile */ nullptr,
@@ -1362,7 +1366,7 @@ Error GenericPluginTy::deinitDevice(int32_t DeviceId) {
     return Plugin::success();
 
   // Deinitialize the device and release its resources.
-  if (auto Err = Devices[DeviceId]->deinit())
+  if (auto Err = Devices[DeviceId]->deinit(*this))
     return Err;
 
   // Delete the device and invalidate its reference.
@@ -1815,6 +1819,12 @@ int32_t __tgt_rtl_init_device_info(int32_t DeviceId,
   return OFFLOAD_SUCCESS;
 }
 
+int32_t __tgt_rtl_set_device_offset(int32_t DeviceIdOffset) {
+  Plugin::get().setDeviceIdStartIndex(DeviceIdOffset);
+
+  return OFFLOAD_SUCCESS;
+}
+
 #ifdef __cplusplus
 }
 #endif

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

@@ -611,7 +611,7 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   /// Deinitialize the device and free all its resources. After this call, the
   /// device is no longer considered ready, so no queries or modifications are
   /// allowed.
-  Error deinit();
+  Error deinit(GenericPluginTy &Plugin);
   virtual Error deinitImpl() = 0;
 
   /// Load the binary image into the device and return the target table.
@@ -946,6 +946,12 @@ struct GenericPluginTy {
   /// Get the number of active devices.
   int32_t getNumDevices() const { return NumDevices; }
 
+  /// Get the plugin-specific device identifier offset.
+  int32_t getDeviceIdStartIndex() const { return DeviceIdStartIndex; }
+
+  /// Set the plugin-specific device identifier offset.
+  void setDeviceIdStartIndex(int32_t Offset) { DeviceIdStartIndex = Offset; }
+
   /// Get the ELF code to recognize the binary image of this plugin.
   virtual uint16_t getMagicElfBits() const = 0;
 
@@ -1010,6 +1016,11 @@ struct GenericPluginTy {
   /// Number of devices available for the plugin.
   int32_t NumDevices = 0;
 
+  /// Index offset, which when added to a DeviceId, will yield a unique
+  /// user-observable device identifier. This is especially important when
+  /// DeviceIds of multiple plugins / RTLs need to be distinguishable.
+  int32_t DeviceIdStartIndex = 0;
+
   /// Array of pointers to the devices. Initially, they are all set to nullptr.
   /// Once a device is initialized, the pointer is stored in the position given
   /// by its device id. A position with nullptr means that the corresponding

openmp/libomptarget/src/device.cpp

@@ -583,7 +583,7 @@ void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {
   void *TargetPtr = nullptr;
   OMPT_IF_BUILT(InterfaceRAII TargetDataAllocRAII(
                     RegionInterface.getCallbacks<ompt_target_data_alloc>(),
-                    RTLDeviceID, HstPtr, &TargetPtr, Size,
+                    DeviceID, HstPtr, &TargetPtr, Size,
                     /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
   TargetPtr = RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);
@@ -594,7 +594,7 @@ int32_t DeviceTy::deleteData(void *TgtAllocBegin, int32_t Kind) {
   /// RAII to establish tool anchors before and after data deletion
   OMPT_IF_BUILT(InterfaceRAII TargetDataDeleteRAII(
                     RegionInterface.getCallbacks<ompt_target_data_delete>(),
-                    RTLDeviceID, TgtAllocBegin,
+                    DeviceID, TgtAllocBegin,
                     /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
   return RTL->data_delete(RTLDeviceID, TgtAllocBegin, Kind);
@@ -632,7 +632,7 @@ int32_t DeviceTy::submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
   OMPT_IF_BUILT(
       InterfaceRAII TargetDataSubmitRAII(
           RegionInterface.getCallbacks<ompt_target_data_transfer_to_device>(),
-          RTLDeviceID, TgtPtrBegin, HstPtrBegin, Size,
+          DeviceID, TgtPtrBegin, HstPtrBegin, Size,
           /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
   if (!AsyncInfo || !RTL->data_submit_async || !RTL->synchronize)
@@ -660,7 +660,7 @@ int32_t DeviceTy::retrieveData(void *HstPtrBegin, void *TgtPtrBegin,
   OMPT_IF_BUILT(
       InterfaceRAII TargetDataRetrieveRAII(
           RegionInterface.getCallbacks<ompt_target_data_transfer_from_device>(),
-          RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size,
+          DeviceID, HstPtrBegin, TgtPtrBegin, Size,
           /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
   if (!RTL->data_retrieve_async || !RTL->synchronize)

openmp/libomptarget/src/rtl.cpp

@@ -249,6 +249,8 @@ bool RTLsTy::attemptLoadRTL(const std::string &RTLName, RTLInfoTy &RTL) {
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_notify_mapped");
   *((void **)&RTL.data_notify_unmapped) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_notify_unmapped");
+  *((void **)&RTL.set_device_offset) =
+      DynLibrary->getAddressOfSymbol("__tgt_rtl_set_device_offset");
 
   // Record Replay RTL
   *((void **)&RTL.activate_record_replay) =
@@ -424,6 +426,10 @@ void RTLsTy::initRTLonce(RTLInfoTy &R) {
     R.IsUsed = true;
     UsedRTLs.push_back(&R);
 
+    // If possible, set the device identifier offset
+    if (R.set_device_offset)
+      R.set_device_offset(Start);
+
     DP("RTL " DPxMOD " has index %d!\n", DPxPTR(R.LibraryHandler.get()), R.Idx);
   }
 }