[llvm] [Offload] Initialize all platforms before plugin device creation (PR #160702)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Piotr Balcer (pbalcer)

<details>
<summary>Changes</summary>

Devices store a raw pointer to back to their owning Platform. Platforms are stored directly inside of a vector. Modifying this vector risks invalidating all the platform pointers stored in devices.

This change moves the host platform emplace to happen before device object creation, eliminating the risk of vector reallocation after platform pointers are set.

This popped up when trying to use https://github.com/llvm/llvm-project/pull/158900 with liboffload. The platforms vector was reallocated and liboffload segfaulted when iterating over the platforms. This is at attempt at the simplest possible fix.

Since the number of platforms is fixed, I initially wanted to refactor this code so that it became impossible to modify the platforms container after Context is created, but, when that patch grew large, I pivoted to this instead. Another option is to use shared_ptr/weak_ptr to allocate the Platform instead of storing it in the vector directly.

Given, https://github.com/llvm/llvm-project/issues/159636, which requires a non-trivial change to this same code anyway, I'm thinking I can do follow-up refactor that tackles both issues at the same time.
Thoughts?

---
Full diff: https://github.com/llvm/llvm-project/pull/160702.diff


1 Files Affected:

- (modified) offload/liboffload/src/OffloadImpl.cpp (+12-2) 


``````````diff
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 08a2e25b97d85..0701ff12c4953 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -256,9 +256,18 @@ Error initPlugins(OffloadContext &Context) {
           pluginNameToBackend(#Name)});                                        \
   } while (false);
 #include "Shared/Targets.def"
+  // Host platform must be added before initializing plugin devices because
+  // devices contain a pointer back to the owning platform, and modifying the
+  // Platforms vector risks reallocating the underlying storage, thus invalidating
+  // all the platform pointers.
+  Context.Platforms.emplace_back(
+    ol_platform_impl_t{nullptr, OL_PLATFORM_BACKEND_HOST});
 
   // Preemptively initialize all devices in the plugin
   for (auto &Platform : Context.Platforms) {
+    if (Platform.BackendType == OL_PLATFORM_BACKEND_HOST)
+      continue;
+
     auto Err = Platform.Plugin->init();
     [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
     for (auto DevNum = 0; DevNum < Platform.Plugin->number_of_devices();
@@ -273,10 +282,11 @@ Error initPlugins(OffloadContext &Context) {
       }
     }
   }
+  // The Context.Platforms cannot be modified after this point without updating
+  // all the Device.Platform pointers.
 
   // Add the special host device
-  auto &HostPlatform = Context.Platforms.emplace_back(
-      ol_platform_impl_t{nullptr, OL_PLATFORM_BACKEND_HOST});
+  auto &HostPlatform = Context.Platforms.back();
   HostPlatform.Devices.emplace_back(
       std::make_unique<ol_device_impl_t>(-1, nullptr, nullptr, InfoTreeNode{}));
   Context.HostDevice()->Platform = &HostPlatform;

``````````

</details>


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