[Openmp-commits] [openmp] 513a02e - [Offload] Save jit image (#212384)

[via Openmp-commits]

Author: jandrovins
Date: 2026-08-03T21:33:56-07:00
New Revision: 513a02ee150335a84c39ab2e2d0f965f1b0e260e

URL: https://github.com/llvm/llvm-project/commit/513a02ee150335a84c39ab2e2d0f965f1b0e260e
DIFF: https://github.com/llvm/llvm-project/commit/513a02ee150335a84c39ab2e2d0f965f1b0e260e.diff

LOG: [Offload] Save jit image (#212384)

This PR adds the envar LIBOMPTARGET_JIT_SAVE_IMAGE_FILENAME
to indicate that the device image generated after JIT compilation
should be saved to a file.

Added: 
    offload/test/jit/save_image.c

Modified: 
    offload/plugins-nextgen/common/include/JIT.h
    offload/plugins-nextgen/common/src/JIT.cpp
    openmp/docs/design/Runtimes.rst

Removed: 
    


################################################################################
diff  --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h
index b4e3712d9c980..96dea06335730 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 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 9153059887757..b78c9b1768139 100644
--- a/offload/plugins-nextgen/common/src/JIT.cpp
+++ b/offload/plugins-nextgen/common/src/JIT.cpp
@@ -300,5 +300,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 (SaveImageFileName.isPresent() && ImageOrError) {
+    std::error_code EC;
+    raw_fd_ostream OS(SaveImageFileName.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/test/jit/save_image.c b/offload/test/jit/save_image.c
new file mode 100644
index 0000000000000..66936e196c528
--- /dev/null
+++ b/offload/test/jit/save_image.c
@@ -0,0 +1,20 @@
+// 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/openmp/docs/design/Runtimes.rst b/openmp/docs/design/Runtimes.rst
index 4e3137abd6fb7..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
 """""""""""""""""""""""""""""""""""""""""""