[llvm] [Offload] Lazily initialize platforms in the Offloading API (PR #227084)
Joseph Huber via llvm-commits
llvm-commits at lists.llvm.org
Mon Sep 28 11:57:16 PDT 2026
https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/227084
>From 083cb0aa9b1880fe3b5b56fde53e06c3dc397338 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Mon, 28 Sep 2026 13:51:33 -0500
Subject: [PATCH] [Offload] Lazily initialize platforms in the Offloading API
Summary:
The Offloading library wraps around the underlying plugins. The problem
is that we currently initialize all plugins we find, even if they are
not needed for the program. This is very expensive for trivial uses, as
fully heterogenous usage is quite rare. In practice this means that you
will always pay a 200 ms penalty for having CUDA installed.
This patch changes the behavior to provide accessors into the plugins
and devices that allows them to be initialized lazily. We use a
once_flag, this should properly take a fast-path check while still
blocking on concurrent use.
Making full use of this will require a way to filter platforms more
specifically. I'm thinking of what this would look like as an API.
I'm thinking that we either have an extra iterate function that takes a
callback on the platform, or we just provide a helper to find all the
devices that can run a given image. Maybe both?
Fixes: #159636
Realnd of: https://github.com/llvm/llvm-project/pull/163272
---
offload/liboffload/API/Common.td | 6 +-
offload/liboffload/src/OffloadImpl.cpp | 110 ++++++++++++++----
.../tools/offload-tblgen/EntryPointGen.cpp | 6 +-
offload/unittests/OffloadAPI/init/olInit.cpp | 20 ++++
4 files changed, 113 insertions(+), 29 deletions(-)
diff --git a/offload/liboffload/API/Common.td b/offload/liboffload/API/Common.td
index f09ee57961d67..d5030a784e964 100644
--- a/offload/liboffload/API/Common.td
+++ b/offload/liboffload/API/Common.td
@@ -167,8 +167,9 @@ def olInit : Function {
let desc = "Perform initialization of the Offload library";
let details = [
"This must be the first API call made by a user of the Offload library",
+ "The underlying platforms are lazily initialized on their first use",
"Each call will increment an internal reference count that is decremented by `olShutDown`",
- "If InitArgs is NULL, default configuration is used which initializes all available platforms"
+ "If InitArgs is NULL, default configuration is used which makes all available platforms usable"
];
let params = [
Param<"const ol_init_args_t*", "InitArgs", "Optional pointer to initialization configuration. NULL uses defaults.", PARAM_IN_OPTIONAL>
@@ -180,7 +181,8 @@ def olShutDown : Function {
let desc = "Release the resources in use by Offload";
let details = [
"This decrements an internal reference count. When this reaches 0, all resources will be released",
- "Subsequent API calls to methods other than `olInit` made after resources are released will return OL_ERRC_UNINITIALIZED"
+ "Subsequent API calls to methods other than `olInit` made after resources are released will return OL_ERRC_UNINITIALIZED",
+ "Resources still held at process exit are released before the vendor runtimes are torn down, as if the reference count reached 0. Calls balancing an earlier `olInit` still succeed afterwards"
];
let params = [];
let returns = [];
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 11c07483b38aa..ffdd2b1a33e51 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -19,6 +19,7 @@
#include <OffloadAPI.h>
#include <cstdint>
+#include <cstdlib>
#include <mutex>
// TODO: Some plugins expect to be linked into libomptarget which defines these
@@ -47,18 +48,32 @@ struct ol_platform_impl_t {
ol_platform_backend_t BackendType;
/// Complete all pending work for this platform and perform any needed
- /// cleanup.
+ /// cleanup. Does nothing if the platform was never initialized.
///
/// After calling this function, no liboffload functions should be called with
/// this platform handle.
llvm::Error destroy();
- /// Initialize the associated plugin and devices.
+ /// Initialize the associated plugin and devices on first use.
llvm::Error init();
+ /// Get the device list, lazily initializing the platform if necessary.
+ llvm::Expected<llvm::ArrayRef<std::unique_ptr<ol_device_impl_t>>>
+ getDevices() {
+ if (llvm::Error Err = init())
+ return std::move(Err);
+ return llvm::ArrayRef(Devices);
+ }
+
/// Direct access to the plugin, may be uninitialized if accessed here.
std::unique_ptr<GenericPluginTy> Plugin;
+private:
+ llvm::Error initImpl();
+
+ // Initialize the platform once per instance of the context.
+ std::once_flag Initialized;
+ std::optional<std::pair<std::error_code, std::string>> InitError;
llvm::SmallVector<std::unique_ptr<ol_device_impl_t>> Devices;
};
@@ -77,9 +92,34 @@ struct ol_device_impl_t {
InfoTreeNode Info;
};
-llvm::Error ol_platform_impl_t::destroy() { return Plugin->deinit(); }
+llvm::Error ol_platform_impl_t::destroy() {
+ if (!Plugin || !Plugin->is_initialized())
+ return llvm::Error::success();
+ return Plugin->deinit();
+}
+
+namespace llvm::offload {
+static void shutDownAtExit();
+} // namespace llvm::offload
llvm::Error ol_platform_impl_t::init() {
+ std::call_once(Initialized, [&]() {
+ if (llvm::Error Err = initImpl())
+ llvm::handleAllErrors(std::move(Err), [&](llvm::StringError &E) {
+ InitError.emplace(E.convertToErrorCode(), E.getMessage());
+ });
+
+ // Vendor runtimes register their exit-time teardown when initialized. Exit
+ // handlers run in reverse order, so ours runs while they are still alive.
+ std::atexit(llvm::offload::shutDownAtExit);
+ });
+ if (InitError)
+ return llvm::make_error<error::OffloadError>(InitError->first,
+ InitError->second);
+ return llvm::Error::success();
+}
+
+llvm::Error ol_platform_impl_t::initImpl() {
if (!Plugin)
return llvm::Error::success();
@@ -255,6 +295,9 @@ struct OffloadContext;
// initialized
static std::atomic<OffloadContext *> OffloadContextVal;
std::mutex OffloadContextValMutex;
+// Set once the vendor runtimes may have been torn down at process exit.
+// Guarded by OffloadContextValMutex.
+static bool ShutDownAtExit = false;
struct OffloadContext {
OffloadContext(OffloadContext &) = delete;
OffloadContext(OffloadContext &&) = delete;
@@ -326,14 +369,6 @@ Error initPlugins(OffloadContext &Context, const ol_init_args_t *InitArgs) {
} while (false);
#include "Shared/Targets.def"
- // Eagerly initialize all of the plugins and devices. We need to make sure
- // that the platform is initialized at a consistent point to maintain the
- // expected teardown order in the vendor libraries.
- for (auto &Platform : Context.Platforms) {
- if (Error Err = Platform->init())
- return Err;
- }
-
Context.TracingEnabled = std::getenv("OFFLOAD_TRACE");
Context.ValidationEnabled = !std::getenv("OFFLOAD_DISABLE_VALIDATION");
@@ -343,6 +378,10 @@ Error initPlugins(OffloadContext &Context, const ol_init_args_t *InitArgs) {
Error olInit_impl(const ol_init_args_t *InitArgs) {
std::lock_guard<std::mutex> Lock(OffloadContextValMutex);
+ if (ShutDownAtExit)
+ return createOffloadError(ErrorCode::UNINITIALIZED,
+ "liboffload was shut down during process exit");
+
if (isOffloadInitialized()) {
OffloadContext::get().RefCount++;
return Plugin::success();
@@ -367,26 +406,42 @@ Error olInit_impl(const ol_init_args_t *InitArgs) {
return InitResult;
}
+static Error destroyContext(OffloadContext *Context) {
+ Error Result = Error::success();
+ for (auto &Platform : Context->Platforms)
+ if (auto Res = Platform->destroy())
+ Result = joinErrors(std::move(Result), std::move(Res));
+
+ delete Context;
+ return Result;
+}
+
Error olShutDown_impl() {
std::lock_guard<std::mutex> Lock(OffloadContextValMutex);
+ // The context may already have been released at process exit, calls that
+ // balance an earlier olInit are still valid.
+ if (!isOffloadInitialized()) {
+ if (ShutDownAtExit)
+ return Error::success();
+ return createOffloadError(ErrorCode::UNINITIALIZED,
+ "liboffload has not been initialized");
+ }
+
if (--OffloadContext::get().RefCount != 0)
return Error::success();
- Error Result = Error::success();
- auto *OldContext = OffloadContextVal.exchange(nullptr);
-
- for (auto &Platform : OldContext->Platforms) {
- // Host plugin is nullptr and has no deinit
- if (!Platform->Plugin || !Platform->Plugin->is_initialized())
- continue;
-
- if (auto Res = Platform->destroy())
- Result = joinErrors(std::move(Result), std::move(Res));
- }
+ return destroyContext(OffloadContextVal.exchange(nullptr));
+}
- delete OldContext;
- return Result;
+// Vendor libraries have specific teardown orders but can be initialized lazily.
+// In these cases we register a specific handler to shut it down with the proper
+// ordering. If this is necessary than the shutdown implementation is skipped.
+static void shutDownAtExit() {
+ std::lock_guard<std::mutex> Lock(OffloadContextValMutex);
+ ShutDownAtExit = true;
+ if (OffloadContext *Context = OffloadContextVal.exchange(nullptr))
+ consumeError(destroyContext(Context));
}
Error olGetPlatformInfoImplDetail(ol_platform_handle_t Platform,
@@ -432,6 +487,8 @@ Error olGetPlatformInfoSize_impl(ol_platform_handle_t Platform,
Error olPlatformRegisterRPCCallback_impl(ol_platform_handle_t Platform,
ol_platform_rpc_cb_t Callback) {
+ if (Error Err = Platform->init())
+ return Err;
Platform->Plugin->getRPCServer().registerCallback(Callback);
return Error::success();
}
@@ -606,7 +663,10 @@ Error olGetDeviceInfoSize_impl(ol_device_handle_t Device,
Error olIterateDevices_impl(ol_device_iterate_cb_t Callback, void *UserData) {
for (auto &Platform : OffloadContext::get().Platforms) {
- for (auto &Device : Platform->Devices) {
+ auto DevicesOrErr = Platform->getDevices();
+ if (!DevicesOrErr)
+ return DevicesOrErr.takeError();
+ for (auto &Device : *DevicesOrErr) {
if (!Callback(Device.get(), UserData)) {
return Error::success();
}
diff --git a/offload/tools/offload-tblgen/EntryPointGen.cpp b/offload/tools/offload-tblgen/EntryPointGen.cpp
index 4f76100ed2dc3..88ceff9f43c05 100644
--- a/offload/tools/offload-tblgen/EntryPointGen.cpp
+++ b/offload/tools/offload-tblgen/EntryPointGen.cpp
@@ -82,10 +82,12 @@ static void EmitEntryPointFunc(const FunctionRec &F, raw_ostream &OS) {
}
OS << ") {\n";
- // Check offload is initialized
- if (F.getName() != "olInit") {
+ // Check offload is initialized. olShutDown checks this itself as the context
+ // may have been released at process exit.
+ if (F.getName() != "olInit" && F.getName() != "olShutDown")
OS << "if (!llvm::offload::isOffloadInitialized()) return &UninitError;";
+ if (F.getName() != "olInit") {
// Emit pre-call prints
// Postpone pre-calls for olInit as tracing requires liboffload to be initialized
OS << TAB_1 "if (llvm::offload::isTracingEnabled()) {\n";
diff --git a/offload/unittests/OffloadAPI/init/olInit.cpp b/offload/unittests/OffloadAPI/init/olInit.cpp
index 4c74122e89d42..2d4e1cb6975d0 100644
--- a/offload/unittests/OffloadAPI/init/olInit.cpp
+++ b/offload/unittests/OffloadAPI/init/olInit.cpp
@@ -33,6 +33,26 @@ TEST_F(olInitTest, RepeatedInit) {
}
}
+TEST_F(olInitTest, RepeatedInitWithDevices) {
+ size_t Expected = 0;
+ for (size_t I = 0; I < 3; I++) {
+ ASSERT_SUCCESS(olInit(nullptr));
+ size_t Count = 0;
+
+ // Force the platform to be initialized.
+ ASSERT_SUCCESS(olIterateDevices(
+ [](ol_device_handle_t, void *Data) {
+ ++*static_cast<size_t *>(Data);
+ return true;
+ },
+ &Count));
+ ASSERT_SUCCESS(olShutDown());
+ if (I == 0)
+ Expected = Count;
+ ASSERT_EQ(Count, Expected);
+ }
+}
+
TEST_F(olInitTest, WithInitArgs) {
ol_init_args_t Args = OL_INIT_ARGS_INIT;
ol_platform_backend_t Backends[] = {OL_PLATFORM_BACKEND_HOST};
More information about the llvm-commits
mailing list