[offload] Add the ability to record and replay IR bitcode modules (#207278)

This PR allows the user to use JIT with kernel record & replay by
expanding the latter to also work with IR bitcode images.

- Add a new field to `DeviceImageTy` for storing the IR image.
- Expand `RecordReplayTy::FileTy` to also include IR bitcode and
associate it with the file extension `.bc`
- Add the `--load-bitcode` command line option to
`llvm-omp-kernel-replay` tool.

---------

Co-authored-by: Giorgi Gvalia <gvalia1@llnl.gov>

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -375,6 +375,9 @@ class DeviceImageTy {
   /// The managed image data.
   std::unique_ptr<MemoryBuffer> Image;
 
+  /// An optional buffer for an IR image.
+  std::unique_ptr<MemoryBuffer> DeviceIRImage;
+
   /// Reference to the device this image is loaded on.
   GenericDeviceTy &Device;
 
@@ -402,6 +405,21 @@ public:
     return MemoryBufferRef(StringRef((const char *)getStart(), getSize()),
                            "Image");
   }
+
+  /// Get a memory buffer reference to the whole IR image.
+  MemoryBufferRef getIRImageMemoryBuffer() const {
+    if (DeviceIRImage)
+      return MemoryBufferRef(*DeviceIRImage);
+
+    return MemoryBufferRef();
+  }
+
+  bool hasIRImage() const { return DeviceIRImage != nullptr; }
+
+  /// Set the IR image.
+  void setIRImage(std::unique_ptr<MemoryBuffer> ImageBuffer) {
+    DeviceIRImage = std::move(ImageBuffer);
+  }
 };
 
 /// Class implementing common functionalities of offload kernels. Each plugin

offload/plugins-nextgen/common/include/RecordReplay.h

@@ -57,7 +57,8 @@ public:
     EpilogueSnapshot,
     Descriptor,
     Globals,
-    Program
+    Program,
+    IRImage
   };
 
   struct HandleTy {
@@ -296,7 +297,8 @@ private:
   Error recordGlobals(StringRef Filename);
 
   /// Record the device image to a file.
-  Error recordImage(const GenericKernelTy &Kernel, StringRef Filename);
+  Error recordImage(const GenericKernelTy &Kernel, StringRef Filename,
+                    StringRef IRImageFilename);
 };
 
 } // namespace plugin

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -695,6 +695,9 @@ Expected<DeviceImageTy *> GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
     return ImageOrErr.takeError();
   DeviceImageTy *Image = *ImageOrErr;
 
+  if (identify_magic(InputTgtImage) == file_magic::bitcode)
+    Image->setIRImage(MemoryBuffer::getMemBufferCopy(InputTgtImage));
+
   // Add the image to list.
   LoadedImages.push_back(Image);
 

offload/plugins-nextgen/common/src/RecordReplay.cpp

@@ -246,7 +246,9 @@ Error NativeRecordReplayTy::recordPrologueImpl(
     return Err;
 
   SmallString<128> ImageFilename = getFilename(Instance, FileTy::Program);
-  return recordImage(Kernel, ImageFilename.c_str());
+
+  SmallString<128> IRImageFilename = getFilename(Instance, FileTy::IRImage);
+  return recordImage(Kernel, ImageFilename.c_str(), IRImageFilename.c_str());
 }
 
 Error NativeRecordReplayTy::recordEpilogueImpl(const GenericKernelTy &Kernel,
@@ -323,6 +325,8 @@ StringRef NativeRecordReplayTy::getExtension(FileTy FileType) {
     return "globals";
   case FileTy::Program:
     return "image";
+  case FileTy::IRImage:
+    return "bc";
   }
   return "";
 }
@@ -367,13 +371,23 @@ Error NativeRecordReplayTy::recordSnapshot(StringRef Filename) {
 }
 
 Error NativeRecordReplayTy::recordImage(const GenericKernelTy &Kernel,
-                                        StringRef Filename) {
+                                        StringRef Filename,
+                                        StringRef IRImageFilename) {
   std::error_code EC;
   raw_fd_ostream OS(Filename, EC);
   if (EC)
     return Plugin::error(ErrorCode::HOST_IO, "saving image file");
   OS << Kernel.getImage().getMemoryBuffer().getBuffer();
   OS.close();
+
+  if (!IRImageFilename.empty() && Kernel.getImage().hasIRImage()) {
+    raw_fd_ostream IROS(IRImageFilename, EC);
+    if (EC)
+      return Plugin::error(ErrorCode::HOST_IO, "saving IR image file");
+    IROS << Kernel.getImage().getIRImageMemoryBuffer().getBuffer();
+    IROS.close();
+  }
+
   return Plugin::success();
 }
 

offload/test/tools/omp-kernel-replay/record-replay-ir-bitcode.cpp

@@ -0,0 +1,44 @@
+// 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: ls -t %t.testdir/*.bc | grep .
+// RUN: ls -t %t.testdir/*.json | sed -n '1p' | grep . | xargs %omp-kernel-replay --load-bitcode --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;
+}

offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp

@@ -68,6 +68,11 @@ static cl::opt<bool>
                     cl::desc("Ignore thread and team limits (unrecommended)."),
                     cl::init(false), cl::cat(ReplayOptions));
 
+static cl::opt<bool>
+    LoadBitcodeOpt("load-bitcode",
+                   cl::desc("Load the recorded IR 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,
@@ -300,7 +305,7 @@ Error replayKernel() {
   }
 
   // Load the device image file.
-  Filepath.replace_extension("image");
+  Filepath.replace_extension(LoadBitcodeOpt ? "bc" : "image");
   auto ImageBufferOrErr =
       MemoryBuffer::getFile(Filepath.c_str(), /*isText=*/false,
                             /*RequiresNullTerminator=*/false);