[llvm] [openmp] Save jit image (PR #212384)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Jul 29 10:48:16 PDT 2026
https://github.com/jandrovins updated https://github.com/llvm/llvm-project/pull/212384
>From ff53181a25d791408310d297485277f5eecb2700 Mon Sep 17 00:00:00 2001
From: "Vincent A. Arcila Larrea" <arcilalarrea1 at llnl.gov>
Date: Mon, 27 Jul 2026 17:40:16 -0700
Subject: [PATCH 1/3] [offload][OpenMP] Add option to save JIT-compiled device
images
---
.../common/src/PluginInterface.cpp | 10 ++++++++++
.../kernelreplay/llvm-omp-kernel-replay.cpp | 16 ++++++++++++++++
openmp/docs/design/Runtimes.rst | 16 ++++++++++++++++
3 files changed, 42 insertions(+)
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 7b821e77df179..db480aab73c5d 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -693,6 +693,16 @@ Expected<DeviceImageTy *> GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
"failure to jit IR image");
}
Buffer = std::move(*CompiledImageOrErr);
+
+ StringEnvar JITSaveImage("LIBOMPTARGET_JIT_SAVE_IMAGE", "");
+ if (JITSaveImage.isPresent()) {
+ std::error_code EC;
+ raw_fd_ostream OS(JITSaveImage.get(), EC);
+ if (EC)
+ return Plugin::error(ErrorCode::HOST_IO, "saving JIT image file");
+ OS << Buffer->getBuffer();
+ OS.close();
+ }
} else {
Buffer = MemoryBuffer::getMemBufferCopy(InputTgtImage);
}
diff --git a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
index 3f821a33a1d9b..a63ad707ca2c1 100644
--- a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
+++ b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
@@ -73,6 +73,12 @@ static cl::opt<bool>
cl::desc("Load the recorded IR bitcode image file."),
cl::init(false), cl::cat(ReplayOptions));
+static cl::opt<bool>
+ SaveJITImageOpt("save-jit-image",
+ cl::desc("Save the JIT-compiled image next to the bitcode "
+ "image file."),
+ cl::init(false), cl::cat(ReplayOptions));
+
template <typename... ArgsTy>
Error createErr(const char *ErrFmt, ArgsTy &&...Args) {
return llvm::createStringError(llvm::inconvertibleErrorCode(), ErrFmt,
@@ -306,6 +312,16 @@ Error replayKernel() {
// Load the device image file.
Filepath.replace_extension(LoadBitcodeOpt ? "bc" : "image");
+ if (SaveJITImageOpt) {
+ if (!LoadBitcodeOpt)
+ return createErr("--save-jit-image requires --load-bitcode");
+
+ std::filesystem::path JITImageFilepath = Filepath;
+ JITImageFilepath.replace_extension("image");
+ if (setenv("LIBOMPTARGET_JIT_SAVE_IMAGE", JITImageFilepath.c_str(),
+ /*Replace=*/1) != 0)
+ return createErr("failed to configure JIT image output file");
+ }
auto ImageBufferOrErr =
MemoryBuffer::getFile(Filepath.c_str(), /*isText=*/false,
/*RequiresNullTerminator=*/false);
diff --git a/openmp/docs/design/Runtimes.rst b/openmp/docs/design/Runtimes.rst
index 4e3137abd6fb7..3b5847cd239e5 100644
--- a/openmp/docs/design/Runtimes.rst
+++ b/openmp/docs/design/Runtimes.rst
@@ -1253,6 +1253,22 @@ others:
(default 1).
* ``--num-threads=N``: Overrides the number of threads per team.
* ``--num-teams=N``: Overrides the number of teams.
+* ``--load-bitcode``: Loads the recorded IR bitcode image instead of the
+ recorded device image. The bitcode is JIT compiled for the selected device.
+* ``--save-jit-image``: Requires ``--load-bitcode`` and saves the resulting
+ JIT-compiled device image alongside the recorded bitcode as an ``.image``
+ file.
+
+For example, use the following command after modifying a recorded bitcode
+image. It JIT compiles the bitcode once, saves the executable device image, and
+replays the kernel:
+
+.. code-block:: console
+
+ $ llvm-omp-kernel-replay --load-bitcode --save-jit-image records/5681756204876336171_6652394454608725381.json
+
+Subsequent replays can omit both options and load the saved ``.image`` file
+directly, avoiding another JIT compilation.
If ``--num-threads`` or ``--num-teams`` are not specified, the replay
automatically defaults to the values used during the original recorded run. The
>From 4f41d07db071a321142731a732e2f1eb4a914b7e Mon Sep 17 00:00:00 2001
From: "Vincent A. Arcila Larrea" <arcilalarrea1 at llnl.gov>
Date: Mon, 27 Jul 2026 17:58:15 -0700
Subject: [PATCH 2/3] [offload][OpenMP] Add test for --save-jit-image
---
.../record-replay-save-jit-image.cpp | 46 +++++++++++++++++++
1 file changed, 46 insertions(+)
create mode 100644 offload/test/tools/omp-kernel-replay/record-replay-save-jit-image.cpp
diff --git a/offload/test/tools/omp-kernel-replay/record-replay-save-jit-image.cpp b/offload/test/tools/omp-kernel-replay/record-replay-save-jit-image.cpp
new file mode 100644
index 0000000000000..beeb5ff279e04
--- /dev/null
+++ b/offload/test/tools/omp-kernel-replay/record-replay-save-jit-image.cpp
@@ -0,0 +1,46 @@
+// clang-format off
+// RUN: %libomptarget-compilexx-generic -fopenmp-target-jit
+// RUN: rm -rf %t.testdir
+// RUN: mkdir -p %t.testdir
+// RUN: env LIBOMPTARGET_RECORD=1 LIBOMPTARGET_RECORD_MEMSIZE=536870912 LIBOMPTARGET_RECORD_DIR=%t.testdir %libomptarget-run-generic 2>&1 | %fcheck-generic
+// RUN: rm %t.testdir/*.image
+// RUN: ls -t %t.testdir/*.json | sed -n '1p' | grep . | xargs %omp-kernel-replay --load-bitcode --save-jit-image --verify
+// RUN: ls -t %t.testdir/*.image | grep .
+// RUN: ls -t %t.testdir/*.json | sed -n '1p' | grep . | xargs %omp-kernel-replay --verify
+// clang-format on
+
+// REQUIRES: gpu
+
+// UNSUPPORTED: aarch64-unknown-linux-gnu
+// UNSUPPORTED: x86_64-unknown-linux-gnu
+// UNSUPPORTED: s390x-ibm-linux-gnu
+// UNSUPPORTED: intelgpu
+
+#include <cstdint>
+#include <cstdio>
+
+int main() {
+ size_t Size = 1000;
+ uint64_t *Data = new uint64_t[Size];
+
+ for (size_t I = 0; I < Size; ++I) {
+ Data[I] = 20;
+ }
+
+#pragma omp target teams distribute parallel for num_teams(256) \
+ thread_limit(128) map(tofrom : Data[0 : Size])
+ for (size_t I = 0; I < Size; ++I) {
+ Data[I] = 10 + (uint64_t)I;
+ }
+
+ uint64_t Sum = 0;
+ for (size_t I = 0; I < Size; ++I) {
+ Sum += Data[I];
+ }
+
+ // CHECK: PASS
+ if (Sum == 509500)
+ printf("PASS\n");
+
+ delete[] Data;
+}
\ No newline at end of file
>From 7437e2e40d78e20ec7f56970a9677d869d3fc96d Mon Sep 17 00:00:00 2001
From: "Vincent A. Arcila Larrea" <arcilalarrea1 at llnl.gov>
Date: Wed, 29 Jul 2026 09:57:32 -0700
Subject: [PATCH 3/3] [offload][OpenMP] Move code for JIT-image saving into
JIT.(h/cpp)
---
offload/plugins-nextgen/common/include/JIT.h | 2 ++
offload/plugins-nextgen/common/src/JIT.cpp | 14 +++++++++++++-
.../plugins-nextgen/common/src/PluginInterface.cpp | 10 ----------
.../tools/kernelreplay/llvm-omp-kernel-replay.cpp | 2 +-
4 files changed, 16 insertions(+), 12 deletions(-)
diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h
index b4e3712d9c980..664b335a651d2 100644
--- a/offload/plugins-nextgen/common/include/JIT.h
+++ b/offload/plugins-nextgen/common/include/JIT.h
@@ -103,6 +103,8 @@ struct JITEngine {
StringEnvar("LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE");
StringEnvar PostOptIRModuleFileName =
StringEnvar("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE");
+ StringEnvar PostOptSaveImageFileName =
+ StringEnvar("LIBOMPTARGET_JIT_POST_OPT_SAVE_IMAGE");
UInt32Envar JITOptLevel = UInt32Envar("LIBOMPTARGET_JIT_OPT_LEVEL", 3);
BoolEnvar JITSkipOpt = BoolEnvar("LIBOMPTARGET_JIT_SKIP_OPT", false);
};
diff --git a/offload/plugins-nextgen/common/src/JIT.cpp b/offload/plugins-nextgen/common/src/JIT.cpp
index 881e27dad384c..934981bcc5648 100644
--- a/offload/plugins-nextgen/common/src/JIT.cpp
+++ b/offload/plugins-nextgen/common/src/JIT.cpp
@@ -292,5 +292,17 @@ JITEngine::process(StringRef Image, target::plugin::GenericDeviceTy &Device) {
return Device.doJITPostProcessing(std::move(MB));
};
- return compile(Image, ComputeUnitKind, PostProcessing);
+ auto ImageOrError = compile(Image, ComputeUnitKind, PostProcessing);
+
+ if (PostOptSaveImageFileName.isPresent() && ImageOrError) {
+ std::error_code EC;
+ raw_fd_ostream OS(PostOptSaveImageFileName.get(), EC);
+ if (EC)
+ return createStringError(
+ error::ErrorCode::HOST_IO,
+ "saving JIT image file\n");
+ OS << ImageOrError.get()->getBuffer();
+ OS.close();
+ }
+ return ImageOrError;
}
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index db480aab73c5d..7b821e77df179 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -693,16 +693,6 @@ Expected<DeviceImageTy *> GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
"failure to jit IR image");
}
Buffer = std::move(*CompiledImageOrErr);
-
- StringEnvar JITSaveImage("LIBOMPTARGET_JIT_SAVE_IMAGE", "");
- if (JITSaveImage.isPresent()) {
- std::error_code EC;
- raw_fd_ostream OS(JITSaveImage.get(), EC);
- if (EC)
- return Plugin::error(ErrorCode::HOST_IO, "saving JIT image file");
- OS << Buffer->getBuffer();
- OS.close();
- }
} else {
Buffer = MemoryBuffer::getMemBufferCopy(InputTgtImage);
}
diff --git a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
index a63ad707ca2c1..231ad0dd75bd1 100644
--- a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
+++ b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
@@ -318,7 +318,7 @@ Error replayKernel() {
std::filesystem::path JITImageFilepath = Filepath;
JITImageFilepath.replace_extension("image");
- if (setenv("LIBOMPTARGET_JIT_SAVE_IMAGE", JITImageFilepath.c_str(),
+ if (setenv("LIBOMPTARGET_JIT_POST_OPT_SAVE_IMAGE", JITImageFilepath.c_str(),
/*Replace=*/1) != 0)
return createErr("failed to configure JIT image output file");
}
More information about the llvm-commits
mailing list