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

Merged
merged 2 commits into from
Jul 25, 2025
Merged

[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
2 commits
Select commit Hold shift + click to select a range
ea5e3b5
[Offload] Erase entries from JIT cache when program is destroyed
RossBrunton Jul 15, 2025
ca0be4c
Tweaked comment a bit
RossBrunton Jul 23, 2025
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
12 changes: 9 additions & 3 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 @@ -89,11 +93,13 @@ struct JITEngine {
/// LLVM Context in which the modules will be constructed.
LLVMContext Context;

/// Output images generated from LLVM backend.
SmallVector<std::unique_ptr<MemoryBuffer>, 4> JITImages;
/// A map of embedded IR images to the buffer used to store JITed code
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.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Re-reading this thread and I think I worded things a bit confusingly; by "Image address" I mean &Image rather than Image->ImageStart.

I'm not sure there's a key that we can use to uniquely identify Images across create/destroy boundaries, nor can I see a use case for that.

Copy link
Contributor

Choose a reason for hiding this comment

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

A SHA value based on the contents of the image could do it. In that way, even the image "handler" can be created and destroyed multiple times, the contents of the image is expected to be the same.

I'll not be the blocker here. This is less ideal but I'm fine with it.

Copy link
Contributor

Choose a reason for hiding this comment

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

We need to seriously rework the image handling as whole.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I'm still unsure of why we need to keep them around after the backing image has been dropped. I can't see the user constantly recreating the same buffer many times with the same contents and expecting high performance.

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