[llvm] [Offload] Do not load images from the same descriptor on the same device (PR #139147)

[Joseph Huber via llvm-commits]

https://github.com/jhuber6 created https://github.com/llvm/llvm-project/pull/139147

Summary:
Right now we generally assume that we have one image per device. The
binary descriptor represents a single 'compilation'. This means that
each image is going to contain the same code built for different
architectures when used through the OpenMP interface. This is
problematic when we have cases where the same code will then be loaded
multiple times (like wiht sm_80, sm_89 or the generic GFX ISAs). This
patch is the quick and dirty slution, we just prevent this from
happening at all. This means we use the first one we find, which might
not be overly optimal, but it should be better than the alternative.
Note that this does not affect shared library loads as it is per binary
descriptor, not per device.


>From a58c1af7f521080d4d912928638712e9d11bdf71 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Thu, 8 May 2025 15:14:26 -0500
Subject: [PATCH] [Offload] Do not load images from the same descriptor on the
 same device

Summary:
Right now we generally assume that we have one image per device. The
binary descriptor represents a single 'compilation'. This means that
each image is going to contain the same code built for different
architectures when used through the OpenMP interface. This is
problematic when we have cases where the same code will then be loaded
multiple times (like wiht sm_80, sm_89 or the generic GFX ISAs). This
patch is the quick and dirty slution, we just prevent this from
happening at all. This means we use the first one we find, which might
not be overly optimal, but it should be better than the alternative.
Note that this does not affect shared library loads as it is per binary
descriptor, not per device.
---
 offload/libomptarget/PluginManager.cpp | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/offload/libomptarget/PluginManager.cpp b/offload/libomptarget/PluginManager.cpp
index d6d529a207587..712458b4de8dd 100644
--- a/offload/libomptarget/PluginManager.cpp
+++ b/offload/libomptarget/PluginManager.cpp
@@ -202,6 +202,7 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
     PM->addDeviceImage(*Desc, Desc->DeviceImages[i]);
 
   // Register the images with the RTLs that understand them, if any.
+  llvm::DenseMap<GenericPluginTy *, llvm::DenseSet<int32_t>> UsedDevices;
   for (DeviceImageTy &DI : PM->deviceImages()) {
     // Obtain the image and information that was previously extracted.
     __tgt_device_image *Img = &DI.getExecutableImage();
@@ -232,6 +233,17 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
         if (!initializeDevice(R, DeviceId))
           continue;
 
+        // We only want a single matching image to be registered for each binary
+        // descriptor. This prevents multiple of the same image from being
+        // registered for the same device in the case that they are mutually
+        // compatible, such as sm_80 and sm_89.
+        if (!UsedDevices[&R].insert(DeviceId).second) {
+          DP("Image " DPxMOD
+             " is a duplicate, not loaded on RTL %s device %d!\n",
+             DPxPTR(Img->ImageStart), R.getName(), DeviceId);
+          continue;
+        }
+
         // Initialize (if necessary) translation table for this library.
         PM->TrlTblMtx.lock();
         if (!PM->HostEntriesBeginToTransTable.count(Desc->HostEntriesBegin)) {