diff --git a/include/kernelDB.h b/include/kernelDB.h index 3700833..32360fd 100644 --- a/include/kernelDB.h +++ b/include/kernelDB.h @@ -53,6 +53,7 @@ THE SOFTWARE. #include #include #include +#include #include #include #include @@ -113,6 +114,7 @@ struct KernelArgument { bool buildDwarfAddressMap(const char* filename, size_t offset, size_t hsaco_length, std::map& addressMap); SourceLocation getSourceLocation(const std::map& addrMap, Dwarf_Addr addr); __attribute__((visibility("default"))) bool getDisassembly(hsa_agent_t agent, const std::string& fileName, std::string& out); +bool getDisassemblyForSymbol(hsa_agent_t agent, const std::string& fileName, const std::string& symbolName, std::string& out); bool invokeProgram(const std::string& programName, const std::vector& params, const std::string& outputFileName); std::string create_temp_file_segment(const std::string& filename, std::streamoff offset, std::streamsize length); __attribute__((visibility("default"))) std::vector extractCodeObjects(hsa_agent_t agent, const std::string& fileName); @@ -210,7 +212,7 @@ class __attribute__((visibility("default"))) kernelDB { ~kernelDB(); bool getBasicBlocks(const std::string& name, std::vector&); CDNAKernel& getKernel(const std::string& name); - bool addFile(const std::string& name, hsa_agent_t agent, const std::string& strFilter); + bool addFile(const std::string& name, hsa_agent_t agent, const std::string& strFilter, bool lazy = false); bool parseDisassembly(const std::string& text); void mapDisassemblyToSource(hsa_agent_t agent, const char *elfFilePath); bool addKernel(std::unique_ptr kernel); @@ -227,9 +229,15 @@ class __attribute__((visibility("default"))) kernelDB { bool scanCodeObject(const std::string& co_file); bool hasKernel(const std::string& name); private: + bool scanCodeObjectForKernel(const std::string& co_file, const std::string& kernelName); + bool parseDisassemblyForKernel(const std::string& text, const std::string& targetKernel); + /// Get kernel symbol names from a .hsaco ELF without disassembling (reads .symtab). + static std::vector getKernelNamesFromElf(const std::string& fileName); + /// If kernel is lazy-loaded, disassemble its code object and fill kernels_; then remove from lazy set. + void ensureKernelLoaded(const std::string& name); void buildLineMap(size_t offset, size_t hsaco_length, const char *elfFilePath); void extractArgumentsFromDwarf(hsa_agent_t agent, const char *elfFilePath, bool resolve_typedefs); - void processKernelsWithAddressMap(const std::map& addrMap); + void processKernelsWithAddressMap(const std::map& addrMap, const std::string& targetKernel = ""); parse_mode getLineType(std::string& line); std::string extractKernelName(const std::string& line); static bool isBranch(const std::string& instruction); @@ -240,6 +248,17 @@ class __attribute__((visibility("default"))) kernelDB { std::string fileName_; std::map> file_map_; std::set scanned_code_objects_; + struct LazyKernelEntry { + std::string hsaco_path; + std::string logical_file; + std::string elf_symbol; // raw (mangled) ELF symbol name for --disassemble-symbols + }; + /// Lazy-loaded kernels: canonical name -> entry. Filled by addFile(..., lazy=true). + std::map lazy_kernels_; + /// Kernels currently being loaded — prevents concurrent disassembly of the same kernel. + std::set loading_kernels_; + std::mutex loading_mutex_; + std::condition_variable loading_cv_; std::shared_mutex mutex_; }; diff --git a/kerneldb/api.py b/kerneldb/api.py index 29b9503..bc1293e 100644 --- a/kerneldb/api.py +++ b/kerneldb/api.py @@ -38,14 +38,16 @@ class KernelDB: to query instruction-level information mapped to source lines. """ - def __init__(self, binary_path: Optional[str] = None, agent_id: Optional[int] = None): + def __init__(self, binary_path: Optional[str] = None, agent_id: Optional[int] = None, lazy: bool = False): """ Initialize KernelDB Args: - binary_path: Path to HSACO or HIP binary file. If empty string or None, - will search in the running process for fat binaries. + binary_path: Path to HSACO or HIP binary file. If None and lazy=False, + empty string loads the running process executable + shared libs. agent_id: HSA agent handle (if None, will use first GPU) + lazy: If True, create an empty DB; load binaries later with add_file() to avoid + loading the whole process. If False, load binary_path (or process) in constructor. """ # Initialize HSA status = _kerneldb.hsa_init() @@ -61,10 +63,53 @@ def __init__(self, binary_path: Optional[str] = None, agent_id: Optional[int] = if self.agent.handle == 0: raise RuntimeError("No GPU agent found") - # Create kernelDB instance (analysis happens in constructor) - binary_path = binary_path or "" - self._kdb = _kerneldb.KernelDB(self.agent, binary_path) - self.binary_path = binary_path + if lazy: + # Empty DB; add files with add_file() for lazy loading + self._kdb = _kerneldb.KernelDB(self.agent) + self.binary_path = None + else: + binary_path = binary_path or "" + self._kdb = _kerneldb.KernelDB(self.agent, binary_path) + self.binary_path = binary_path or None + + def add_file(self, path: str, filter: str = "", lazy: bool = True) -> bool: + """ + Add a binary (HIP executable or .hsaco). + + With lazy=True (default): only indexes kernel names and their code-object + locations—no disassembly. Disassembly is done on demand when you call + get_kernel(), get_kernel_lines(), get_instructions_for_line(), or access + assembly/arguments for a kernel. + + With lazy=False: full load (disassemble all code objects and map to source), + same as the previous behavior. + + Args: + path: Path to HIP fat binary or .hsaco file + filter: Optional kernel name filter (currently unused in C++) + lazy: If True (default), only index; disassemble on first use per code object. + + Returns: + True on success + """ + return self._kdb.add_file(path, self.agent, filter, lazy) + + def scan_code_object(self, co_file: str) -> bool: + """ + Scan a single .hsaco code object (disassembly + DWARF + args). + Idempotent if the code object was already scanned. + + Args: + co_file: Path to a .hsaco file + + Returns: + True on success + """ + return self._kdb.scan_code_object(co_file) + + def has_kernel(self, name: str) -> bool: + """Return True if a kernel with the given name exists.""" + return self._kdb.has_kernel(name) def get_kernels(self) -> List[str]: """ diff --git a/src/disassemble.cc b/src/disassemble.cc index 57f5fb6..cf0b6d0 100644 --- a/src/disassemble.cc +++ b/src/disassemble.cc @@ -24,7 +24,7 @@ THE SOFTWARE. #include #include "include/kernelDB.h" -std::vector disassembly_params = {"-d ", "--arch-name=amdgcn"}; +std::vector disassembly_params = {"-d", "--arch-name=amdgcn"}; void readFileToString(const std::string& filename, std::string& content) { std::ifstream file(filename, std::ios::binary); @@ -72,19 +72,19 @@ bool getDisassembly(hsa_agent_t agent, const std::string& fileName, std::string& ss << "--mcpu=" << name; parms.push_back(ss.str()); parms.push_back(fileName); - // Create a temporary file using tmpnam (note: tmpnam is not the most secure option) - char temp_filename[L_tmpnam]; - if (tmpnam(temp_filename) == nullptr) - throw std::runtime_error("Failed to generate temporary filename"); - else if (invokeProgram(disassembler, parms, temp_filename)) + char temp_template[] = "/tmp/kdb_dis_XXXXXX"; + int fd = mkstemp(temp_template); + if (fd < 0) + throw std::runtime_error("Failed to create temporary file"); + close(fd); + if (invokeProgram(disassembler, parms, temp_template)) { - // read file contents here - readFileToString(temp_filename, out); - unlink(temp_filename); + readFileToString(temp_template, out); + unlink(temp_template); return out.length() != 0; } - else - return false; + unlink(temp_template); + return false; } else return false; @@ -92,13 +92,53 @@ bool getDisassembly(hsa_agent_t agent, const std::string& fileName, std::string& return true; } +bool getDisassemblyForSymbol(hsa_agent_t agent, const std::string& fileName, + const std::string& symbolName, std::string& out) +{ + // Use --disassemble-symbols to extract only the requested kernel. + std::vector parms = {"--arch-name=amdgcn"}; + char name[64]; + memset(name, 0, sizeof(name)); + hsa_status_t status = hsa_agent_get_info(agent, HSA_AGENT_INFO_NAME, name); + if (status != HSA_STATUS_SUCCESS) + return false; + + std::stringstream ss; + ss << "--mcpu=" << name; + parms.push_back(ss.str()); + parms.push_back("--disassemble-symbols=" + symbolName); + parms.push_back(fileName); + + char temp_template[] = "/tmp/kdb_dis_XXXXXX"; + int fd = mkstemp(temp_template); + if (fd < 0) + throw std::runtime_error("Failed to create temporary file"); + close(fd); + + if (invokeProgram(disassembler, parms, temp_template)) + { + readFileToString(temp_template, out); + unlink(temp_template); + return out.length() != 0; + } + unlink(temp_template); + return false; +} + bool invokeProgram(const std::string& programName, const std::vector& params, const std::string& outputFileName) { - // Construct the command string + // Construct the command string with shell-safe quoting. + // Parameters are single-quoted so kernel names containing parentheses, + // spaces, or other shell metacharacters are passed through safely. std::stringstream command; command << programName; for (const auto& param : params) { - // Basic escaping of parameters (assumes no spaces in params; enhance if needed) - command << " " << param; + std::string escaped = param; + size_t pos = 0; + while ((pos = escaped.find('\'', pos)) != std::string::npos) { + escaped.replace(pos, 1, "'\\''"); + pos += 4; + } + command << " '" << escaped << "'"; } // Redirect stdout to outputFileName command << " > " << outputFileName; diff --git a/src/kernelDB.cc b/src/kernelDB.cc index 731f43c..f7d8b7f 100644 --- a/src/kernelDB.cc +++ b/src/kernelDB.cc @@ -256,14 +256,12 @@ kernelDB::~kernelDB() CDNAKernel& kernelDB::getKernel(const std::string& name) { + ensureKernelLoaded(name); std::shared_lock lock(mutex_); auto it = kernels_.find(getKernelName(name)); if (it != kernels_.end()) - { return *(it->second.get()); - } - else - throw std::runtime_error(name + " kernel does not exist."); + throw std::runtime_error(name + " kernel does not exist."); } @@ -277,27 +275,68 @@ bool kernelDB::addKernel(std::unique_ptr kernel) bool result = true; std::unique_lock lock(mutex_); std::string strName = kernel.get()->getName(); - // std::cout << "Adding kernel \"" << strName << "\"" << std::endl; - if (kernels_.find(strName) == kernels_.end()) + std::string canonical = getKernelName(strName); // same key form as lazy_kernels_ and getKernel lookup + if (kernels_.find(canonical) == kernels_.end()) { - kernels_[strName] = std::move(kernel); + kernels_[canonical] = std::move(kernel); } else { - // std::cout << "You're adding kernel \"" << strName << "\" which we've seen before. Something may be wrong." << std::endl; - kernels_[strName] = std::move(kernel); + kernels_[canonical] = std::move(kernel); result = false; } return result; } -bool kernelDB::addFile(const std::string& name, hsa_agent_t agent, const std::string& strFilter) +void kernelDB::ensureKernelLoaded(const std::string& name) +{ + std::string canonical = getKernelName(name); + std::string hsaco_path, elf_symbol; + + // Use loading_mutex_ to serialize the check-and-claim of loading_kernels_. + { + std::unique_lock lk(loading_mutex_); + + // If another thread is currently loading this kernel, wait for it. + while (loading_kernels_.count(canonical)) + loading_cv_.wait(lk); + + // Check lazy_kernels_ under shared_mutex. + { + std::shared_lock rlock(mutex_); + auto it = lazy_kernels_.find(canonical); + if (it == lazy_kernels_.end()) + return; // Already loaded (or was never lazy). + hsaco_path = it->second.hsaco_path; + elf_symbol = it->second.elf_symbol; + } + + loading_kernels_.insert(canonical); + } + + bool ok = scanCodeObjectForKernel(hsaco_path, elf_symbol); + + { + std::unique_lock wlock(mutex_); + // Always remove from lazy_kernels_: on success the kernel is now in + // kernels_; on failure we avoid retrying a doomed disassembly on every + // subsequent getKernel() call. + lazy_kernels_.erase(canonical); + } + { + std::lock_guard lk(loading_mutex_); + loading_kernels_.erase(canonical); + } + loading_cv_.notify_all(); +} + +bool kernelDB::addFile(const std::string& name, hsa_agent_t agent, const std::string& strFilter, bool lazy) { bool bReturn = true; amd_comgr_data_t executable; bool bValidExecutable = false; std::vector isas = ::kernelDB::getIsaList(agent); - std::cout << "Adding " << name << std::endl; + std::cout << "Adding " << name << (lazy ? " (lazy)" : "") << std::endl; if (name.ends_with(".hsaco")) { @@ -320,29 +359,47 @@ bool kernelDB::addFile(const std::string& name, hsa_agent_t agent, const std::st } } - if(bValidExecutable && isas.size()) + if (!bValidExecutable || !isas.size()) + { + bReturn = false; + return bReturn; + } + + if (lazy) { - // Disassemble and parse each code object to discover all kernels + // Only index: get kernel names from ELF symbol table; no disassembly yet. for (const auto& hsaco : file_map_[name]) { - std::string strDisassembly; - getDisassembly(agent, hsaco, strDisassembly); - parseDisassembly(strDisassembly); + std::vector rawNames = getKernelNamesFromElf(hsaco); + std::unique_lock lock(mutex_); + for (const auto& raw : rawNames) + { + std::string demangled = demangleName(raw.c_str()); + std::string canonical = getKernelName(demangled); + if (canonical.empty()) + continue; + lazy_kernels_[canonical] = {hsaco, name, raw}; + } } + return true; + } - // Then map all kernels to source - try - { - mapDisassemblyToSource(agent, name.c_str()); - } - catch (const std::runtime_error& e) - { - std::cerr << "Error adding " << name << "\n\t" << e.what() << std::endl; - bReturn = false; - } + // Eager: disassemble and parse each code object, then map to source + for (const auto& hsaco : file_map_[name]) + { + std::string strDisassembly; + getDisassembly(agent, hsaco, strDisassembly); + parseDisassembly(strDisassembly); } - else + try + { + mapDisassemblyToSource(agent, name.c_str()); + } + catch (const std::runtime_error& e) + { + std::cerr << "Error adding " << name << "\n\t" << e.what() << std::endl; bReturn = false; + } return bReturn; } @@ -400,10 +457,84 @@ bool kernelDB::scanCodeObject(const std::string& co_file) return true; } +bool kernelDB::scanCodeObjectForKernel(const std::string& co_file, const std::string& kernelName) +{ + std::string strDisassembly; + // Use targeted disassembly (--disassemble-symbols) to extract only the + // requested kernel instead of disassembling the entire code object. + // Try the FUNC symbol name first, then with .kd suffix (in case + // the caller passed a kernel descriptor name). + bool gotDisasm = false; + std::string funcName = kernelName; + if (funcName.size() > 3 && funcName.substr(funcName.size() - 3) == ".kd") + funcName = funcName.substr(0, funcName.size() - 3); + for (const auto& sym : {funcName, funcName + ".kd"}) + { + try { + gotDisasm = getDisassemblyForSymbol(agent_, co_file, sym, strDisassembly); + } catch (const std::exception&) { + gotDisasm = false; + } catch (...) { + gotDisasm = false; + } + if (gotDisasm) + break; + strDisassembly.clear(); + } + if (!gotDisasm) + return false; + + parseDisassemblyForKernel(strDisassembly, kernelName); + + std::map addrMap; + buildDwarfAddressMap(co_file.c_str(), 0, 0, addrMap); + + if (!addrMap.empty()) + { + processKernelsWithAddressMap(addrMap, kernelName); + } + + // Extract arguments only for the target kernel + std::map> kernelArgsMap; + extractKernelArguments(co_file.c_str(), 0, 0, kernelArgsMap, false); + + { + std::unique_lock lock(mutex_); + for (const auto& entry : kernelArgsMap) + { + std::string demangledName = demangleName(entry.first.c_str()); + std::string searchName = getKernelName(demangledName); + + auto it = kernels_.find(searchName); + if (it == kernels_.end()) + { + for (auto& k : kernels_) + { + if (k.first.compare(0, searchName.length(), searchName) == 0 && + (k.first.length() == searchName.length() || k.first[searchName.length()] == '(')) + { + k.second.get()->setArguments(entry.second); + break; + } + } + } + else + { + it->second.get()->setArguments(entry.second); + } + } + // Do NOT mark co_file in scanned_code_objects_ — other kernels from this + // code object may still need loading later. + } + + return true; +} + bool kernelDB::hasKernel(const std::string& name) { std::shared_lock lock(mutex_); - return kernels_.find(getKernelName(name)) != kernels_.end(); + std::string key = getKernelName(name); + return kernels_.find(key) != kernels_.end() || lazy_kernels_.find(key) != lazy_kernels_.end(); } std::string kernelDB::extractKernelName(const std::string& line) @@ -435,6 +566,8 @@ std::string kernelDB::extractKernelName(const std::string& line) parse_mode kernelDB::getLineType(std::string& line) { parse_mode result = BBLOCK; + if (line.empty()) + return result; auto it = line.begin(); if (*it == ':' || line.starts_with(".text:")) result = BEGIN; @@ -444,7 +577,7 @@ parse_mode kernelDB::getLineType(std::string& line) if (*it == ':') { // It's only a valid kernel if there are no spaces in lines that end with ':' - if(line.find_first_of(" ") == std::string::npos || *(--it) == '>') + if(line.find_first_of(" ") == std::string::npos || (line.size() > 1 && *(--it) == '>')) { if (line.ends_with("<.text>:")) result = BEGIN; @@ -463,11 +596,16 @@ void kernelDB::getBlockMarkers(const std::string& disassembly, std::map>::iterator it; + std::map>::iterator it = markers.end(); + bool haveKernel = false; do { parse_mode mode = getLineType(line); @@ -482,8 +620,9 @@ void kernelDB::getBlockMarkers(const std::string& disassembly, std::map tokens; split(line, tokens, " ", false); @@ -492,19 +631,33 @@ void kernelDB::getBlockMarkers(const std::string& disassembly, std::map tmp; split(addr, tmp, "+", false); - assert(tmp.size() == 2); - tmp[1].pop_back(); - it->second.insert(base_addr + std::stoull(tmp[1], nullptr, 16)); + if (tmp.size() == 2 && !tmp[1].empty()) + { + tmp[1].pop_back(); + try { + it->second.insert(base_addr + std::stoull(tmp[1], nullptr, 16)); + } catch (...) { + // malformed address — skip this branch marker + } + } } - else if (base_addr == 0) + else if (base_addr == 0 && tokens.size() > 1) { size_t i = 1; - while (tokens[i].find("//") == std::string::npos) + while (i < tokens.size() && tokens[i].find("//") == std::string::npos) i++; - std::string strAddress = tokens[++i]; - // remove the ending colon - strAddress.pop_back(); - base_addr = std::stoull(strAddress, nullptr, 16); + if (i + 1 < tokens.size()) + { + std::string strAddress = tokens[i + 1]; + // remove the ending colon + if (!strAddress.empty()) + strAddress.pop_back(); + try { + base_addr = std::stoull(strAddress, nullptr, 16); + } catch (...) { + base_addr = 0; + } + } } } }while(std::getline(in,line)); @@ -513,6 +666,18 @@ void kernelDB::getBlockMarkers(const std::string& disassembly, std::map> markers; getBlockMarkers(text, markers); - //std::cout << "Found " << markers.size() << " in getBlockMarkers"; std::map>::iterator mit; bool bDoingKernels = false; + bool skip = false; while(std::getline(in,line)) { bool blockCreated = false; @@ -540,26 +705,41 @@ bool kernelDB::parseDisassembly(const std::string& text) block_count = 0; break; case KERNEL: + { bDoingKernels = true; strKernel = extractKernelName(line); - //strKernel = line.substr(0, line.length() - 1); - kernel = std::make_unique(demangleName(strKernel.c_str())); - mit = markers.find(demangleName(strKernel.c_str())); - assert(mit != markers.end()); + std::string demangledName = demangleName(strKernel.c_str()); + std::string canonical = getKernelName(demangledName); + if (filterByKernel && canonical != targetCanonical) + { + skip = true; + current_kernel = nullptr; + current_block = nullptr; + break; + } + skip = false; + kernel = std::make_unique(demangledName); + mit = markers.find(demangledName); + if (mit == markers.end()) + { + skip = true; + current_kernel = nullptr; + current_block = nullptr; + break; + } current_kernel = kernel.get(); - mode=BBLOCK; + mode = BBLOCK; addKernel(std::move(kernel)); - break; + } case BBLOCK: - if (!bDoingKernels) + if (!bDoingKernels || skip) break; split(line, tokens, " ", false); if (tokens.size()) { if (!current_block) { - //std::cout << "Starting a new block:\n\t" << line << std::endl; block = std::make_unique(); block_count++; current_block = block.get(); @@ -568,24 +748,21 @@ bool kernelDB::parseDisassembly(const std::string& text) trim(tokens[0]); if (isBranch(tokens[0])) { - if (current_kernel) + if (current_kernel && block) + { current_kernel->addBlock(block_count, std::move(block)); - else + } + else if (!current_kernel) { - std::cout << "Disassembly parsing error. Processing a branch instruction when there's not a kernel currently defined.\n"; - std::cout << line << std::endl; - abort(); + current_block = nullptr; + continue; } if (tokens[0].find("s_endpgm") != std::string::npos) { blockCreated = true; - // std::cout << "New Block at endpgm\n\t" << line << std::endl; block = std::make_unique(); block_count++; current_block = block.get(); - // Going to let this drop down to add the s_endpgrm instruction to the block. - // This is the only branch instruction we include - // This should always be the last block and last instruction in the kernel } else { @@ -596,7 +773,6 @@ bool kernelDB::parseDisassembly(const std::string& text) std::vector inst_tokens; instruction_t inst; - // If there is more than one token and the first token contains underscores that means it's an instruction line if (tokens.size() > 1 && tokens[0].find("_") != std::string::npos) { split(tokens[0], inst_tokens, "_", false); @@ -611,21 +787,25 @@ bool kernelDB::parseDisassembly(const std::string& text) for (size_t i = 1; i < tokens.size() && tokens[i].find("//") == std::string::npos; i++) inst.operands_.push_back(tokens[i]); size_t i = 1; - while (tokens[i].find("//") == std::string::npos) + while (i < tokens.size() && tokens[i].find("//") == std::string::npos) i++; - std::string strAddress = tokens[++i]; - // remove the ending colon - strAddress.pop_back(); - inst.address_ = std::stoull(strAddress, nullptr, 16); + if (i + 1 < tokens.size()) + { + std::string strAddress = tokens[i + 1]; + if (!strAddress.empty()) + strAddress.pop_back(); + try { + inst.address_ = std::stoull(strAddress, nullptr, 16); + } catch (...) { + inst.address_ = 0; + } + } + else + { + inst.address_ = 0; + } if (!blockCreated && mit != markers.end() && (mit->second.find(inst.address_) != mit->second.end())) { - // This is the first line of a new block - // So add the current block and create a new one - // Not all blocks end in branches, so when we come to - // an address identified as the target of a conditional branch - // we don't care if the current block ends with a branch instruction - // we just save it and create a new one. - //std::cout << "Starting block at " << strAddress << std::endl; if (current_block) current_kernel->addBlock(block_count, std::move(block)); block = std::make_unique(); @@ -633,7 +813,8 @@ bool kernelDB::parseDisassembly(const std::string& text) current_block = block.get(); } inst.block_ = current_block; - current_block->addInstruction(inst); + if (current_block) + current_block->addInstruction(inst); if (inst.inst_ == "s_endpgm") { if (current_kernel && current_block) @@ -655,8 +836,6 @@ bool kernelDB::parseDisassembly(const std::string& text) } } - std::cout << "Marker Count: " << markers.size() << " Kernel Count: " << kernels_.size() << std::endl; - return bReturn; } @@ -705,6 +884,92 @@ void kernelDB::getElfSectionBits(const std::string &fileName, const std::string throw std::runtime_error("Section not found: " + sectionName); } +// AMDGPU kernel symbol type (STT_LOOS = 10); standard ELF STT_FUNC = 2 +#ifndef STT_AMDGPU_HSA_KERNEL +#define STT_AMDGPU_HSA_KERNEL 10 +#endif + +std::vector kernelDB::getKernelNamesFromElf(const std::string& fileName) +{ + std::vector names; + std::ifstream file(fileName, std::ios::binary); + if (!file) + return names; + + Elf64_Ehdr elfHeader; + file.read(reinterpret_cast(&elfHeader), sizeof(elfHeader)); + if (file.gcount() != sizeof(elfHeader) || memcmp(elfHeader.e_ident, ELFMAG, SELFMAG) != 0) + return names; + + file.seekg(elfHeader.e_shoff, std::ios::beg); + std::vector sectionHeaders(elfHeader.e_shnum); + file.read(reinterpret_cast(sectionHeaders.data()), elfHeader.e_shnum * sizeof(Elf64_Shdr)); + if (file.gcount() != static_cast(elfHeader.e_shnum * sizeof(Elf64_Shdr))) + return names; + + const Elf64_Shdr& shstrtab = sectionHeaders[elfHeader.e_shstrndx]; + std::vector shstrtabData(shstrtab.sh_size); + file.seekg(shstrtab.sh_offset, std::ios::beg); + file.read(shstrtabData.data(), shstrtab.sh_size); + + Elf64_Word textShndx = SHN_UNDEF; + size_t symtabOffset = 0, symtabSize = 0, strtabOffset = 0, strtabSize = 0; + + for (Elf64_Word i = 0; i < elfHeader.e_shnum; ++i) + { + std::string secName(&shstrtabData[sectionHeaders[i].sh_name]); + if (secName == ".text") + textShndx = i; + else if (secName == ".symtab") + { + symtabOffset = sectionHeaders[i].sh_offset; + symtabSize = sectionHeaders[i].sh_size; + } + else if (secName == ".strtab") + { + strtabOffset = sectionHeaders[i].sh_offset; + strtabSize = sectionHeaders[i].sh_size; + } + } + + if (textShndx == SHN_UNDEF || symtabSize == 0 || strtabSize == 0) + return names; + + std::vector strtabData(strtabSize); + file.seekg(strtabOffset, std::ios::beg); + file.read(strtabData.data(), strtabSize); + + const size_t numSyms = symtabSize / sizeof(Elf64_Sym); + file.seekg(symtabOffset, std::ios::beg); + for (size_t s = 0; s < numSyms; ++s) + { + Elf64_Sym sym; + file.read(reinterpret_cast(&sym), sizeof(sym)); + if (file.gcount() != static_cast(sizeof(sym))) + break; + if (sym.st_shndx == SHN_UNDEF) + continue; + uint8_t type = ELF64_ST_TYPE(sym.st_info); + bool isCode = (type == STT_FUNC || type == STT_AMDGPU_HSA_KERNEL); + bool isObject = (type == STT_OBJECT); + if (!isCode && !isObject) + continue; + if (isCode && sym.st_shndx != textShndx) + continue; + if (sym.st_name >= strtabSize) + continue; + const char* nameStr = &strtabData[sym.st_name]; + if (!nameStr[0]) + continue; + std::string symName(nameStr); + // Only keep .kd kernel descriptors from OBJECT symbols + if (isObject && (symName.size() < 3 || symName.substr(symName.size() - 3) != ".kd")) + continue; + names.push_back(symName); + } + return names; +} + //using namespace llvm; //using namespace llvm::object; @@ -858,12 +1123,19 @@ void CDNAKernel::getSourceCode(std::vector& outputLines) } -void kernelDB::processKernelsWithAddressMap(const std::map& addrMap) +void kernelDB::processKernelsWithAddressMap(const std::map& addrMap, const std::string& targetKernel) { + // Normalize the target so raw ELF symbol names match demangled kernel keys. + std::string targetCanonical = targetKernel.empty() ? "" : getKernelName(demangleName(targetKernel.c_str())); std::unique_lock lock(mutex_); auto it = kernels_.begin(); while(it != kernels_.end()) { + if (!targetCanonical.empty() && it->first != targetCanonical) + { + it++; + continue; + } const auto& blocks = it->second.get()->getBasicBlocks(); for (const auto& block : blocks) { @@ -889,7 +1161,6 @@ void kernelDB::processKernelsWithAddressMap(const std::mapsecond.get()->addLine(MISSING_SOURCE_INFO, inst); - //std::cout << "No match for " << std::hex << "0x" << instruction.address_ << std::dec << std::endl; } } } @@ -962,6 +1233,7 @@ void kernelDB::mapDisassemblyToSource(hsa_agent_t agent, const char *elfFilePath std::string kernelDB::getFileName(const std::string& kernel, size_t index) { + ensureKernelLoaded(kernel); std::shared_lock lock(mutex_); auto it = kernels_.find(getKernelName(kernel)); if (it != kernels_.end()) @@ -977,6 +1249,7 @@ std::string kernelDB::getFileName(const std::string& kernel, size_t index) std::vector kernelDB::getInstructionsForLine(const std::string& kernel_name, uint32_t line, const std::string& match) { + ensureKernelLoaded(kernel_name); std::shared_lock lock(mutex_); auto it = kernels_.find(getKernelName(kernel_name)); if (it != kernels_.end()) @@ -987,6 +1260,7 @@ std::vector kernelDB::getInstructionsForLine(const std::string& k const std::vector& kernelDB::getInstructionsForLine(const std::string& kernel_name, uint32_t line) { + ensureKernelLoaded(kernel_name); std::shared_lock lock(mutex_); auto it = kernels_.find(getKernelName(kernel_name)); if (it != kernels_.end()) @@ -998,30 +1272,35 @@ const std::vector& kernelDB::getInstructionsForLine(const std::st void kernelDB::getKernels(std::vector& out) { std::shared_lock lock(mutex_); - auto it = kernels_.begin(); - while (it != kernels_.end()) - { - out.push_back(it->first); - it++; - } + for (const auto& p : kernels_) + out.push_back(p.first); + for (const auto& p : lazy_kernels_) + out.push_back(p.first); } void kernelDB::getKernelLines(const std::string& kernel, std::vector& out) { + ensureKernelLoaded(kernel); std::shared_lock lock(mutex_); auto it = kernels_.find(getKernelName(kernel)); if (it != kernels_.end()) - { - it->second.get()->getLineNumbers(out); - } + it->second.get()->getLineNumbers(out); } std::vector kernelDB::getKernelArguments(const std::string& kernel_name, bool resolve_typedefs) { - // If we need to resolve typedefs and haven't done so yet, re-extract with resolution - if (resolve_typedefs) { - extractArgumentsFromDwarf(agent_, fileName_.c_str(), true); + std::string logical_file; + { + std::shared_lock lock(mutex_); + auto lit = lazy_kernels_.find(getKernelName(kernel_name)); + if (lit != lazy_kernels_.end()) + logical_file = lit->second.logical_file; } + ensureKernelLoaded(kernel_name); + if (resolve_typedefs && !logical_file.empty()) + extractArgumentsFromDwarf(agent_, logical_file.c_str(), true); + else if (resolve_typedefs && !fileName_.empty()) + extractArgumentsFromDwarf(agent_, fileName_.c_str(), true); std::shared_lock lock(mutex_); auto it = kernels_.find(getKernelName(kernel_name)); diff --git a/src/pybind11_wrapper.cc b/src/pybind11_wrapper.cc index fbb2344..072ee86 100644 --- a/src/pybind11_wrapper.cc +++ b/src/pybind11_wrapper.cc @@ -127,8 +127,18 @@ PYBIND11_MODULE(_kerneldb, m) { "Create KernelDB instance with agent and filename", py::arg("agent"), py::arg("filename")) .def(py::init(), - "Create KernelDB instance with agent (searches process)", + "Create empty KernelDB for lazy loading (add files with add_file)", py::arg("agent")) + .def("add_file", &kernelDB::kernelDB::addFile, + "Add a binary. If lazy=True (default), only index kernel names (no disassembly). " + "Disassembly runs on demand when you get assembly/source/args for a kernel.", + py::arg("path"), py::arg("agent"), py::arg("filter") = "", py::arg("lazy") = true) + .def("scan_code_object", &kernelDB::kernelDB::scanCodeObject, + "Scan a single .hsaco code object (disassembly + DWARF + args). Idempotent if already scanned.", + py::arg("co_file")) + .def("has_kernel", &kernelDB::kernelDB::hasKernel, + "Return True if a kernel with the given name exists", + py::arg("name")) .def("get_kernel", &kernelDB::kernelDB::getKernel, "Get a kernel by name", py::arg("name"), diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py new file mode 100644 index 0000000..4db9836 --- /dev/null +++ b/tests/test_lazy_loading.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Tests for KernelDB lazy loading path: add_file(path, lazy=True) should +index kernel names without disassembly, then load on demand and produce +results identical to eager loading. +""" + +import subprocess +import tempfile +from pathlib import Path + +import pytest + +from conftest import requires_rocm + +kerneldb = pytest.importorskip("kerneldb", reason="kernelDB C++ extension not available") +KernelDB = kerneldb.KernelDB + +_HIP_SOURCE = r""" +#include + +__global__ void lazy_add(float* a, float* b, float* c, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + c[idx] = a[idx] + b[idx]; + } +} + +__global__ void lazy_mul(float* a, float* b, float* c, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + c[idx] = a[idx] * b[idx]; + } +} + +int main() { return 0; } +""" + +_binary_path = None + + +def _get_binary(): + global _binary_path + if _binary_path is not None: + return _binary_path + tmp = Path(tempfile.mkdtemp(prefix="kerneldb_lazy_test_")) + src = tmp / "lazy.cpp" + exe = tmp / "lazy" + src.write_text(_HIP_SOURCE) + r = subprocess.run(["hipcc", "-g", str(src), "-o", str(exe)], + capture_output=True, text=True) + if r.returncode != 0: + pytest.skip(f"hipcc compilation failed:\n{r.stderr}") + _binary_path = str(exe) + return _binary_path + + +def _find_kernel(kernels, fragment): + matches = [k for k in kernels if fragment in k] + if not matches: + pytest.skip(f"No kernel containing {fragment!r} found") + return matches[0] + + +@requires_rocm +def test_lazy_kernel_discovery(): + """Lazy add_file should discover the same kernels as eager loading.""" + path = _get_binary() + + eager = KernelDB(path) + eager_kernels = sorted(eager.get_kernels()) + + lazy = KernelDB(lazy=True) + lazy.add_file(path, lazy=True) + lazy_kernels = sorted(lazy.get_kernels()) + + assert lazy_kernels == eager_kernels + + +@requires_rocm +def test_lazy_has_kernel(): + """has_kernel() should return True for lazy-indexed kernels before loading.""" + path = _get_binary() + kdb = KernelDB(lazy=True) + kdb.add_file(path, lazy=True) + + kernels = kdb.get_kernels() + assert len(kernels) >= 2 + + for name in kernels: + assert kdb.has_kernel(name), f"has_kernel({name!r}) returned False" + + +@requires_rocm +def test_lazy_get_kernel_matches_eager(): + """Lazy-loaded kernel assembly should match eager-loaded assembly.""" + path = _get_binary() + + eager = KernelDB(path) + eager_kernel = eager.get_kernel(_find_kernel(eager.get_kernels(), "lazy_add")) + + lazy = KernelDB(lazy=True) + lazy.add_file(path, lazy=True) + lazy_kernel = lazy.get_kernel(_find_kernel(lazy.get_kernels(), "lazy_add")) + + assert lazy_kernel.assembly == eager_kernel.assembly + assert lazy_kernel.lines == eager_kernel.lines + + +@requires_rocm +def test_lazy_instructions_for_line(): + """Instructions retrieved via lazy loading should match eager loading.""" + path = _get_binary() + + eager = KernelDB(path) + eager_name = _find_kernel(eager.get_kernels(), "lazy_add") + eager_lines = eager.get_kernel_lines(eager_name) + + lazy = KernelDB(lazy=True) + lazy.add_file(path, lazy=True) + lazy_name = _find_kernel(lazy.get_kernels(), "lazy_add") + lazy_lines = lazy.get_kernel_lines(lazy_name) + + assert lazy_lines == eager_lines + + for line in eager_lines: + eager_insts = [i.disassembly for i in eager.get_instructions_for_line(eager_name, line)] + lazy_insts = [i.disassembly for i in lazy.get_instructions_for_line(lazy_name, line)] + assert lazy_insts == eager_insts, f"Mismatch at line {line}"