Skip to content

Commit 89cdd48

Browse files
authored
[Libomptarget] Remove temporary files in AMDGPU JIT impl (#77980)
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.
1 parent c58bc24 commit 89cdd48

File tree

1 file changed

+32
-18
lines changed
  • openmp/libomptarget/plugins-nextgen/amdgpu/src

1 file changed

+32
-18
lines changed

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
#include "llvm/Frontend/OpenMP/OMPConstants.h"
3939
#include "llvm/Frontend/OpenMP/OMPGridValues.h"
4040
#include "llvm/Support/Error.h"
41+
#include "llvm/Support/FileOutputBuffer.h"
4142
#include "llvm/Support/FileSystem.h"
4243
#include "llvm/Support/MemoryBuffer.h"
4344
#include "llvm/Support/Program.h"
@@ -1999,21 +2000,27 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
19992000

20002001
// TODO: We should try to avoid materialization but there seems to be no
20012002
// good linker interface w/o file i/o.
2002-
SmallString<128> LinkerOutputFilePath;
2003-
std::error_code EC = sys::fs::createTemporaryFile(
2004-
"amdgpu-pre-link-jit", ".out", LinkerOutputFilePath);
2003+
SmallString<128> LinkerInputFilePath;
2004+
std::error_code EC = sys::fs::createTemporaryFile("amdgpu-pre-link-jit",
2005+
"o", LinkerInputFilePath);
20052006
if (EC)
2006-
return createStringError(EC,
2007-
"Failed to create temporary file for linker");
2008-
2009-
SmallString<128> LinkerInputFilePath = LinkerOutputFilePath;
2010-
LinkerInputFilePath.pop_back_n(2);
2007+
return Plugin::error("Failed to create temporary file for linker");
2008+
2009+
// Write the file's contents to the output file.
2010+
Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
2011+
FileOutputBuffer::create(LinkerInputFilePath, MB->getBuffer().size());
2012+
if (!OutputOrErr)
2013+
return OutputOrErr.takeError();
2014+
std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
2015+
llvm::copy(MB->getBuffer(), Output->getBufferStart());
2016+
if (Error E = Output->commit())
2017+
return std::move(E);
20112018

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

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

20272034
std::string MCPU = "-plugin-opt=mcpu=" + getComputeUnitKind();
2028-
20292035
StringRef Args[] = {LLDPath,
20302036
"-flavor",
20312037
"gnu",
@@ -2039,12 +2045,20 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
20392045
std::string Error;
20402046
int RC = sys::ExecuteAndWait(LLDPath, Args, std::nullopt, {}, 0, 0, &Error);
20412047
if (RC)
2042-
return createStringError(inconvertibleErrorCode(),
2043-
"Linking optimized bitcode failed: %s",
2044-
Error.c_str());
2048+
return Plugin::error("Linking optimized bitcode failed: %s",
2049+
Error.c_str());
2050+
2051+
auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath);
2052+
if (!BufferOrErr)
2053+
return Plugin::error("Failed to open temporary file for lld");
2054+
2055+
// Clean up the temporary files afterwards.
2056+
if (sys::fs::remove(LinkerOutputFilePath))
2057+
return Plugin::error("Failed to remove temporary output file for lld");
2058+
if (sys::fs::remove(LinkerInputFilePath))
2059+
return Plugin::error("Failed to remove temporary input file for lld");
20452060

2046-
return std::move(
2047-
MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath.data()).get());
2061+
return std::move(*BufferOrErr);
20482062
}
20492063

20502064
/// See GenericDeviceTy::getComputeUnitKind().

0 commit comments

Comments
 (0)