Skip to content

[Offload] Erase entries from JIT cache when program is destroyed #148847

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+28 −11
Open

[Offload] Erase entries from JIT cache when program is destroyed #148847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions offload/plugins-nextgen/common/include/JIT.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -90,10 +94,12 @@ struct JITEngine {
LLVMContext Context;

/// Output images generated from LLVM backend.
SmallVector<std::unique_ptr<MemoryBuffer>, 4> JITImages;
DenseMap<const __tgt_device_image *, std::unique_ptr<MemoryBuffer>>
JITImages;

/// A map of embedded IR images to JITed images.
DenseMap<const __tgt_device_image *, __tgt_device_image *> TgtImageMap;
DenseMap<const __tgt_device_image *, std::unique_ptr<__tgt_device_image>>
TgtImageMap;
};

/// Map from (march) "CPUs" (e.g., sm_80, or gfx90a), which we call compute
Expand Down
26 changes: 17 additions & 9 deletions offload/plugins-nextgen/common/src/JIT.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't follow in what cases this Image is not a unique one?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The data Image points to happens to be inside ol_program_impl_t, but something similar to this:

Image *MyImage = new Image();
delete MyImage;
Image *MyImage2 = new Image();
// MyImage may equal MyImage2

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So the image can be created on the fly. In that case, we probably still want to cache that, but use a different key that can effectively tell the two images apart.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The JIT engine uses the Image address to identify input images; once the Image is dropped, we lose the ability to look it up in the cache, so there's no reason to keep the entry around. What use case are you thinking of?

Copy link
Contributor

@shiltian shiltian Jul 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

but your case demonstrates that image address could be not unique, even for the "same" image.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As long as the image pointer is actually alive, it's a unique identifier for the JIT'ed binary. The issue only happens once the input image is free'd.

return JITedImage;
if (CUI.TgtImageMap.contains(&Image))
return CUI.TgtImageMap[&Image].get();

auto ObjMBOrErr = getOrCreateObjFile(Image, CUI.Context, ComputeUnitKind);
if (!ObjMBOrErr)
Expand All @@ -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<char *>(ImageMB->getBufferStart());
JITedImage->ImageEnd = const_cast<char *>(ImageMB->getBufferEnd());

return JITedImage;
return JITedImage.get();
}

Expected<const __tgt_device_image *>
Expand All @@ -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<std::mutex> Lock(ComputeUnitMapMutex);
const std::string &ComputeUnitKind = Device.getComputeUnitKind();
ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];

CUI.TgtImageMap.erase(&Image);
CUI.JITImages.erase(&Image);
}
3 changes: 3 additions & 0 deletions offload/plugins-nextgen/common/src/PluginInterface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -854,6 +854,9 @@ Error GenericDeviceTy::unloadBinary(DeviceImageTy *Image) {
return Err;
}

if (Image->getTgtImageBitcode())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I forget how this works, we have the image passed in by the user and the one created by the backend right? I'm wondering if we should just check the magic bytes there.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we are compiling bitcode, then this will be set to a __tgt_device_image * for the original bitcode that olCreateProgram copied. It is nullptr if no compilation took place (and so we don't need to tell the JIT to remove anything.

Plugin.getJIT().erase(*Image->getTgtImageBitcode(), Image->getDevice());

return unloadBinaryImpl(Image);
}

Expand Down
Loading