[llvm] [OFFLOAD][L0] Improve cleanup on errors (PR #188251)

[Alex Duran via llvm-commits]

https://github.com/adurang created https://github.com/llvm/llvm-project/pull/188251

Additional cleanup improvements on error conditions (in addition to those in #187597):

  * Fixed incomplete cleanup in L0Context::init()
  * Fixed build log leak in addModule()
  * Fixed context inconsistent state in findDevices()

Disclaimer: The base of this PR was generated by Claude and adjusted by me afterwards.

>From fdb03dc2f27cfd51a28cf345e053dab7ff58b770 Mon Sep 17 00:00:00 2001
From: "Duran, Alex" <alejandro.duran at intel.com>
Date: Tue, 24 Mar 2026 06:41:38 -0700
Subject: [PATCH] [OFFLOAD][L0] Improve cleanup on errors

Additional cleanup improvements (in addition to those in #187597):

  * Fixed incomplete cleanup in L0Context::init()
  * Fixed build log leak in addModule()
  * Fixed context inconsistent state in findDevices()
---
 .../plugins-nextgen/level_zero/src/L0Context.cpp | 16 ++++++++++++++--
 .../plugins-nextgen/level_zero/src/L0Plugin.cpp  |  5 ++++-
 .../plugins-nextgen/level_zero/src/L0Program.cpp | 15 +++++++++++++--
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/offload/plugins-nextgen/level_zero/src/L0Context.cpp b/offload/plugins-nextgen/level_zero/src/L0Context.cpp
index 3d1588bc5fecc..9608476871310 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Context.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Context.cpp
@@ -16,16 +16,28 @@
 namespace llvm::omp::target::plugin {
 
 Error L0ContextTy::init() {
+  auto cleanupOnError = [&]() {
+    if (zeContext) {
+      zeContextDestroy(zeContext);
+      zeContext = nullptr;
+    }
+  };
   CALL_ZE_RET_ERROR(zeDriverGetApiVersion, zeDriver, &APIVersion);
   ODBG(OLDT_Init) << "Driver API version is "
                   << llvm::format(PRIx32, APIVersion);
 
   ze_context_desc_t Desc{ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
   CALL_ZE_RET_ERROR(zeContextCreate, zeDriver, &Desc, &zeContext);
-  if (auto Err = EventPool.init(zeContext, 0))
+  if (auto Err = EventPool.init(zeContext, 0)) {
+    cleanupOnError();
     return Err;
-  if (auto Err = HostMemAllocator.initHostPool(*this, Plugin.getOptions()))
+  }
+  if (auto Err = HostMemAllocator.initHostPool(*this, Plugin.getOptions())) {
+    if (auto DeinitErr = EventPool.deinit())
+      Err = joinErrors(std::move(Err), std::move(DeinitErr));
+    cleanupOnError();
     return Err;
+  }
   return Plugin::success();
 }
 
diff --git a/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp b/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp
index 7dc0b919ca5a5..4900b55eeab7b 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp
@@ -61,8 +61,11 @@ Expected<int32_t> LevelZeroPluginTy::findDevices() {
     // We have a driver that supports at least one device.
     ContextList.emplace_back(*this, Driver, DriverId);
     auto &DrvInfo = ContextList.back();
-    if (auto Err = DrvInfo.init())
+    if (auto Err = DrvInfo.init()) {
+      // Remove the partially initialized context from the list
+      ContextList.pop_back();
       return std::move(Err);
+    }
     llvm::SmallVector<ze_device_handle_t> FoundDevices(DeviceCount);
     CALL_ZE_RET_ERROR(zeDeviceGet, Driver, &DeviceCount, FoundDevices.data());
 
diff --git a/offload/plugins-nextgen/level_zero/src/L0Program.cpp b/offload/plugins-nextgen/level_zero/src/L0Program.cpp
index 60a71c9ad65a2..7adbcbe56fa80 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Program.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Program.cpp
@@ -83,8 +83,19 @@ Error L0ProgramBuilderTy::addModule(size_t Size, const uint8_t *Image,
   ModuleDesc.pInputModule = Image;
   ModuleDesc.pBuildFlags = BuildOptions.c_str();
   ModuleDesc.pConstants = &SpecConstants;
-  CALL_ZE_RET_ERROR(zeModuleCreate, l0Device.getZeContext(),
-                    l0Device.getZeDevice(), &ModuleDesc, &Module, &BuildLog);
+  Error CreateErrors = Error::success();
+  auto handleError = [&](Error Err) {
+    if (BuildLog)
+      zeModuleBuildLogDestroy(BuildLog);
+    CreateErrors = joinErrors(std::move(CreateErrors), std::move(Err));
+  };
+  CALL_ZE_HANDLE_ERROR(handleError, zeModuleCreate, l0Device.getZeContext(),
+                       l0Device.getZeDevice(), &ModuleDesc, &Module, &BuildLog);
+  if (CreateErrors)
+    return CreateErrors;
+
+  if (BuildLog)
+    zeModuleBuildLogDestroy(BuildLog);
 
   // Check if module link is required. We do not need this check for
   // library module.