From ea5e3b5655d319ff3f1b17bcbea10f8173fddd3f Mon Sep 17 00:00:00 2001 From: Ross Brunton Date: Tue, 15 Jul 2025 14:08:35 +0100 Subject: [PATCH 1/2] [Offload] Erase entries from JIT cache when program is destroyed When `unloadBinary` is called, any entries in the JITEngine's cache for that binary will be cleared. This fixes a nasty issue with liboffload program handles. If two handles happen to have had the same address (after one was free'd, for example), the cache would be hit and return the wrong program. --- offload/plugins-nextgen/common/include/JIT.h | 10 +++++-- offload/plugins-nextgen/common/src/JIT.cpp | 26 ++++++++++++------- .../common/src/PluginInterface.cpp | 3 +++ 3 files changed, 28 insertions(+), 11 deletions(-) diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h index 8c530436a754b..68a6d039c4641 100644 --- a/offload/plugins-nextgen/common/include/JIT.h +++ b/offload/plugins-nextgen/common/include/JIT.h @@ -55,6 +55,10 @@ struct JITEngine { process(const __tgt_device_image &Image, target::plugin::GenericDeviceTy &Device); + /// Remove \p Image from the jit engine's cache + void erase(const __tgt_device_image &Image, + target::plugin::GenericDeviceTy &Device); + private: /// Compile the bitcode image \p Image and generate the binary image that can /// be loaded to the target device of the triple \p Triple architecture \p @@ -90,10 +94,12 @@ struct JITEngine { LLVMContext Context; /// Output images generated from LLVM backend. - SmallVector, 4> JITImages; + DenseMap> + JITImages; /// A map of embedded IR images to JITed images. - DenseMap TgtImageMap; + DenseMap> + TgtImageMap; }; /// Map from (march) "CPUs" (e.g., sm_80, or gfx90a), which we call compute diff --git a/offload/plugins-nextgen/common/src/JIT.cpp b/offload/plugins-nextgen/common/src/JIT.cpp index c82a06e36d8f9..00720fa2d8103 100644 --- a/offload/plugins-nextgen/common/src/JIT.cpp +++ b/offload/plugins-nextgen/common/src/JIT.cpp @@ -285,8 +285,8 @@ JITEngine::compile(const __tgt_device_image &Image, // Check if we JITed this image for the given compute unit kind before. ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind]; - if (__tgt_device_image *JITedImage = CUI.TgtImageMap.lookup(&Image)) - return JITedImage; + if (CUI.TgtImageMap.contains(&Image)) + return CUI.TgtImageMap[&Image].get(); auto ObjMBOrErr = getOrCreateObjFile(Image, CUI.Context, ComputeUnitKind); if (!ObjMBOrErr) @@ -296,17 +296,15 @@ JITEngine::compile(const __tgt_device_image &Image, if (!ImageMBOrErr) return ImageMBOrErr.takeError(); - CUI.JITImages.push_back(std::move(*ImageMBOrErr)); - __tgt_device_image *&JITedImage = CUI.TgtImageMap[&Image]; - JITedImage = new __tgt_device_image(); + CUI.JITImages.insert({&Image, std::move(*ImageMBOrErr)}); + auto &ImageMB = CUI.JITImages[&Image]; + CUI.TgtImageMap.insert({&Image, std::make_unique<__tgt_device_image>()}); + auto &JITedImage = CUI.TgtImageMap[&Image]; *JITedImage = Image; - - auto &ImageMB = CUI.JITImages.back(); - JITedImage->ImageStart = const_cast(ImageMB->getBufferStart()); JITedImage->ImageEnd = const_cast(ImageMB->getBufferEnd()); - return JITedImage; + return JITedImage.get(); } Expected @@ -324,3 +322,13 @@ JITEngine::process(const __tgt_device_image &Image, return &Image; } + +void JITEngine::erase(const __tgt_device_image &Image, + target::plugin::GenericDeviceTy &Device) { + std::lock_guard Lock(ComputeUnitMapMutex); + const std::string &ComputeUnitKind = Device.getComputeUnitKind(); + ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind]; + + CUI.TgtImageMap.erase(&Image); + CUI.JITImages.erase(&Image); +} diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp index 81b9d423e13d8..94a050b559efe 100644 --- a/offload/plugins-nextgen/common/src/PluginInterface.cpp +++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp @@ -854,6 +854,9 @@ Error GenericDeviceTy::unloadBinary(DeviceImageTy *Image) { return Err; } + if (Image->getTgtImageBitcode()) + Plugin.getJIT().erase(*Image->getTgtImageBitcode(), Image->getDevice()); + return unloadBinaryImpl(Image); } From ca0be4c6ac4e4b1a9c266f7e114fa5daf6d9c61c Mon Sep 17 00:00:00 2001 From: Ross Brunton Date: Wed, 23 Jul 2025 11:27:56 +0100 Subject: [PATCH 2/2] Tweaked comment a bit --- offload/plugins-nextgen/common/include/JIT.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h index 68a6d039c4641..d62516d20764a 100644 --- a/offload/plugins-nextgen/common/include/JIT.h +++ b/offload/plugins-nextgen/common/include/JIT.h @@ -93,7 +93,7 @@ struct JITEngine { /// LLVM Context in which the modules will be constructed. LLVMContext Context; - /// Output images generated from LLVM backend. + /// A map of embedded IR images to the buffer used to store JITed code DenseMap> JITImages;