From dd686a29afdb05793e0d8c8bd76ef59437b4c23a Mon Sep 17 00:00:00 2001 From: Harsh Menon Date: Wed, 1 Jul 2026 05:21:38 -0700 Subject: [PATCH] AMDGPU: harden hotswap ELF growth and descriptors Key entry trampoline descriptor rewrites by descriptor virtual address so duplicate kernel descriptor names update the intended descriptor. Make SGPR metadata fallback explicit, keep descriptor SGPR fields in sync with the effective metadata count, and adjust grown ELF addresses using virtual address ordering. Add unit coverage for duplicate descriptors, metadata fallback, vaddr-based ELF growth, and entry-only debug growth. --- amd/comgr/src/comgr-hotswap-b0a0.cpp | 24 +- amd/comgr/src/comgr-hotswap-elf.cpp | 417 ++++++++++++------ .../src/comgr-hotswap-entry-trampoline.cpp | 48 +- amd/comgr/src/comgr-hotswap-internal.h | 28 +- amd/comgr/test-unit/HotswapElfTest.cpp | 117 ++++- amd/comgr/test-unit/HotswapMCTest.cpp | 158 +++++++ amd/comgr/test-unit/comgr-test-elf-utils.h | 46 +- 7 files changed, 660 insertions(+), 178 deletions(-) diff --git a/amd/comgr/src/comgr-hotswap-b0a0.cpp b/amd/comgr/src/comgr-hotswap-b0a0.cpp index 65f1afde36d01..167fe4457c2ed 100644 --- a/amd/comgr/src/comgr-hotswap-b0a0.cpp +++ b/amd/comgr/src/comgr-hotswap-b0a0.cpp @@ -475,23 +475,25 @@ fixupTrampolineBranches(std::vector &Trampolines, uint8_t *Text, } /// Fix up DWARF sections of the grown ELF after trampolines have been -/// appended: adds trampoline symbols to the symbol table, shifts -/// .debug_line / .debug_ranges / .debug_info / .debug_frame addresses by -/// the total trampoline footprint, and reports per-section failures via -/// log(). Individual patchDebug* helpers are weak stubs here; concrete +/// appended. \p SymbolTrampolines contains only B0/A0 trampolines that +/// correspond to rewritten source instructions and therefore may get +/// trampoline symbols / line entries. \p TextGrowthBytes is the full .text +/// growth, including entry stubs and padding, used to shift post-.text debug +/// addresses. Individual patchDebug* helpers are weak stubs here; concrete /// implementations land in separate PRs. static void patchDebugSections(WritableMemoryBuffer &ElfBuf, - ArrayRef Trampolines, - const ElfView &Elf, size_t GrowthTotal) { + ArrayRef SymbolTrampolines, + const ElfView &Elf, size_t TextGrowthBytes) { uint8_t *Data = reinterpret_cast(ElfBuf.getBufferStart()); size_t Size = ElfBuf.getBufferSize(); - if (!addTrampolineSymbols(ElfBuf, Trampolines, Elf.textSize(), + if (!addTrampolineSymbols(ElfBuf, SymbolTrampolines, Elf.textSize(), Elf.textSectionIndex())) log() << "hotswap: error: addTrampolineSymbols failed\n"; - patchDebugRanges(Data, Size, Elf.textAddr(), Elf.textSize(), GrowthTotal); - patchDebugInfo(Data, Size, Elf.textAddr(), Elf.textSize(), GrowthTotal); - patchDebugFrame(Data, Size, Elf.textAddr(), Elf.textSize(), GrowthTotal); - if (!patchDebugLine(ElfBuf, Trampolines, Elf.textSize(), Elf.textAddr())) + patchDebugRanges(Data, Size, Elf.textAddr(), Elf.textSize(), TextGrowthBytes); + patchDebugInfo(Data, Size, Elf.textAddr(), Elf.textSize(), TextGrowthBytes); + patchDebugFrame(Data, Size, Elf.textAddr(), Elf.textSize(), TextGrowthBytes); + if (!patchDebugLine(ElfBuf, SymbolTrampolines, Elf.textSize(), + Elf.textAddr())) log() << "hotswap: error: patchDebugLine failed\n"; } diff --git a/amd/comgr/src/comgr-hotswap-elf.cpp b/amd/comgr/src/comgr-hotswap-elf.cpp index a80a7bb8264f6..e8052e1f5799a 100644 --- a/amd/comgr/src/comgr-hotswap-elf.cpp +++ b/amd/comgr/src/comgr-hotswap-elf.cpp @@ -93,7 +93,7 @@ readSgprCountMetadataNode(const msgpack::DocNode &SgprNode, if (SgprNode.getKind() == msgpack::Type::UInt) { uint64_t SgprCount = SgprNode.getUInt(); if (SgprCount > std::numeric_limits::max()) { - log() << "hotswap: error: " << Context << ": .sgpr_count for '" + log() << "hotswap: warning: " << Context << ": .sgpr_count for '" << KernelName << "' exceeds unsigned.\n"; return std::nullopt; } @@ -104,21 +104,145 @@ readSgprCountMetadataNode(const msgpack::DocNode &SgprNode, int64_t SgprCount = SgprNode.getInt(); if (SgprCount < 0 || static_cast(SgprCount) > std::numeric_limits::max()) { - log() << "hotswap: error: " << Context << ": .sgpr_count for '" + log() << "hotswap: warning: " << Context << ": .sgpr_count for '" << KernelName << "' is outside unsigned range.\n"; return std::nullopt; } return static_cast(SgprCount); } - log() << "hotswap: error: " << Context << ": .sgpr_count for '" << KernelName - << "' is not an integer.\n"; + log() << "hotswap: warning: " << Context << ": .sgpr_count for '" + << KernelName << "' is not an integer.\n"; return std::nullopt; } +static std::optional encodeSgprCountGranule(unsigned RequiredSgprs, + StringRef KernelName, + StringRef Context) { + if (RequiredSgprs == 0) + return 0; + + uint64_t RequiredGranulated64 = + (static_cast(RequiredSgprs) + SgprEncodingGranule - 1) / + SgprEncodingGranule - + 1; + uint32_t MaxGranulated = static_cast( + amdhsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT >> + amdhsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT_SHIFT); + if (RequiredGranulated64 > MaxGranulated) { + log() << "hotswap: error: " << Context << ": kernel '" << KernelName + << "' needs " << RequiredSgprs + << " SGPRs, which exceeds the descriptor encoding limit.\n"; + return std::nullopt; + } + + return static_cast(RequiredGranulated64); +} + +static std::optional readDescriptorSgprCount(const uint8_t *Kd, + StringRef KernelName, + StringRef Context) { + if (!Kd) { + log() << "hotswap: error: " << Context << ": kernel descriptor symbol '" + << KernelName << ".kd' not found for SGPR count fallback.\n"; + return std::nullopt; + } + + uint32_t Rsrc1; + std::memcpy(&Rsrc1, + Kd + offsetof(amdhsa::kernel_descriptor_t, compute_pgm_rsrc1), + sizeof(Rsrc1)); + uint32_t Granulated = AMDHSA_BITS_GET( + Rsrc1, amdhsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT); + uint64_t SgprCount = + (static_cast(Granulated) + 1) * SgprEncodingGranule; + if (SgprCount > std::numeric_limits::max()) { + log() << "hotswap: error: " << Context << ": descriptor SGPR count for '" + << KernelName << "' exceeds unsigned.\n"; + return std::nullopt; + } + return static_cast(SgprCount); +} + +static void logSgprMetadataFallback(StringRef Context, StringRef KernelName, + StringRef Reason) { + log() << "hotswap: warning: " << Context << ": " << Reason << " for kernel '" + << KernelName + << "'; falling back to the kernel descriptor SGPR count.\n"; +} + +static uint8_t * +findKernelDescriptorSymbol(const ELFFileT &File, ELFT::ShdrRange Sections, + uint8_t *Data, size_t Size, StringRef KernelName, + std::optional VAddr, StringRef Context) { + std::string KdName; + if (!KernelName.empty()) + KdName = (KernelName + ".kd").str(); + + for (const ELFT::Shdr &SymShdr : Sections) { + if (SymShdr.sh_type != ELF::SHT_SYMTAB && + SymShdr.sh_type != ELF::SHT_DYNSYM) + continue; + + Expected SymsOrErr = File.symbols(&SymShdr); + if (!SymsOrErr) { + log() << "hotswap: error: " << Context + << ": failed to read descriptor symbol table: " + << toString(SymsOrErr.takeError()) << "\n"; + continue; + } + Expected StrTabOrErr = + File.getStringTableForSymtab(SymShdr, Sections); + if (!StrTabOrErr) { + log() << "hotswap: error: " << Context + << ": failed to read descriptor symbol string table: " + << toString(StrTabOrErr.takeError()) << "\n"; + continue; + } + + for (const ELFT::Sym &Sym : *SymsOrErr) { + Expected NameOrErr = Sym.getName(*StrTabOrErr); + if (!NameOrErr) { + log() << "hotswap: error: " << Context + << ": failed to read descriptor symbol name: " + << toString(NameOrErr.takeError()) << "\n"; + continue; + } + + if (KernelName.empty()) { + if (!NameOrErr->ends_with(".kd")) + continue; + } else if (*NameOrErr != KdName) { + continue; + } + + if (VAddr && Sym.st_value != *VAddr) + continue; + + unsigned Shndx = Sym.st_shndx; + Expected HostShdrOrErr = File.getSection(Shndx); + if (!HostShdrOrErr) { + log() << "hotswap: error: " << Context << ": descriptor symbol '" + << *NameOrErr << "' has unreadable section index " << Shndx + << ": " << toString(HostShdrOrErr.takeError()) << "\n"; + continue; + } + const ELFT::Shdr &HostShdr = **HostShdrOrErr; + std::optional FileOffset = checkedSectionFileOffset( + HostShdr, Sym.st_value, KdSize, Size, + (Twine(Context) + ": descriptor symbol '" + *NameOrErr + "'").str()); + if (!FileOffset) + continue; + return Data + *FileOffset; + } + } + return nullptr; +} + static MetadataSgprUpdateStatus updateKernelMetadataSgprCount(uint8_t *Elf, const ELFFileT &File, - StringRef KernelName, unsigned RequiredSgprs) { + StringRef KernelName, unsigned RequiredSgprs, + unsigned &EffectiveSgprs) { Expected PhdrsOrErr = File.program_headers(); if (!PhdrsOrErr) { log() << "hotswap: error: updateKernelMetadataSgprCount: failed to read " @@ -179,19 +303,22 @@ updateKernelMetadataSgprCount(uint8_t *Elf, const ELFFileT &File, msgpack::DocNode::MapTy::iterator SgprIt = KMap.find(".sgpr_count"); if (SgprIt == KMap.end()) { - log() << "hotswap: error: updateKernelMetadataSgprCount: metadata " - << "for kernel '" << KernelName << "' has no .sgpr_count.\n"; - return MetadataSgprUpdateStatus::Error; + logSgprMetadataFallback("updateKernelMetadataSgprCount", KernelName, + "metadata has no .sgpr_count"); + return MetadataSgprUpdateStatus::NotFound; } std::optional CurrentSgprs = readSgprCountMetadataNode( SgprIt->second, KernelName, "updateKernelMetadataSgprCount"); if (!CurrentSgprs) - return MetadataSgprUpdateStatus::Error; - if (RequiredSgprs <= *CurrentSgprs) + return MetadataSgprUpdateStatus::NotFound; + if (RequiredSgprs <= *CurrentSgprs) { + EffectiveSgprs = *CurrentSgprs; return MetadataSgprUpdateStatus::Found; + } SgprIt->second = static_cast(RequiredSgprs); + EffectiveSgprs = RequiredSgprs; std::string NewBlob; Doc.writeToBlob(NewBlob); if (NewBlob.size() != Blob.size()) { @@ -229,9 +356,9 @@ updateKernelMetadataSgprCount(uint8_t *Elf, const ELFFileT &File, } if (SawMetadataNote) { - log() << "hotswap: error: updateKernelMetadataSgprCount: AMDGPU metadata " - << "has no entry for kernel '" << KernelName << "'.\n"; - return MetadataSgprUpdateStatus::Error; + logSgprMetadataFallback("updateKernelMetadataSgprCount", KernelName, + "AMDGPU metadata has no matching entry"); + return MetadataSgprUpdateStatus::NotFound; } return MetadataSgprUpdateStatus::NotFound; } @@ -379,50 +506,14 @@ std::string ElfView::findKernelAtOffset(uint64_t TextOffset) const { // -- ElfView::findKernelDescriptor -------------------------------------------- uint8_t *ElfView::findKernelDescriptor(StringRef KernelName) { - std::string KdName = (KernelName + ".kd").str(); - for (const ELFT::Shdr &SymShdr : Sections) { - if (SymShdr.sh_type != ELF::SHT_SYMTAB && - SymShdr.sh_type != ELF::SHT_DYNSYM) - continue; - - Expected SymsOrErr = File.symbols(&SymShdr); - if (!SymsOrErr) { - consumeError(SymsOrErr.takeError()); - continue; - } - Expected StrTabOrErr = - File.getStringTableForSymtab(SymShdr, Sections); - if (!StrTabOrErr) { - consumeError(StrTabOrErr.takeError()); - continue; - } + return findKernelDescriptorSymbol(File, Sections, data(), size(), KernelName, + std::nullopt, "findKernelDescriptor"); +} - for (const ELFT::Sym &Sym : *SymsOrErr) { - Expected NameOrErr = Sym.getName(*StrTabOrErr); - if (!NameOrErr) { - consumeError(NameOrErr.takeError()); - continue; - } - if (*NameOrErr != KdName) - continue; - unsigned Shndx = Sym.st_shndx; - Expected HostShdrOrErr = File.getSection(Shndx); - if (!HostShdrOrErr) { - consumeError(HostShdrOrErr.takeError()); - continue; - } - const ELFT::Shdr &HostShdr = **HostShdrOrErr; - std::optional FileOffset = checkedSectionFileOffset( - HostShdr, Sym.st_value, KdSize, size(), - (Twine("findKernelDescriptor: descriptor symbol '") + *NameOrErr + - "'") - .str()); - if (!FileOffset) - continue; - return data() + *FileOffset; - } - } - return nullptr; +uint8_t *ElfView::findKernelDescriptorByVAddr(uint64_t DescriptorVAddr) { + return findKernelDescriptorSymbol(File, Sections, data(), size(), "", + DescriptorVAddr, + "findKernelDescriptorByVAddr"); } // -- ElfView::kernelDescriptors ----------------------------------------------- @@ -507,13 +598,32 @@ ElfView::getKernelDescriptorVAddr(StringRef KernelName) const { bool ElfView::updateKernelDescriptorEntryOffset(StringRef KernelName, int64_t NewEntryOffset) { - namespace hsa = amdhsa; uint8_t *Kd = findKernelDescriptor(KernelName); if (!Kd) { log() << "hotswap: error: updateKernelDescriptorEntryOffset: kernel " << "descriptor symbol '" << KernelName << ".kd' not found.\n"; return false; } + namespace hsa = amdhsa; + std::memcpy( + Kd + offsetof(hsa::kernel_descriptor_t, kernel_code_entry_byte_offset), + &NewEntryOffset, sizeof(NewEntryOffset)); + return true; +} + +bool ElfView::updateKernelDescriptorEntryOffset(uint64_t DescriptorVAddr, + StringRef KernelName, + int64_t NewEntryOffset) { + namespace hsa = amdhsa; + uint8_t *Kd = findKernelDescriptorSymbol(File, Sections, data(), size(), + KernelName, DescriptorVAddr, + "updateKernelDescriptorEntryOffset"); + if (!Kd) { + log() << "hotswap: error: updateKernelDescriptorEntryOffset: kernel " + << "descriptor symbol '" << KernelName << ".kd' at vaddr 0x" + << utohexstr(DescriptorVAddr) << " not found.\n"; + return false; + } std::memcpy( Kd + offsetof(hsa::kernel_descriptor_t, kernel_code_entry_byte_offset), &NewEntryOffset, sizeof(NewEntryOffset)); @@ -522,55 +632,56 @@ bool ElfView::updateKernelDescriptorEntryOffset(StringRef KernelName, bool ElfView::updateKernelDescriptorSgprCount(StringRef KernelName, unsigned RequiredSgprs) { + std::optional DescriptorVAddr = + getKernelDescriptorVAddr(KernelName); + if (!DescriptorVAddr) { + log() << "hotswap: error: updateKernelDescriptorSgprCount: kernel " + << "descriptor symbol '" << KernelName << ".kd' not found.\n"; + return false; + } + + return updateKernelDescriptorSgprCount(*DescriptorVAddr, KernelName, + RequiredSgprs); +} + +bool ElfView::updateKernelDescriptorSgprCount(uint64_t DescriptorVAddr, + StringRef KernelName, + unsigned RequiredSgprs) { namespace hsa = amdhsa; if (RequiredSgprs == 0) return true; - uint8_t *Kd = findKernelDescriptor(KernelName); + if (!encodeSgprCountGranule(RequiredSgprs, KernelName, + "updateKernelDescriptorSgprCount")) + return false; + + uint8_t *Kd = findKernelDescriptorSymbol(File, Sections, data(), size(), + KernelName, DescriptorVAddr, + "updateKernelDescriptorSgprCount"); if (!Kd) { log() << "hotswap: error: updateKernelDescriptorSgprCount: kernel " - << "descriptor symbol '" << KernelName << ".kd' not found.\n"; + << "descriptor symbol '" << KernelName << ".kd' at vaddr 0x" + << utohexstr(DescriptorVAddr) << " not found.\n"; return false; } - uint32_t Rsrc1 = 0; - std::memcpy(&Rsrc1, - Kd + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc1), - sizeof(Rsrc1)); - - uint32_t CurrentGranulated = AMDHSA_BITS_GET( - Rsrc1, hsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT); - uint64_t CurrentSgprs = - (static_cast(CurrentGranulated) + 1) * SgprEncodingGranule; - - std::optional RequiredGranulated; - if (RequiredSgprs > CurrentSgprs) { - uint64_t RequiredGranulated64 = - (static_cast(RequiredSgprs) + SgprEncodingGranule - 1) / - SgprEncodingGranule - - 1; - uint32_t MaxGranulated = static_cast( - hsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT >> - hsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT_SHIFT); - if (RequiredGranulated64 > MaxGranulated) { - log() << "hotswap: error: updateKernelDescriptorSgprCount: kernel '" - << KernelName << "' needs " << RequiredSgprs - << " SGPRs, which exceeds the descriptor encoding limit.\n"; - return false; - } - RequiredGranulated = static_cast(RequiredGranulated64); - } - - MetadataSgprUpdateStatus MetadataStatus = - updateKernelMetadataSgprCount(data(), File, KernelName, RequiredSgprs); + unsigned EffectiveSgprs = RequiredSgprs; + MetadataSgprUpdateStatus MetadataStatus = updateKernelMetadataSgprCount( + data(), File, KernelName, RequiredSgprs, EffectiveSgprs); if (MetadataStatus == MetadataSgprUpdateStatus::Error) return false; // NotFound is allowed for minimal code objects without AMDGPU metadata; in // that case the descriptor field remains the only SGPR count to update. + std::optional RequiredGranulated = encodeSgprCountGranule( + EffectiveSgprs, KernelName, "updateKernelDescriptorSgprCount"); if (!RequiredGranulated) - return true; + return false; + uint32_t Rsrc1 = 0; + std::memcpy(&Rsrc1, + Kd + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc1), + sizeof(Rsrc1)); AMDHSA_BITS_SET(Rsrc1, hsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT, *RequiredGranulated); std::memcpy(Kd + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc1), @@ -732,12 +843,20 @@ ElfView::getKernelSgprCount(StringRef KernelName) const { msgpack::DocNode::MapTy::iterator SgprIt = KMap.find(".sgpr_count"); if (SgprIt == KMap.end()) { - log() << "hotswap: error: getKernelSgprCount: metadata for kernel '" - << KernelName << "' has no .sgpr_count.\n"; - return std::nullopt; - } - return readSgprCountMetadataNode(SgprIt->second, KernelName, + logSgprMetadataFallback("getKernelSgprCount", KernelName, + "metadata has no .sgpr_count"); + uint8_t *Kd = + const_cast(this)->findKernelDescriptor(KernelName); + return readDescriptorSgprCount(Kd, KernelName, "getKernelSgprCount"); + } + std::optional MetadataSgprs = readSgprCountMetadataNode( + SgprIt->second, KernelName, "getKernelSgprCount"); + if (MetadataSgprs) + return MetadataSgprs; + uint8_t *Kd = + const_cast(this)->findKernelDescriptor(KernelName); + return readDescriptorSgprCount(Kd, KernelName, "getKernelSgprCount"); } } if (Err) { @@ -753,33 +872,18 @@ ElfView::getKernelSgprCount(StringRef KernelName) const { } if (SawMetadataNote) { - log() << "hotswap: error: getKernelSgprCount: AMDGPU metadata has no " - << ".sgpr_count entry for kernel '" << KernelName << "'.\n"; - return std::nullopt; + logSgprMetadataFallback("getKernelSgprCount", KernelName, + "AMDGPU metadata has no matching .sgpr_count"); + uint8_t *Kd = const_cast(this)->findKernelDescriptor(KernelName); + return readDescriptorSgprCount(Kd, KernelName, "getKernelSgprCount"); } // --- Fallback: read the KD field. --- // The LLVM assembler populates GRANULATED_WAVEFRONT_SGPR_COUNT even on // GFX10+ where the hardware ignores it, so this is still usable for // ROCm-compiled code objects that lack a metadata note. - namespace hsa = amdhsa; uint8_t *Kd = const_cast(this)->findKernelDescriptor(KernelName); - if (!Kd) - return std::nullopt; - uint32_t Rsrc1; - std::memcpy(&Rsrc1, - Kd + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc1), - sizeof(Rsrc1)); - uint32_t Granulated = AMDHSA_BITS_GET( - Rsrc1, hsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT); - uint64_t SgprCount = - (static_cast(Granulated) + 1) * SgprEncodingGranule; - if (SgprCount > std::numeric_limits::max()) { - log() << "hotswap: error: getKernelSgprCount: descriptor SGPR count for '" - << KernelName << "' exceeds unsigned.\n"; - return std::nullopt; - } - return static_cast(SgprCount); + return readDescriptorSgprCount(Kd, KernelName, "getKernelSgprCount"); } // -- ElfView::updateKernelDescriptor ------------------------------------------ @@ -820,8 +924,8 @@ void ElfView::updateKernelDescriptor(StringRef KernelName, unsigned ExtraVgprs, // -- Section/program header adjustment for trampoline growth ------------------ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, - uint64_t TextOffset, uint64_t TextSize, - size_t TrampTotal) { + uint64_t TextOffset, uint64_t TextAddr, + uint64_t TextSize, size_t TrampTotal) { if (ElfSize < sizeof(Ehdr)) return true; @@ -829,6 +933,10 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, checkedAddUint64(TextOffset, TextSize, "section header .text end"); if (!TextEnd) return false; + std::optional TextEndAddr = + checkedAddUint64(TextAddr, TextSize, "section header .text vaddr end"); + if (!TextEndAddr) + return false; uint64_t Shoff; uint16_t Shentsize; uint16_t Shnum; @@ -860,6 +968,8 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, uint8_t *Sh = Elf + *ShPos; uint64_t ShOffset; std::memcpy(&ShOffset, Sh + offsetof(Shdr, sh_offset), sizeof(ShOffset)); + uint64_t ShFlags; + std::memcpy(&ShFlags, Sh + offsetof(Shdr, sh_flags), sizeof(ShFlags)); if (ShOffset == TextOffset) { std::optional NewTextSize = @@ -877,11 +987,12 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, uint64_t NewOffsetValue = *NewOffset; std::memcpy(Sh + offsetof(Shdr, sh_offset), &NewOffsetValue, sizeof(NewOffsetValue)); - uint64_t ShFlags; - std::memcpy(&ShFlags, Sh + offsetof(Shdr, sh_flags), sizeof(ShFlags)); - if (ShFlags & ELF::SHF_ALLOC) { - uint64_t ShAddr; - std::memcpy(&ShAddr, Sh + offsetof(Shdr, sh_addr), sizeof(ShAddr)); + } + + if (ShOffset != TextOffset && (ShFlags & ELF::SHF_ALLOC)) { + uint64_t ShAddr; + std::memcpy(&ShAddr, Sh + offsetof(Shdr, sh_addr), sizeof(ShAddr)); + if (ShAddr >= *TextEndAddr) { std::optional NewAddr = checkedAddUint64(ShAddr, TrampTotal, "post-.text section address"); if (!NewAddr) @@ -895,8 +1006,8 @@ static bool adjustSectionHeaders(uint8_t *Elf, size_t ElfSize, } static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, - uint64_t TextOffset, uint64_t TextSize, - size_t TrampTotal) { + uint64_t TextOffset, uint64_t TextAddr, + uint64_t TextSize, size_t TrampTotal) { if (ElfSize < sizeof(Ehdr)) return true; @@ -904,6 +1015,10 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, checkedAddUint64(TextOffset, TextSize, "program header .text end"); if (!TextEnd) return false; + std::optional TextEndAddr = + checkedAddUint64(TextAddr, TextSize, "program header .text vaddr end"); + if (!TextEndAddr) + return false; uint64_t Phoff; uint16_t Phentsize; uint16_t Phnum; @@ -925,42 +1040,60 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, uint64_t POffset; uint64_t PFilesz; uint64_t PMemsz; + uint64_t PVaddr; + uint64_t PPaddr; std::memcpy(&POffset, Ph + offsetof(Phdr, p_offset), sizeof(POffset)); std::memcpy(&PFilesz, Ph + offsetof(Phdr, p_filesz), sizeof(PFilesz)); std::memcpy(&PMemsz, Ph + offsetof(Phdr, p_memsz), sizeof(PMemsz)); + std::memcpy(&PVaddr, Ph + offsetof(Phdr, p_vaddr), sizeof(PVaddr)); + std::memcpy(&PPaddr, Ph + offsetof(Phdr, p_paddr), sizeof(PPaddr)); std::optional PEnd = checkedAddUint64(POffset, PFilesz, "program header file end"); if (!PEnd) return false; - if (POffset <= TextOffset && *PEnd >= *TextEnd) { + std::optional PVaddrEnd = + checkedAddUint64(PVaddr, PMemsz, "program header vaddr end"); + if (!PVaddrEnd) + return false; + + bool FileCoversText = POffset <= TextOffset && *PEnd >= *TextEnd; + bool VAddrCoversText = PVaddr <= TextAddr && *PVaddrEnd >= *TextEndAddr; + if (FileCoversText) { std::optional NewPFilesz = checkedAddUint64(PFilesz, TrampTotal, "program header file size"); + if (!NewPFilesz) + return false; + PFilesz = *NewPFilesz; + std::memcpy(Ph + offsetof(Phdr, p_filesz), &PFilesz, sizeof(PFilesz)); + } + if (VAddrCoversText) { std::optional NewPMemsz = checkedAddUint64(PMemsz, TrampTotal, "program header memory size"); - if (!NewPFilesz || !NewPMemsz) + if (!NewPMemsz) return false; - PFilesz = *NewPFilesz; PMemsz = *NewPMemsz; - std::memcpy(Ph + offsetof(Phdr, p_filesz), &PFilesz, sizeof(PFilesz)); std::memcpy(Ph + offsetof(Phdr, p_memsz), &PMemsz, sizeof(PMemsz)); - } else if (POffset > TextOffset) { + } + + if (POffset > TextOffset) { std::optional NewPOffset = checkedAddUint64(POffset, TrampTotal, "post-.text program offset"); if (!NewPOffset) return false; POffset = *NewPOffset; std::memcpy(Ph + offsetof(Phdr, p_offset), &POffset, sizeof(POffset)); - uint64_t PVaddr; - std::memcpy(&PVaddr, Ph + offsetof(Phdr, p_vaddr), sizeof(PVaddr)); + } + + if (PVaddr >= *TextEndAddr) { std::optional NewPVaddr = checkedAddUint64(PVaddr, TrampTotal, "post-.text program vaddr"); if (!NewPVaddr) return false; PVaddr = *NewPVaddr; std::memcpy(Ph + offsetof(Phdr, p_vaddr), &PVaddr, sizeof(PVaddr)); - uint64_t PPaddr; - std::memcpy(&PPaddr, Ph + offsetof(Phdr, p_paddr), sizeof(PPaddr)); + } + if (PPaddr >= *TextEndAddr) { std::optional NewPPaddr = checkedAddUint64(PPaddr, TrampTotal, "post-.text program paddr"); if (!NewPPaddr) @@ -972,11 +1105,16 @@ static bool adjustProgramHeaders(uint8_t *Elf, size_t ElfSize, return true; } -static bool adjustSymbolValues(uint8_t *Elf, size_t ElfSize, - uint64_t TextOffset, size_t TrampTotal) { +static bool adjustSymbolValues(uint8_t *Elf, size_t ElfSize, uint64_t TextAddr, + uint64_t TextSize, size_t TrampTotal) { if (TrampTotal == 0) return true; + std::optional TextEndAddr = + checkedAddUint64(TextAddr, TextSize, "symbol .text vaddr end"); + if (!TextEndAddr) + return false; + Expected FileOrErr = ELFFileT::create(StringRef(reinterpret_cast(Elf), ElfSize)); if (!FileOrErr) { @@ -1027,7 +1165,7 @@ static bool adjustSymbolValues(uint8_t *Elf, size_t ElfSize, } const ELFT::Shdr &DefShdr = **DefShdrOrErr; if (!(DefShdr.sh_flags & ELF::SHF_ALLOC) || - DefShdr.sh_offset <= TextOffset) + DefShdr.sh_addr < *TextEndAddr) continue; const uint8_t *SymBytes = reinterpret_cast(&Sym); @@ -1140,11 +1278,12 @@ ElfView::growWithTrampolines(ArrayRef Trampolines, if (TextEndSize < InputSize) std::memcpy(Out + Pos, Input + TextEndSize, InputSize - TextEndSize); - if (!adjustSectionHeaders(Out, NewSize, textOffset(), textSize(), + if (!adjustSectionHeaders(Out, NewSize, textOffset(), textAddr(), textSize(), PaddedTrampTotal) || - !adjustProgramHeaders(Out, NewSize, textOffset(), textSize(), + !adjustProgramHeaders(Out, NewSize, textOffset(), textAddr(), textSize(), PaddedTrampTotal) || - !adjustSymbolValues(Out, NewSize, textOffset(), PaddedTrampTotal)) + !adjustSymbolValues(Out, NewSize, textAddr(), textSize(), + PaddedTrampTotal)) return nullptr; log() << "hotswap: growWithTrampolines: grew ELF from " << InputSize << " to " << NewSize << " bytes (" << Trampolines.size() << " trampoline" diff --git a/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp b/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp index 86286b3d9db3a..f09bf381483b3 100644 --- a/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp +++ b/amd/comgr/src/comgr-hotswap-entry-trampoline.cpp @@ -484,7 +484,12 @@ std::optional appendKernelEntryTrampolines( Trampoline T; T.Bytes.assign(Stub.begin(), Stub.end()); LocalGrowth.push_back(std::move(T)); - LocalFixups.push_back({KD.KernelName, AppendOffset, *ScratchSgpr + 2}); + KernelEntryTrampolineFixup Fixup; + Fixup.KernelName = KD.KernelName; + Fixup.KernelDescriptorVAddr = KD.VAddr; + Fixup.StubTextOffset = AppendOffset; + Fixup.RequiredSgprs = *ScratchSgpr + 2; + LocalFixups.push_back(std::move(Fixup)); std::optional NewAppendOffset = checkedAddUint64( AppendOffset, KernelEntryStubStride, (Twine("entry trampoline append offset after '") + KD.KernelName + "'") @@ -538,12 +543,37 @@ bool rewriteKernelEntryDescriptorOffsets( bool Ok = true; ElfView &OutElf = *ViewOrErr; + if (OutElf.textSize() < OldTextSize) { + log() << "hotswap: error: grown ELF .text size " << OutElf.textSize() + << " is smaller than original .text size " << OldTextSize << ".\n"; + return false; + } + uint64_t GrowthBytes = OutElf.textSize() - OldTextSize; + std::optional OldTextEndVAddr = + checkedAddUint64(OutElf.textAddr(), OldTextSize, + "entry trampoline original text end vaddr"); + if (!OldTextEndVAddr) + return false; + for (const KernelEntryTrampolineFixup &Fixup : Fixups) { - std::optional KdVAddr = - OutElf.getKernelDescriptorVAddr(Fixup.KernelName); - if (!KdVAddr) { + uint64_t DescriptorVAddr = Fixup.KernelDescriptorVAddr; + if (DescriptorVAddr >= *OldTextEndVAddr) { + std::optional ShiftedDescriptorVAddr = + checkedAddUint64(DescriptorVAddr, GrowthBytes, + (Twine("entry trampoline descriptor vaddr for '") + + Fixup.KernelName + "'") + .str()); + if (!ShiftedDescriptorVAddr) { + Ok = false; + continue; + } + DescriptorVAddr = *ShiftedDescriptorVAddr; + } + + if (!OutElf.findKernelDescriptorByVAddr(DescriptorVAddr)) { log() << "hotswap: error: missing kernel descriptor for entry " - << "trampoline fixup '" << Fixup.KernelName << "'.\n"; + << "trampoline fixup '" << Fixup.KernelName << "' at vaddr 0x" + << utohexstr(DescriptorVAddr) << ".\n"; Ok = false; continue; } @@ -563,7 +593,7 @@ bool rewriteKernelEntryDescriptorOffsets( continue; } std::optional NewOffset = checkedSignedDifference( - *StubVAddr, *KdVAddr, + *StubVAddr, DescriptorVAddr, (Twine("entry trampoline descriptor offset for '") + Fixup.KernelName + "'") .str()); @@ -571,10 +601,10 @@ bool rewriteKernelEntryDescriptorOffsets( Ok = false; continue; } - bool UpdatedEntry = - OutElf.updateKernelDescriptorEntryOffset(Fixup.KernelName, *NewOffset); + bool UpdatedEntry = OutElf.updateKernelDescriptorEntryOffset( + DescriptorVAddr, Fixup.KernelName, *NewOffset); bool UpdatedSgprs = OutElf.updateKernelDescriptorSgprCount( - Fixup.KernelName, Fixup.RequiredSgprs); + DescriptorVAddr, Fixup.KernelName, Fixup.RequiredSgprs); Ok = UpdatedEntry && UpdatedSgprs && Ok; } return Ok; diff --git a/amd/comgr/src/comgr-hotswap-internal.h b/amd/comgr/src/comgr-hotswap-internal.h index 9b59fa74e9e44..b34d8619c05b4 100644 --- a/amd/comgr/src/comgr-hotswap-internal.h +++ b/amd/comgr/src/comgr-hotswap-internal.h @@ -211,6 +211,10 @@ class ElfView { /// or nullptr if not found. uint8_t *findKernelDescriptor(llvm::StringRef KernelName); + /// Pointer to the kernel_descriptor at \p DescriptorVAddr inside the buffer, + /// or nullptr if not found. + uint8_t *findKernelDescriptorByVAddr(uint64_t DescriptorVAddr); + /// Enumerate kernel descriptor symbols named ".kd" and read their /// current kernel_code_entry_byte_offset values. std::vector kernelDescriptors() const; @@ -224,10 +228,22 @@ class ElfView { bool updateKernelDescriptorEntryOffset(llvm::StringRef KernelName, int64_t NewEntryOffset); + /// Rewrite kernel_code_entry_byte_offset for the descriptor identified by + /// \p DescriptorVAddr. + bool updateKernelDescriptorEntryOffset(uint64_t DescriptorVAddr, + llvm::StringRef KernelName, + int64_t NewEntryOffset); + /// Ensure the kernel descriptor reserves at least \p RequiredSgprs SGPRs. bool updateKernelDescriptorSgprCount(llvm::StringRef KernelName, unsigned RequiredSgprs); + /// Ensure the descriptor identified by \p DescriptorVAddr reserves at least + /// \p RequiredSgprs SGPRs. + bool updateKernelDescriptorSgprCount(uint64_t DescriptorVAddr, + llvm::StringRef KernelName, + unsigned RequiredSgprs); + /// Read COMPUTE_PGM_RSRC3.INST_PREF_SIZE for \p KernelName. std::optional getKernelDescriptorInstPrefSize(llvm::StringRef KernelName, @@ -258,11 +274,12 @@ class ElfView { /// Read the SGPR count for \p KernelName from the \c amdhsa.kernels /// msgpack metadata note (\c .sgpr_count key), falling back to the kernel - /// descriptor when the metadata note is absent. On GFX10+ the kernel - /// descriptor's \c GRANULATED_WAVEFRONT_SGPR_COUNT is architecturally - /// reserved, so metadata is the only reliable source when present. - /// Returns std::nullopt if the matching metadata is malformed, the kernel is - /// missing from present metadata, or the descriptor fallback is unavailable. + /// descriptor when the metadata note is absent or the matching kernel entry + /// is incomplete. On GFX10+ the kernel descriptor's + /// \c GRANULATED_WAVEFRONT_SGPR_COUNT is architecturally reserved, so + /// metadata is the preferred source when present. Returns std::nullopt if + /// the metadata note itself is corrupt or the descriptor fallback is + /// unavailable. std::optional getKernelSgprCount(llvm::StringRef KernelName) const; /// Update the RSRC1 VGPR granule count in the kernel descriptor for @@ -689,6 +706,7 @@ HotswapPatchVTable &getHotswapPatchVTable(); struct KernelEntryTrampolineFixup { std::string KernelName; + uint64_t KernelDescriptorVAddr = 0; uint64_t StubTextOffset = 0; unsigned RequiredSgprs = 0; }; diff --git a/amd/comgr/test-unit/HotswapElfTest.cpp b/amd/comgr/test-unit/HotswapElfTest.cpp index 0121f25606239..08e8068f7965a 100644 --- a/amd/comgr/test-unit/HotswapElfTest.cpp +++ b/amd/comgr/test-unit/HotswapElfTest.cpp @@ -34,6 +34,21 @@ static unsigned readReservedSgprs(const std::vector &Bytes, 8; } +static void writeReservedSgprs(std::vector &Bytes, + uint64_t KernelDescriptorOffset, + unsigned Sgprs) { + namespace hsa = llvm::amdhsa; + + ASSERT_NE(Sgprs, 0u); + uint32_t Granulated = (Sgprs + 7) / 8 - 1; + uint32_t Rsrc1 = 0; + AMDHSA_BITS_SET(Rsrc1, hsa::COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT, + Granulated); + std::memcpy(Bytes.data() + KernelDescriptorOffset + + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc1), + &Rsrc1, sizeof(Rsrc1)); +} + // -- ElfView::create ---------------------------------------------------------- TEST(ElfView, RejectsTruncatedInput) { @@ -164,6 +179,38 @@ TEST(ElfView, GrowWithTrampolinesShiftsAllocSectionSymbols) { EXPECT_EQ(KDs[0].VAddr, Obj.RodataAddr + GrowthBytes); } +TEST(ElfView, GrowWithTrampolinesDoesNotShiftAllocSymbolBeforeTextByVAddr) { + static constexpr uint64_t GrowthBytes = 8; + + comgr_test::KernelDescriptorElfOptions Opts; + Opts.KernelName = "entry_kernel"; + Opts.TextAddr = 0x3000; + Opts.RodataAddr = 0x1000; + comgr_test::KernelDescriptorElf Obj = + comgr_test::makeKernelDescriptorElf(makeText(), Opts); + + llvm::Expected ViewOrErr = + ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); + ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); + + Trampoline T; + T.Bytes.assign(GrowthBytes, 0); + std::vector Trampolines; + Trampolines.push_back(T); + const uint8_t SNop[4] = {}; + std::unique_ptr Out = + ViewOrErr->growWithTrampolines(Trampolines, SNop); + ASSERT_NE(Out, nullptr); + + uint8_t *OutData = reinterpret_cast(Out->getBufferStart()); + llvm::Expected OutView = + ElfView::create(OutData, Out->getBufferSize()); + ASSERT_TRUE((bool)OutView) << llvm::toString(OutView.takeError()); + std::vector KDs = OutView->kernelDescriptors(); + ASSERT_EQ(KDs.size(), 1u); + EXPECT_EQ(KDs[0].VAddr, Obj.RodataAddr); +} + TEST(ElfView, UpdateKernelDescriptorSgprCountUpdatesMetadataAndDescriptor) { comgr_test::KernelDescriptorElfOptions Opts; Opts.KernelName = "entry_kernel"; @@ -183,7 +230,46 @@ TEST(ElfView, UpdateKernelDescriptorSgprCountUpdatesMetadataAndDescriptor) { EXPECT_GE(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 10u); } -TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsMissingMetadataCount) { +TEST(ElfView, UpdateKernelDescriptorSgprCountConvergesDescriptorToRequired) { + comgr_test::KernelDescriptorElfOptions Opts; + Opts.KernelName = "entry_kernel"; + Opts.MetadataSgprCount = 9; + comgr_test::KernelDescriptorElf Obj = + comgr_test::makeKernelDescriptorElf(makeText(), Opts); + writeReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset, 64); + + llvm::Expected ViewOrErr = + ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); + ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); + + ASSERT_TRUE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); + std::optional MetadataSgprs = + ViewOrErr->getKernelSgprCount("entry_kernel"); + ASSERT_TRUE(MetadataSgprs.has_value()); + EXPECT_EQ(*MetadataSgprs, 10u); + EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 16u); +} + +TEST(ElfView, UpdateKernelDescriptorSgprCountPreservesLargerMetadataCount) { + comgr_test::KernelDescriptorElfOptions Opts; + Opts.KernelName = "entry_kernel"; + Opts.MetadataSgprCount = 24; + comgr_test::KernelDescriptorElf Obj = + comgr_test::makeKernelDescriptorElf(makeText(), Opts); + + llvm::Expected ViewOrErr = + ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); + ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); + + ASSERT_TRUE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); + std::optional MetadataSgprs = + ViewOrErr->getKernelSgprCount("entry_kernel"); + ASSERT_TRUE(MetadataSgprs.has_value()); + EXPECT_EQ(*MetadataSgprs, 24u); + EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 24u); +} + +TEST(ElfView, UpdateKernelDescriptorSgprCountFallsBackOnMissingMetadataCount) { comgr_test::KernelDescriptorElfOptions Opts; Opts.KernelName = "entry_kernel"; Opts.MetadataOmitSgprCount = true; @@ -194,11 +280,14 @@ TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsMissingMetadataCount) { ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); - EXPECT_FALSE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); - EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 8u); + EXPECT_TRUE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); + EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 16u); + std::optional Sgprs = ViewOrErr->getKernelSgprCount("entry_kernel"); + ASSERT_TRUE(Sgprs.has_value()); + EXPECT_EQ(*Sgprs, 16u); } -TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsMissingMetadataKernel) { +TEST(ElfView, UpdateKernelDescriptorSgprCountFallsBackOnMissingMetadataKernel) { comgr_test::KernelDescriptorElfOptions Opts; Opts.KernelName = "entry_kernel"; Opts.MetadataKernelName = "other_kernel"; @@ -210,12 +299,15 @@ TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsMissingMetadataKernel) { ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); - EXPECT_EQ(ViewOrErr->getKernelSgprCount("entry_kernel"), std::nullopt); - EXPECT_FALSE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); - EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 8u); + std::optional InitialSgprs = + ViewOrErr->getKernelSgprCount("entry_kernel"); + ASSERT_TRUE(InitialSgprs.has_value()); + EXPECT_EQ(*InitialSgprs, 8u); + EXPECT_TRUE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); + EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 16u); } -TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsNonIntegerMetadataCount) { +TEST(ElfView, UpdateKernelDescriptorSgprCountFallsBackOnNonIntegerMetadata) { comgr_test::KernelDescriptorElfOptions Opts; Opts.KernelName = "entry_kernel"; Opts.MetadataSgprCountAsString = true; @@ -226,9 +318,12 @@ TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsNonIntegerMetadataCount) { ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); - EXPECT_EQ(ViewOrErr->getKernelSgprCount("entry_kernel"), std::nullopt); - EXPECT_FALSE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); - EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 8u); + std::optional InitialSgprs = + ViewOrErr->getKernelSgprCount("entry_kernel"); + ASSERT_TRUE(InitialSgprs.has_value()); + EXPECT_EQ(*InitialSgprs, 8u); + EXPECT_TRUE(ViewOrErr->updateKernelDescriptorSgprCount("entry_kernel", 10)); + EXPECT_EQ(readReservedSgprs(Obj.Bytes, Obj.KernelDescriptorOffset), 16u); } TEST(ElfView, UpdateKernelDescriptorSgprCountRejectsMetadataSizeChange) { diff --git a/amd/comgr/test-unit/HotswapMCTest.cpp b/amd/comgr/test-unit/HotswapMCTest.cpp index 38515ab66d1ff..ca0c055f30dd2 100644 --- a/amd/comgr/test-unit/HotswapMCTest.cpp +++ b/amd/comgr/test-unit/HotswapMCTest.cpp @@ -56,6 +56,72 @@ void ensureLLVMInitialized() { } } // namespace COMGR +namespace { +struct DebugPatchCapture { + bool Enabled = false; + unsigned AddSymbolsCalls = 0; + unsigned DebugLineCalls = 0; + size_t LastSymbolTrampolineCount = 0; + size_t LastDebugLineTrampolineCount = 0; + uint64_t LastRangesGrowth = 0; + uint64_t LastInfoGrowth = 0; + uint64_t LastFrameGrowth = 0; +}; + +DebugPatchCapture DebugCapture; + +static uint64_t descriptorEntryTarget(const KernelDescriptorInfo &KD) { + if (KD.EntryOffset >= 0) + return KD.VAddr + static_cast(KD.EntryOffset); + + uint64_t Magnitude = + static_cast(-(KD.EntryOffset + 1)) + static_cast(1); + return KD.VAddr - Magnitude; +} +} // namespace + +namespace COMGR { +namespace hotswap { +bool addTrampolineSymbols(llvm::WritableMemoryBuffer &, + llvm::ArrayRef Trampolines, uint64_t, + unsigned) { + if (DebugCapture.Enabled) { + ++DebugCapture.AddSymbolsCalls; + DebugCapture.LastSymbolTrampolineCount = Trampolines.size(); + } + return true; +} + +bool patchDebugLine(llvm::WritableMemoryBuffer &, + llvm::ArrayRef Trampolines, uint64_t, + uint64_t) { + if (DebugCapture.Enabled) { + ++DebugCapture.DebugLineCalls; + DebugCapture.LastDebugLineTrampolineCount = Trampolines.size(); + } + return true; +} + +void patchDebugRanges(uint8_t *, size_t, uint64_t, uint64_t, + uint64_t TrampTotal) { + if (DebugCapture.Enabled) + DebugCapture.LastRangesGrowth = TrampTotal; +} + +void patchDebugInfo(uint8_t *, size_t, uint64_t, uint64_t, + uint64_t TrampTotal) { + if (DebugCapture.Enabled) + DebugCapture.LastInfoGrowth = TrampTotal; +} + +void patchDebugFrame(uint8_t *, size_t, uint64_t, uint64_t, + uint64_t TrampTotal) { + if (DebugCapture.Enabled) + DebugCapture.LastFrameGrowth = TrampTotal; +} +} // namespace hotswap +} // namespace COMGR + // Build a TargetIdentifier for the gfx1250 test subtarget without features -- // production callers go through parseTargetIdentifier; here we populate // directly so the tests stay self-contained. @@ -520,6 +586,98 @@ TEST(KernelEntryTrampoline, PreservesInstPrefSizeAndAddsPrefetchGuard) { EXPECT_EQ(KDs[0].EntryOffset, static_cast(StubVAddr - *KdVAddr)); } +TEST(KernelEntryTrampoline, RewritesDuplicateKernelNamesByDescriptorVAddr) { + LLVMState S = initLLVM(makeGfx1250Ident()); + ASSERT_TRUE(S.Valid); + + llvm::SmallVector Text = assembleSingleInst("s_endpgm", S); + ASSERT_EQ(Text.size(), MinInstSize); + + comgr_test::KernelDescriptorElfOptions Opts; + Opts.EmitDuplicateKernelDescriptorSymbol = true; + comgr_test::KernelDescriptorElf Obj = + comgr_test::makeKernelDescriptorElf(Text, Opts); + llvm::Expected ViewOrErr = + ElfView::create(Obj.Bytes.data(), Obj.Bytes.size()); + ASSERT_TRUE((bool)ViewOrErr) << llvm::toString(ViewOrErr.takeError()); + + std::vector InputKDs = ViewOrErr->kernelDescriptors(); + ASSERT_EQ(InputKDs.size(), 2u); + EXPECT_EQ(InputKDs[0].KernelName, InputKDs[1].KernelName); + EXPECT_NE(InputKDs[0].VAddr, InputKDs[1].VAddr); + + std::vector Growth; + std::vector Fixups; + std::optional Count = appendKernelEntryTrampolines( + *ViewOrErr, S, /*MaxSgprs=*/106, Growth, Fixups); + ASSERT_TRUE(Count.has_value()); + EXPECT_EQ(*Count, 2u); + ASSERT_EQ(Fixups.size(), 2u); + EXPECT_NE(Fixups[0].KernelDescriptorVAddr, Fixups[1].KernelDescriptorVAddr); + + uint64_t OldTextSize = ViewOrErr->textSize(); + std::unique_ptr Out = + ViewOrErr->growWithTrampolines(Growth, S.SNopBytes); + ASSERT_NE(Out, nullptr); + ASSERT_TRUE(rewriteKernelEntryDescriptorOffsets(*Out, OldTextSize, Fixups)); + + uint8_t *OutData = reinterpret_cast(Out->getBufferStart()); + llvm::Expected OutView = + ElfView::create(OutData, Out->getBufferSize()); + ASSERT_TRUE((bool)OutView) << llvm::toString(OutView.takeError()); + ASSERT_GE(OutView->textSize(), OldTextSize); + uint64_t GrowthBytes = OutView->textSize() - OldTextSize; + + std::vector OutputKDs = OutView->kernelDescriptors(); + ASSERT_EQ(OutputKDs.size(), 2u); + for (const KernelEntryTrampolineFixup &Fixup : Fixups) { + uint64_t ExpectedKdVAddr = Fixup.KernelDescriptorVAddr + GrowthBytes; + uint64_t ExpectedStubVAddr = + ViewOrErr->textAddr() + OldTextSize + Fixup.StubTextOffset; + bool Found = false; + for (const KernelDescriptorInfo &KD : OutputKDs) { + if (KD.VAddr != ExpectedKdVAddr) + continue; + Found = true; + EXPECT_EQ(descriptorEntryTarget(KD), ExpectedStubVAddr); + } + EXPECT_TRUE(Found); + } +} + +TEST(KernelEntryTrampoline, EntryOnlyDebugShiftUsesTotalGrowthWithoutSymbols) { + LLVMState S = initLLVM(makeGfx1250Ident()); + ASSERT_TRUE(S.Valid); + + llvm::SmallVector Text = assembleSingleInst("s_endpgm", S); + ASSERT_EQ(Text.size(), MinInstSize); + + comgr_test::KernelDescriptorElf Obj = + comgr_test::makeKernelDescriptorElf(Text); + + DebugCapture = DebugPatchCapture(); + DebugCapture.Enabled = true; + + Gfx1250RewriteOptions Options; + Options.RunB0A0Patches = false; + Options.RunEntryTrampolines = true; + std::unique_ptr Out; + amd_comgr_status_t Status = retargetCodeObject( + Obj.Bytes.data(), Obj.Bytes.size(), makeGfx1250Ident(), Options, Out); + + DebugCapture.Enabled = false; + + ASSERT_EQ(Status, AMD_COMGR_STATUS_SUCCESS); + ASSERT_NE(Out, nullptr); + EXPECT_EQ(DebugCapture.AddSymbolsCalls, 1u); + EXPECT_EQ(DebugCapture.DebugLineCalls, 1u); + EXPECT_EQ(DebugCapture.LastSymbolTrampolineCount, 0u); + EXPECT_EQ(DebugCapture.LastDebugLineTrampolineCount, 0u); + EXPECT_GT(DebugCapture.LastRangesGrowth, 0u); + EXPECT_EQ(DebugCapture.LastInfoGrowth, DebugCapture.LastRangesGrowth); + EXPECT_EQ(DebugCapture.LastFrameGrowth, DebugCapture.LastRangesGrowth); +} + TEST(KernelEntryTrampoline, AlignsStubByVirtualAddress) { LLVMState S = initLLVM(makeGfx1250Ident()); ASSERT_TRUE(S.Valid); diff --git a/amd/comgr/test-unit/comgr-test-elf-utils.h b/amd/comgr/test-unit/comgr-test-elf-utils.h index 6bfc01f347538..0b50ba169d252 100644 --- a/amd/comgr/test-unit/comgr-test-elf-utils.h +++ b/amd/comgr/test-unit/comgr-test-elf-utils.h @@ -54,7 +54,9 @@ struct KernelDescriptorElfOptions { uint64_t TextAddr = 0x1000; uint64_t RodataAddr = 0x2000; bool EmitKernelDescriptorSymbol = true; + bool EmitDuplicateKernelDescriptorSymbol = false; std::optional KernelDescriptorSymbolValue; + std::optional DuplicateKernelDescriptorSymbolValue; uint32_t GroupSegmentFixedSize = 0; uint32_t ComputePgmRsrc3 = 0; std::optional MetadataKernelName; @@ -67,7 +69,9 @@ struct KernelDescriptorElf { std::vector Bytes; uint64_t RodataAddr = 0; uint64_t KernelDescriptorOffset = 0; + uint64_t DuplicateKernelDescriptorOffset = 0; int64_t EntryOffset = 0; + int64_t DuplicateEntryOffset = 0; }; inline std::string @@ -156,6 +160,9 @@ makeKernelDescriptorElf(llvm::ArrayRef Text, const bool HasMetadataNote = Options.MetadataSgprCount || Options.MetadataOmitSgprCount || Options.MetadataSgprCountAsString; + const bool HasDuplicateKd = Options.EmitKernelDescriptorSymbol && + Options.EmitDuplicateKernelDescriptorSymbol; + const uint64_t RodataBytes = KdBytes * (HasDuplicateKd ? 2 : 1); std::vector MetadataNote; if (HasMetadataNote) { std::string MetadataBlob = makeAmdgpuMetadataBlob(Options); @@ -163,9 +170,10 @@ makeKernelDescriptorElf(llvm::ArrayRef Text, } const uint64_t RodataOff = alignTo8(TextOffset + Text.size()); - const uint64_t StrTabOff = alignTo8(RodataOff + KdBytes); + const uint64_t StrTabOff = alignTo8(RodataOff + RodataBytes); const uint64_t SymTabOff = alignTo8(StrTabOff + StrTab.size()); - const uint64_t SymCount = Options.EmitKernelDescriptorSymbol ? 3 : 2; + const uint64_t SymCount = + Options.EmitKernelDescriptorSymbol ? (HasDuplicateKd ? 4 : 3) : 2; const uint64_t ShStrTabOff = alignTo8(SymTabOff + SymCount * sizeof(Elf64_Sym)); const uint64_t NoteOff = @@ -180,11 +188,19 @@ makeKernelDescriptorElf(llvm::ArrayRef Text, Result.Bytes.assign(BufSize, 0); Result.RodataAddr = Options.RodataAddr; Result.KernelDescriptorOffset = RodataOff; + Result.DuplicateKernelDescriptorOffset = + HasDuplicateKd ? RodataOff + KdBytes : 0; const uint64_t KernelDescriptorAddr = Options.KernelDescriptorSymbolValue.value_or( Options.ElfType == ET_REL ? 0 : Options.RodataAddr); + const uint64_t DuplicateKernelDescriptorAddr = + Options.DuplicateKernelDescriptorSymbolValue.value_or( + Options.ElfType == ET_REL ? KdBytes : Options.RodataAddr + KdBytes); Result.EntryOffset = static_cast(Options.TextAddr) - static_cast(Options.RodataAddr); + Result.DuplicateEntryOffset = + static_cast(Options.TextAddr) - + static_cast(DuplicateKernelDescriptorAddr); uint8_t *Buf = Result.Bytes.data(); std::memcpy(Buf + TextOffset, Text.data(), Text.size()); @@ -225,7 +241,7 @@ makeKernelDescriptorElf(llvm::ArrayRef Text, RodataSh.sh_flags = SHF_ALLOC; RodataSh.sh_offset = RodataOff; RodataSh.sh_addr = Options.RodataAddr; - RodataSh.sh_size = KdBytes; + RodataSh.sh_size = RodataBytes; RodataSh.sh_addralign = 8; std::memcpy(Buf + ShOff + 2 * sizeof(Elf64_Shdr), &RodataSh, sizeof(RodataSh)); @@ -269,13 +285,27 @@ makeKernelDescriptorElf(llvm::ArrayRef Text, Buf + RodataOff + offsetof(hsa::kernel_descriptor_t, kernel_code_entry_byte_offset), &Result.EntryOffset, sizeof(Result.EntryOffset)); + if (HasDuplicateKd) + std::memcpy( + Buf + Result.DuplicateKernelDescriptorOffset + + offsetof(hsa::kernel_descriptor_t, kernel_code_entry_byte_offset), + &Result.DuplicateEntryOffset, sizeof(Result.DuplicateEntryOffset)); std::memcpy(Buf + RodataOff + offsetof(hsa::kernel_descriptor_t, group_segment_fixed_size), &Options.GroupSegmentFixedSize, sizeof(Options.GroupSegmentFixedSize)); + if (HasDuplicateKd) + std::memcpy( + Buf + Result.DuplicateKernelDescriptorOffset + + offsetof(hsa::kernel_descriptor_t, group_segment_fixed_size), + &Options.GroupSegmentFixedSize, sizeof(Options.GroupSegmentFixedSize)); std::memcpy(Buf + RodataOff + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc3), &Options.ComputePgmRsrc3, sizeof(Options.ComputePgmRsrc3)); + if (HasDuplicateKd) + std::memcpy(Buf + Result.DuplicateKernelDescriptorOffset + + offsetof(hsa::kernel_descriptor_t, compute_pgm_rsrc3), + &Options.ComputePgmRsrc3, sizeof(Options.ComputePgmRsrc3)); Elf64_Sym KernelSym{}; KernelSym.st_name = KernelNameOff; @@ -294,6 +324,16 @@ makeKernelDescriptorElf(llvm::ArrayRef Text, KdSym.st_value = KernelDescriptorAddr; KdSym.st_size = KdBytes; std::memcpy(Buf + SymTabOff + 2 * sizeof(Elf64_Sym), &KdSym, sizeof(KdSym)); + if (HasDuplicateKd) { + Elf64_Sym DuplicateKdSym{}; + DuplicateKdSym.st_name = KdNameOff; + DuplicateKdSym.setBindingAndType(STB_GLOBAL, STT_OBJECT); + DuplicateKdSym.st_shndx = 2; + DuplicateKdSym.st_value = DuplicateKernelDescriptorAddr; + DuplicateKdSym.st_size = KdBytes; + std::memcpy(Buf + SymTabOff + 3 * sizeof(Elf64_Sym), &DuplicateKdSym, + sizeof(DuplicateKdSym)); + } } return Result;