[llvm] [openmp] [Offload] Save jit image (PR #212384)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Jul 31 18:05:27 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/4] [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/4] [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 98b1ef8fea917611f077c3434609557b45847436 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/4] [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 | 13 ++++++++++++-
.../plugins-nextgen/common/src/PluginInterface.cpp | 10 ----------
.../tools/kernelreplay/llvm-omp-kernel-replay.cpp | 2 +-
4 files changed, 15 insertions(+), 12 deletions(-)
diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h
index b4e3712d9c980..e22a7e435c0e1 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..1d5cfc77f0e75 100644
--- a/offload/plugins-nextgen/common/src/JIT.cpp
+++ b/offload/plugins-nextgen/common/src/JIT.cpp
@@ -292,5 +292,16 @@ 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");
}
>From 8b2626650f47067a72670d6e252b697a9c59d716 Mon Sep 17 00:00:00 2001
From: "Vincent A. Arcila Larrea" <arcilalarrea1 at llnl.gov>
Date: Fri, 31 Jul 2026 18:04:47 -0700
Subject: [PATCH 4/4] [offload] Save JIT image via environment variable only
---
offload/plugins-nextgen/common/include/JIT.h | 4 +-
offload/plugins-nextgen/common/src/JIT.cpp | 4 +-
offload/test/jit/save_image.c | 18 ++++++++
.../record-replay-save-jit-image.cpp | 46 -------------------
.../kernelreplay/llvm-omp-kernel-replay.cpp | 16 -------
openmp/docs/design/Runtimes.rst | 25 ++++------
6 files changed, 31 insertions(+), 82 deletions(-)
create mode 100644 offload/test/jit/save_image.c
delete mode 100644 offload/test/tools/omp-kernel-replay/record-replay-save-jit-image.cpp
diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h
index e22a7e435c0e1..96dea06335730 100644
--- a/offload/plugins-nextgen/common/include/JIT.h
+++ b/offload/plugins-nextgen/common/include/JIT.h
@@ -103,8 +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");
+ StringEnvar SaveImageFileName =
+ StringEnvar("LIBOMPTARGET_JIT_SAVE_IMAGE_FILENAME");
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 1d5cfc77f0e75..15f4fc843188d 100644
--- a/offload/plugins-nextgen/common/src/JIT.cpp
+++ b/offload/plugins-nextgen/common/src/JIT.cpp
@@ -294,9 +294,9 @@ JITEngine::process(StringRef Image, target::plugin::GenericDeviceTy &Device) {
auto ImageOrError = compile(Image, ComputeUnitKind, PostProcessing);
- if (PostOptSaveImageFileName.isPresent() && ImageOrError) {
+ if (SaveImageFileName.isPresent() && ImageOrError) {
std::error_code EC;
- raw_fd_ostream OS(PostOptSaveImageFileName.get(), EC);
+ raw_fd_ostream OS(SaveImageFileName.get(), EC);
if (EC)
return createStringError(error::ErrorCode::HOST_IO,
"saving JIT image file\n");
diff --git a/offload/test/jit/save_image.c b/offload/test/jit/save_image.c
new file mode 100644
index 0000000000000..8b3a57152e397
--- /dev/null
+++ b/offload/test/jit/save_image.c
@@ -0,0 +1,18 @@
+// clang-format off
+// RUN: %libomptarget-compileopt-generic -fopenmp-target-jit
+// RUN: rm -f %t.image
+// RUN: env LIBOMPTARGET_JIT_SAVE_IMAGE_FILENAME=%t.image %libomptarget-run-generic
+// RUN: test -s %t.image
+// clang-format on
+
+// REQUIRES: gpu
+// XFAIL: intelgpu
+
+int main() {
+ int X = 0;
+
+#pragma omp target map(tofrom : X)
+ { X = 1; }
+
+ return X != 1;
+}
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
deleted file mode 100644
index beeb5ff279e04..0000000000000
--- a/offload/test/tools/omp-kernel-replay/record-replay-save-jit-image.cpp
+++ /dev/null
@@ -1,46 +0,0 @@
-// 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
diff --git a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
index 231ad0dd75bd1..3f821a33a1d9b 100644
--- a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
+++ b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
@@ -73,12 +73,6 @@ 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,
@@ -312,16 +306,6 @@ 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_POST_OPT_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 3b5847cd239e5..14e300a0f531f 100644
--- a/openmp/docs/design/Runtimes.rst
+++ b/openmp/docs/design/Runtimes.rst
@@ -741,6 +741,7 @@ variables is defined below.
* ``LIBOMPTARGET_JIT_REPLACEMENT_MODULE=<in:Filename> (LLVM-IR file)``
* ``LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE=<out:Filename> (LLVM-IR file)``
* ``LIBOMPTARGET_JIT_POST_OPT_IR_MODULE=<out:Filename> (LLVM-IR file)``
+ * ``LIBOMPTARGET_JIT_SAVE_IMAGE_FILENAME=<out:Filename> (device image file)``
* ``LIBOMPTARGET_MIN_THREADS_FOR_LOW_TRIP_COUNT=<Num> (default: 32)``
* ``LIBOMPTARGET_REUSE_BLOCKS_FOR_HIGH_TRIP_COUNT=[TRUE/FALSE] (default TRUE)``
* ``OFFLOAD_TRACK_ALLOCATION_TRACES=[TRUE/FALSE] (default FALSE)``
@@ -1162,6 +1163,14 @@ which the LLVM-IR module is written. The module can be the analyzed, and
transformed and loaded back into the JIT pipeline via
:ref:`LIBOMPTARGET_JIT_REPLACEMENT_MODULE`.
+.. _libomptarget_jit_save_image_filename:
+
+LIBOMPTARGET_JIT_SAVE_IMAGE_FILENAME
+""""""""""""""""""""""""""""""""""""
+
+This environment variable can be used to save the device image produced by the
+device JIT after target-specific post-processing. The value is expected to be a
+filename into which the binary device image is written.
LIBOMPTARGET_MIN_THREADS_FOR_LOW_TRIP_COUNT
"""""""""""""""""""""""""""""""""""""""""""
@@ -1253,22 +1262,6 @@ 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
More information about the llvm-commits
mailing list