[llvm] [Offload] Ensure to load images when the device is used (PR #103002)

Joseph Huber via llvm-commits llvm-commits at lists.llvm.org
Tue Aug 13 13:49:12 PDT 2024


================
@@ -286,13 +291,194 @@ void PluginManager::unregisterLib(__tgt_bin_desc *Desc) {
   DP("Done unregistering library!\n");
 }
 
+/// Map global data and execute pending ctors
+static int loadImagesOntoDevice(DeviceTy &Device) {
+  /*
+   * Map global data
+   */
+  int32_t DeviceId = Device.DeviceID;
+  int Rc = OFFLOAD_SUCCESS;
+  {
+    std::lock_guard<decltype(PM->TrlTblMtx)> LG(PM->TrlTblMtx);
+    for (auto *HostEntriesBegin : PM->HostEntriesBeginRegistrationOrder) {
+      TranslationTable *TransTable =
+          &PM->HostEntriesBeginToTransTable[HostEntriesBegin];
+      DP("Trans table %p : %p\n", TransTable->HostTable.EntriesBegin,
+         TransTable->HostTable.EntriesEnd);
+      if (TransTable->HostTable.EntriesBegin ==
+          TransTable->HostTable.EntriesEnd) {
+        // No host entry so no need to proceed
+        continue;
+      }
+
+      if (TransTable->TargetsTable[DeviceId] != 0) {
+        // Library entries have already been processed
+        continue;
+      }
+
+      // 1) get image.
+      assert(TransTable->TargetsImages.size() > (size_t)DeviceId &&
+             "Not expecting a device ID outside the table's bounds!");
+      __tgt_device_image *Img = TransTable->TargetsImages[DeviceId];
+      if (!Img) {
+        REPORT("No image loaded for device id %d.\n", DeviceId);
+        Rc = OFFLOAD_FAIL;
+        break;
+      }
+
+      // 2) Load the image onto the given device.
+      auto BinaryOrErr = Device.loadBinary(Img);
+      if (llvm::Error Err = BinaryOrErr.takeError()) {
+        REPORT("Failed to load image %s\n",
+               llvm::toString(std::move(Err)).c_str());
+        Rc = OFFLOAD_FAIL;
+        break;
+      }
+
+      // 3) Create the translation table.
+      llvm::SmallVector<__tgt_offload_entry> &DeviceEntries =
+          TransTable->TargetsEntries[DeviceId];
+      for (__tgt_offload_entry &Entry :
+           llvm::make_range(Img->EntriesBegin, Img->EntriesEnd)) {
+        __tgt_device_binary &Binary = *BinaryOrErr;
+
+        __tgt_offload_entry DeviceEntry = Entry;
+        if (Entry.size) {
+          if (Device.RTL->get_global(Binary, Entry.size, Entry.name,
+                                     &DeviceEntry.addr) != OFFLOAD_SUCCESS)
+            REPORT("Failed to load symbol %s\n", Entry.name);
+
+          // If unified memory is active, the corresponding global is a device
+          // reference to the host global. We need to initialize the pointer on
+          // the deive to point to the memory on the host.
----------------
jhuber6 wrote:

```suggestion
          // the device to point to the memory on the host.
```

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


More information about the llvm-commits mailing list