Skip to content

Commit

Permalink
[Libomptarget] Remove temporary files in AMDGPU JIT impl (#77980)
Browse files Browse the repository at this point in the history
Summary:
This patch cleans up some of the JIT handling for AMDGPU as well as
removing its temporary files. Previously these would be left in the
temporary directory after the program was run. This costs some extra
time, but the correct solution to avoid that is to create a sufficient
entrypoint into `ld.lld` that we can simply pass a memory buffer into.
  • Loading branch information
jhuber6 committed Jan 16, 2024
1 parent c58bc24 commit 89cdd48
Showing 1 changed file with 32 additions and 18 deletions.
50 changes: 32 additions & 18 deletions openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
#include "llvm/Frontend/OpenMP/OMPConstants.h"
#include "llvm/Frontend/OpenMP/OMPGridValues.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FileOutputBuffer.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Program.h"
Expand Down Expand Up @@ -1999,21 +2000,27 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {

// TODO: We should try to avoid materialization but there seems to be no
// good linker interface w/o file i/o.
SmallString<128> LinkerOutputFilePath;
std::error_code EC = sys::fs::createTemporaryFile(
"amdgpu-pre-link-jit", ".out", LinkerOutputFilePath);
SmallString<128> LinkerInputFilePath;
std::error_code EC = sys::fs::createTemporaryFile("amdgpu-pre-link-jit",
"o", LinkerInputFilePath);
if (EC)
return createStringError(EC,
"Failed to create temporary file for linker");

SmallString<128> LinkerInputFilePath = LinkerOutputFilePath;
LinkerInputFilePath.pop_back_n(2);
return Plugin::error("Failed to create temporary file for linker");

// Write the file's contents to the output file.
Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
FileOutputBuffer::create(LinkerInputFilePath, MB->getBuffer().size());
if (!OutputOrErr)
return OutputOrErr.takeError();
std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
llvm::copy(MB->getBuffer(), Output->getBufferStart());
if (Error E = Output->commit())
return std::move(E);

auto FD = raw_fd_ostream(LinkerInputFilePath.data(), EC);
SmallString<128> LinkerOutputFilePath;
EC = sys::fs::createTemporaryFile("amdgpu-pre-link-jit", "so",
LinkerOutputFilePath);
if (EC)
return createStringError(EC, "Failed to open temporary file for linker");
FD.write(MB->getBufferStart(), MB->getBufferSize());
FD.close();
return Plugin::error("Failed to create temporary file for linker");

const auto &ErrorOrPath = sys::findProgramByName("lld");
if (!ErrorOrPath)
Expand All @@ -2025,7 +2032,6 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
"Using `%s` to link JITed amdgcn ouput.", LLDPath.c_str());

std::string MCPU = "-plugin-opt=mcpu=" + getComputeUnitKind();

StringRef Args[] = {LLDPath,
"-flavor",
"gnu",
Expand All @@ -2039,12 +2045,20 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
std::string Error;
int RC = sys::ExecuteAndWait(LLDPath, Args, std::nullopt, {}, 0, 0, &Error);
if (RC)
return createStringError(inconvertibleErrorCode(),
"Linking optimized bitcode failed: %s",
Error.c_str());
return Plugin::error("Linking optimized bitcode failed: %s",
Error.c_str());

auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath);
if (!BufferOrErr)
return Plugin::error("Failed to open temporary file for lld");

// Clean up the temporary files afterwards.
if (sys::fs::remove(LinkerOutputFilePath))
return Plugin::error("Failed to remove temporary output file for lld");
if (sys::fs::remove(LinkerInputFilePath))
return Plugin::error("Failed to remove temporary input file for lld");

return std::move(
MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath.data()).get());
return std::move(*BufferOrErr);
}

/// See GenericDeviceTy::getComputeUnitKind().
Expand Down

0 comments on commit 89cdd48

Please sign in to comment.