diff --git a/amd/comgr/src/comgr-hotswap-b0a0.cpp b/amd/comgr/src/comgr-hotswap-b0a0.cpp index 9551ae1f13f83..65f1afde36d01 100644 --- a/amd/comgr/src/comgr-hotswap-b0a0.cpp +++ b/amd/comgr/src/comgr-hotswap-b0a0.cpp @@ -522,6 +522,22 @@ static void runScratchVerification(WritableMemoryBuffer &OutBuf, << "scratch conflicts\n"; } +static std::unique_ptr +copyOutputBuffer(const void *Data, size_t Size, StringRef CopyKind) { + std::unique_ptr Result = + WritableMemoryBuffer::getNewUninitMemBuffer(Size); + if (!Result) { + log() << "hotswap: error: retargetCodeObject: " + << "getNewUninitMemBuffer(" << Size + << ") failed (out of memory) for the " << CopyKind + << " output copy.\n"; + return nullptr; + } + + std::memcpy(Result->getBufferStart(), Data, Size); + return Result; +} + // -- retargetCodeObject ------------------------------------------------------- amd_comgr_status_t retargetCodeObject(const void *ElfData, size_t ElfSize, @@ -535,14 +551,9 @@ amd_comgr_status_t retargetCodeObject(const void *ElfData, size_t ElfSize, if (!Options.RunB0A0Patches && !Options.RunEntryTrampolines) { std::unique_ptr Result = - WritableMemoryBuffer::getNewUninitMemBuffer(ElfSize); - if (!Result) { - log() << "hotswap: error: retargetCodeObject: " - << "getNewUninitMemBuffer(" << ElfSize - << ") failed (out of memory) for the no-op output copy.\n"; + copyOutputBuffer(ElfData, ElfSize, "no-op"); + if (!Result) return AMD_COMGR_STATUS_ERROR_OUT_OF_RESOURCES; - } - std::memcpy(Result->getBufferStart(), ElfData, ElfSize); Out = std::move(Result); return AMD_COMGR_STATUS_SUCCESS; } @@ -638,14 +649,9 @@ amd_comgr_status_t retargetCodeObject(const void *ElfData, size_t ElfSize, EntryFixups)) return AMD_COMGR_STATUS_ERROR; } else { - Result = WritableMemoryBuffer::getNewUninitMemBuffer(ElfSize); - if (!Result) { - log() << "hotswap: error: retargetCodeObject: " - << "getNewUninitMemBuffer(" << ElfSize - << ") failed (out of memory) for the patched output copy.\n"; + Result = copyOutputBuffer(Buf.data(), ElfSize, "patched"); + if (!Result) return AMD_COMGR_STATUS_ERROR_OUT_OF_RESOURCES; - } - std::memcpy(Result->getBufferStart(), Buf.data(), ElfSize); } if (!ScratchPatches.empty()) diff --git a/amd/comgr/src/comgr-hotswap-elf.cpp b/amd/comgr/src/comgr-hotswap-elf.cpp index 51e1538f6c9b8..a80a7bb8264f6 100644 --- a/amd/comgr/src/comgr-hotswap-elf.cpp +++ b/amd/comgr/src/comgr-hotswap-elf.cpp @@ -18,7 +18,6 @@ #include "llvm/ADT/StringExtras.h" #include "llvm/ADT/Twine.h" #include "llvm/BinaryFormat/MsgPackDocument.h" -#include "llvm/Support/CheckedArithmetic.h" #include #include @@ -42,16 +41,6 @@ enum class MetadataSgprUpdateStatus { Error, }; -static std::optional checkedAdd(uint64_t LHS, uint64_t RHS, - StringRef Context) { - std::optional Result = checkedAddUnsigned(LHS, RHS); - if (Result) - return Result; - - log() << "hotswap: error: " << Context << " overflows uint64_t.\n"; - return std::nullopt; -} - static std::optional checkedAlignToSize(size_t Value, uint64_t Alignment, StringRef Context) { if (Alignment <= 1) @@ -61,7 +50,7 @@ checkedAlignToSize(size_t Value, uint64_t Alignment, StringRef Context) { if (Remainder == 0) return Value; std::optional Aligned = - checkedAdd(Value64, Alignment - Remainder, Context); + checkedAddUint64(Value64, Alignment - Remainder, Context); if (!Aligned) return std::nullopt; if (*Aligned > static_cast(std::numeric_limits::max())) { @@ -84,8 +73,8 @@ static std::optional checkedSectionFileOffset(const ELFT::Shdr &Sec, } uint64_t Delta = VAddr - Sec.sh_addr; - std::optional FileOffset = - checkedAdd(Sec.sh_offset, Delta, (Twine(Context) + " file offset").str()); + std::optional FileOffset = checkedAddUint64( + Sec.sh_offset, Delta, (Twine(Context) + " file offset").str()); if (!FileOffset) return std::nullopt; @@ -837,7 +826,7 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, return true; std::optional TextEnd = - checkedAdd(TextOffset, TextSize, "section header .text end"); + checkedAddUint64(TextOffset, TextSize, "section header .text end"); if (!TextEnd) return false; uint64_t Shoff; @@ -851,7 +840,7 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, if (Shoff >= *TextEnd) { std::optional NewShoff = - checkedAdd(Shoff, TrampTotal, "section header table offset"); + checkedAddUint64(Shoff, TrampTotal, "section header table offset"); if (!NewShoff) return false; uint64_t NewShoffValue = *NewShoff; @@ -863,7 +852,7 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, for (uint16_t I = 0; I < Shnum; ++I) { uint64_t ShTableDelta = static_cast(I) * Shentsize; std::optional ShPos = - checkedAdd(Shoff, ShTableDelta, "section header entry offset"); + checkedAddUint64(Shoff, ShTableDelta, "section header entry offset"); if (!ShPos) return false; if (*ShPos > ElfSize || sizeof(Shdr) > ElfSize - *ShPos) @@ -874,7 +863,7 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, if (ShOffset == TextOffset) { std::optional NewTextSize = - checkedAdd(TextSize, TrampTotal, ".text section size"); + checkedAddUint64(TextSize, TrampTotal, ".text section size"); if (!NewTextSize) return false; uint64_t NewTextSizeValue = *NewTextSize; @@ -882,7 +871,7 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, sizeof(NewTextSizeValue)); } else if (ShOffset > TextOffset) { std::optional NewOffset = - checkedAdd(ShOffset, TrampTotal, "post-.text section offset"); + checkedAddUint64(ShOffset, TrampTotal, "post-.text section offset"); if (!NewOffset) return false; uint64_t NewOffsetValue = *NewOffset; @@ -894,7 +883,7 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, uint64_t ShAddr; std::memcpy(&ShAddr, Sh + offsetof(Shdr, sh_addr), sizeof(ShAddr)); std::optional NewAddr = - checkedAdd(ShAddr, TrampTotal, "post-.text section address"); + checkedAddUint64(ShAddr, TrampTotal, "post-.text section address"); if (!NewAddr) return false; ShAddr = *NewAddr; @@ -912,7 +901,7 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, return true; std::optional TextEnd = - checkedAdd(TextOffset, TextSize, "program header .text end"); + checkedAddUint64(TextOffset, TextSize, "program header .text end"); if (!TextEnd) return false; uint64_t Phoff; @@ -927,7 +916,7 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, for (uint16_t I = 0; I < Phnum; ++I) { uint64_t PhTableDelta = static_cast(I) * Phentsize; std::optional PhPos = - checkedAdd(Phoff, PhTableDelta, "program header entry offset"); + checkedAddUint64(Phoff, PhTableDelta, "program header entry offset"); if (!PhPos) return false; if (*PhPos > ElfSize || sizeof(Phdr) > ElfSize - *PhPos) @@ -941,14 +930,14 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, std::memcpy(&PMemsz, Ph + offsetof(Phdr, p_memsz), sizeof(PMemsz)); std::optional PEnd = - checkedAdd(POffset, PFilesz, "program header file end"); + checkedAddUint64(POffset, PFilesz, "program header file end"); if (!PEnd) return false; if (POffset <= TextOffset && *PEnd >= *TextEnd) { std::optional NewPFilesz = - checkedAdd(PFilesz, TrampTotal, "program header file size"); + checkedAddUint64(PFilesz, TrampTotal, "program header file size"); std::optional NewPMemsz = - checkedAdd(PMemsz, TrampTotal, "program header memory size"); + checkedAddUint64(PMemsz, TrampTotal, "program header memory size"); if (!NewPFilesz || !NewPMemsz) return false; PFilesz = *NewPFilesz; @@ -957,7 +946,7 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, std::memcpy(Ph + offsetof(Phdr, p_memsz), &PMemsz, sizeof(PMemsz)); } else if (POffset > TextOffset) { std::optional NewPOffset = - checkedAdd(POffset, TrampTotal, "post-.text program offset"); + checkedAddUint64(POffset, TrampTotal, "post-.text program offset"); if (!NewPOffset) return false; POffset = *NewPOffset; @@ -965,7 +954,7 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, uint64_t PVaddr; std::memcpy(&PVaddr, Ph + offsetof(Phdr, p_vaddr), sizeof(PVaddr)); std::optional NewPVaddr = - checkedAdd(PVaddr, TrampTotal, "post-.text program vaddr"); + checkedAddUint64(PVaddr, TrampTotal, "post-.text program vaddr"); if (!NewPVaddr) return false; PVaddr = *NewPVaddr; @@ -973,7 +962,7 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, uint64_t PPaddr; std::memcpy(&PPaddr, Ph + offsetof(Phdr, p_paddr), sizeof(PPaddr)); std::optional NewPPaddr = - checkedAdd(PPaddr, TrampTotal, "post-.text program paddr"); + checkedAddUint64(PPaddr, TrampTotal, "post-.text program paddr"); if (!NewPPaddr) return false; PPaddr = *NewPPaddr; @@ -1050,7 +1039,7 @@ static bool adjustSymbolValues(uint8_t *Elf, size_t ElfSize, uint64_t SymOffset = SymBytes - File.base(); std::optional Value = - checkedAdd(Sym.st_value, TrampTotal, "post-.text symbol value"); + checkedAddUint64(Sym.st_value, TrampTotal, "post-.text symbol value"); if (!Value) return false; uint64_t Value64 = *Value; @@ -1091,8 +1080,8 @@ ElfView::growWithTrampolines(ArrayRef Trampolines, return nullptr; } - std::optional TextEnd = - checkedAdd(textOffset(), textSize(), "growWithTrampolines .text end"); + std::optional TextEnd = checkedAddUint64( + textOffset(), textSize(), "growWithTrampolines .text end"); if (!TextEnd || *TextEnd > InputSize) { log() << "hotswap: error: growWithTrampolines: .text range exceeds input " << "ELF size.\n"; diff --git a/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp b/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp index b8497fea7665e..86286b3d9db3a 100644 --- a/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp +++ b/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp @@ -17,7 +17,6 @@ #include "llvm/ADT/StringExtras.h" #include "llvm/ADT/Twine.h" -#include "llvm/Support/CheckedArithmetic.h" #include #include @@ -47,16 +46,6 @@ static SmallVector getCodeEndBytes(const LLVMState &LS) { return CodeEnd; } -static std::optional checkedAdd(uint64_t LHS, uint64_t RHS, - StringRef Context) { - std::optional Result = checkedAddUnsigned(LHS, RHS); - if (Result) - return Result; - - log() << "hotswap: error: " << Context << " overflows uint64_t.\n"; - return std::nullopt; -} - SmallVector buildKernelEntryTrampoline(uint64_t StubVAddr, uint64_t EntryVAddr, unsigned ScratchSgpr, @@ -87,8 +76,8 @@ SmallVector buildKernelEntryTrampoline(uint64_t StubVAddr, // Materialize the original entry with a 64-bit PC-relative add so the code // object can be rewritten before ROCR knows final device addresses. std::optional PcBase = - checkedAdd(StubVAddr, static_cast(Bytes.size()), - "kernel-entry stub PC base"); + checkedAddUint64(StubVAddr, static_cast(Bytes.size()), + "kernel-entry stub PC base"); if (!PcBase) return {}; // Unsigned subtraction is intentional: the immediate pair materializes the @@ -235,12 +224,12 @@ static std::optional decodeEntryStubTargetVAddr(ArrayRef Decoded, uint64_t StubVAddr) { std::optional PcBaseOffset = - checkedAdd(Decoded[2].Offset, Decoded[2].Size, - "decoded kernel-entry stub PC-base offset"); + checkedAddUint64(Decoded[2].Offset, Decoded[2].Size, + "decoded kernel-entry stub PC-base offset"); if (!PcBaseOffset) return std::nullopt; - std::optional PcBase = - checkedAdd(StubVAddr, *PcBaseOffset, "decoded kernel-entry stub PC base"); + std::optional PcBase = checkedAddUint64( + StubVAddr, *PcBaseOffset, "decoded kernel-entry stub PC base"); if (!PcBase) return std::nullopt; @@ -266,12 +255,12 @@ checkedAlignTo(uint64_t Value, uint64_t Alignment, StringRef Context) { uint64_t Remainder = Value % Alignment; if (Remainder == 0) return Value; - return checkedAdd(Value, Alignment - Remainder, Context); + return checkedAddUint64(Value, Alignment - Remainder, Context); } static std::optional entryVAddr(const KernelDescriptorInfo &KD) { if (KD.EntryOffset >= 0) - return checkedAdd( + return checkedAddUint64( KD.VAddr, static_cast(KD.EntryOffset), (Twine("kernel entry vaddr for '") + KD.KernelName + "'").str()); @@ -287,18 +276,16 @@ static std::optional entryVAddr(const KernelDescriptorInfo &KD) { return KD.VAddr - Magnitude; } -static std::optional -descriptorAlreadyTargetsEntryStub(const ElfView &Elf, - const KernelDescriptorInfo &KD, - const LLVMState &LS) { +static std::optional descriptorAlreadyTargetsEntryStub( + const ElfView &Elf, const KernelDescriptorInfo &KD, const LLVMState &LS) { std::optional Entry = entryVAddr(KD); if (!Entry) return std::nullopt; if (*Entry < Elf.textAddr()) return false; - std::optional TextEnd = - checkedAdd(Elf.textAddr(), Elf.textSize(), "entry trampoline text end"); + std::optional TextEnd = checkedAddUint64( + Elf.textAddr(), Elf.textSize(), "entry trampoline text end"); if (!TextEnd) return std::nullopt; @@ -327,8 +314,8 @@ totalTrampolineBytes(ArrayRef Trampolines) { uint64_t Total = 0; for (const Trampoline &T : Trampolines) { std::optional NewTotal = - checkedAdd(Total, static_cast(T.Bytes.size()), - "existing trampoline byte count"); + checkedAddUint64(Total, static_cast(T.Bytes.size()), + "existing trampoline byte count"); if (!NewTotal) return std::nullopt; Total = *NewTotal; @@ -445,12 +432,12 @@ std::optional appendKernelEntryTrampolines( if (!ExistingGrowthBytes) return std::nullopt; uint64_t AppendOffset = *ExistingGrowthBytes; - std::optional TextEndVAddr = - checkedAdd(Elf.textAddr(), Elf.textSize(), "entry trampoline text end"); + std::optional TextEndVAddr = checkedAddUint64( + Elf.textAddr(), Elf.textSize(), "entry trampoline text end"); if (!TextEndVAddr) return std::nullopt; - std::optional StubPoolBaseVAddr = - checkedAdd(*TextEndVAddr, AppendOffset, "entry trampoline stub-pool base"); + std::optional StubPoolBaseVAddr = checkedAddUint64( + *TextEndVAddr, AppendOffset, "entry trampoline stub-pool base"); if (!StubPoolBaseVAddr) return std::nullopt; std::optional AlignedStubPoolBaseVAddr = @@ -467,13 +454,13 @@ std::optional appendKernelEntryTrampolines( AppendOffset = StubStart; for (const KernelDescriptorInfo &KD : Work) { - std::optional StubTextEnd = checkedAdd( + std::optional StubTextEnd = checkedAddUint64( Elf.textSize(), AppendOffset, (Twine("entry trampoline append offset for '") + KD.KernelName + "'") .str()); if (!StubTextEnd) return std::nullopt; - std::optional StubVAddr = checkedAdd( + std::optional StubVAddr = checkedAddUint64( Elf.textAddr(), *StubTextEnd, (Twine("entry trampoline vaddr for '") + KD.KernelName + "'").str()); if (!StubVAddr) @@ -498,7 +485,7 @@ std::optional appendKernelEntryTrampolines( T.Bytes.assign(Stub.begin(), Stub.end()); LocalGrowth.push_back(std::move(T)); LocalFixups.push_back({KD.KernelName, AppendOffset, *ScratchSgpr + 2}); - std::optional NewAppendOffset = checkedAdd( + std::optional NewAppendOffset = checkedAddUint64( AppendOffset, KernelEntryStubStride, (Twine("entry trampoline append offset after '") + KD.KernelName + "'") .str()); @@ -560,7 +547,7 @@ bool rewriteKernelEntryDescriptorOffsets( Ok = false; continue; } - std::optional StubTextOffset = checkedAdd( + std::optional StubTextOffset = checkedAddUint64( OldTextSize, Fixup.StubTextOffset, (Twine("entry trampoline text offset for '") + Fixup.KernelName + "'") .str()); @@ -568,7 +555,7 @@ bool rewriteKernelEntryDescriptorOffsets( Ok = false; continue; } - std::optional StubVAddr = checkedAdd( + std::optional StubVAddr = checkedAddUint64( OutElf.textAddr(), *StubTextOffset, (Twine("entry trampoline vaddr for '") + Fixup.KernelName + "'").str()); if (!StubVAddr) { diff --git a/amd/comgr/src/comgr-hotswap-internal.h b/amd/comgr/src/comgr-hotswap-internal.h index fd00d4e08baa2..9b59fa74e9e44 100644 --- a/amd/comgr/src/comgr-hotswap-internal.h +++ b/amd/comgr/src/comgr-hotswap-internal.h @@ -52,6 +52,7 @@ #include "llvm/Object/ELF.h" #include "llvm/Object/ELFTypes.h" #include "llvm/Support/AMDHSAKernelDescriptor.h" +#include "llvm/Support/CheckedArithmetic.h" #include "llvm/Support/Error.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/raw_ostream.h" @@ -73,6 +74,16 @@ inline llvm::raw_ostream &log() { return COMGR::env::shouldEmitVerboseLogs() ? llvm::errs() : llvm::nulls(); } +inline std::optional checkedAddUint64(uint64_t LHS, uint64_t RHS, + llvm::StringRef Context) { + std::optional Result = llvm::checkedAddUnsigned(LHS, RHS); + if (Result) + return Result; + + log() << "hotswap: error: " << Context << " overflows uint64_t.\n"; + return std::nullopt; +} + // -- Trampoline and NOP sled -------------------------------------------------- struct Trampoline { @@ -537,7 +548,7 @@ struct VgprAllocator { }; /// Allocates scratch SGPRs for a patch point. Unlike VGPRs (which have full -/// dataflow liveness), SGPRs have no liveness analysis — we always allocate +/// dataflow liveness), SGPRs have no liveness analysis, so we always allocate /// above the kernel descriptor's reported SGPR count. This is conservative /// but safe: no SGPR currently in use by the kernel can be clobbered. struct SgprAllocator {