Merge lazy loading support for code objects

Add lazy loading support for code objects

Author: rwvo

include/kernelDB.h

@@ -53,6 +53,7 @@ THE SOFTWARE.
 #include <vector>
 #include <thread>
 #include <mutex>
+#include <condition_variable>
 #include <utility>
 #include <shared_mutex>
 #include <filesystem>
@@ -113,6 +114,7 @@ struct KernelArgument {
 bool buildDwarfAddressMap(const char* filename, size_t offset, size_t hsaco_length, std::map<Dwarf_Addr, SourceLocation>& addressMap);
 SourceLocation getSourceLocation(const std::map<Dwarf_Addr, SourceLocation>& 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<std::string>& 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<std::string> 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<basicBlock>&);
     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<CDNAKernel> 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<std::string> 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<Dwarf_Addr, SourceLocation>& addrMap);
+    void processKernelsWithAddressMap(const std::map<Dwarf_Addr, SourceLocation>& 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<std::string, std::vector<std::string>> file_map_;
     std::set<std::string> 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<std::string, LazyKernelEntry> lazy_kernels_;
+    /// Kernels currently being loaded — prevents concurrent disassembly of the same kernel.
+    std::set<std::string> loading_kernels_;
+    std::mutex loading_mutex_;
+    std::condition_variable loading_cv_;
     std::shared_mutex mutex_;
 };
 

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]:
         """

src/disassemble.cc

@@ -24,7 +24,7 @@ THE SOFTWARE.
 #include <sstream>
 #include "include/kernelDB.h"
 
-std::vector<std::string> disassembly_params = {"-d ", "--arch-name=amdgcn"};
+std::vector<std::string> disassembly_params = {"-d", "--arch-name=amdgcn"};
 
 void readFileToString(const std::string& filename, std::string& content) {
     std::ifstream file(filename, std::ios::binary);
@@ -72,33 +72,73 @@ 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;
 
     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<std::string> 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<std::string>& 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;