Address comments

Author: stellaraccident

shared/kpack/python/rocm_kpack/coff/surgery.py

@@ -179,6 +179,54 @@ def load(cls, path: Path) -> "CoffSurgery":
         data = bytearray(path.read_bytes())
         return cls(data, path)
 
+    @staticmethod
+    def has_fatbin_section(file_path: Path) -> bool:
+        """Fast check for .hip_fat section without loading the full binary.
+
+        Reads only the DOS header, PE signature, COFF header, and section
+        headers — a few KB total even for large binaries.
+        Returns False for non-PE files (wrong magic).
+        Raises on I/O errors or corrupt PE headers.
+        """
+        with open(file_path, "rb") as f:
+            dos = f.read(DOS_HEADER_SIZE)
+            if len(dos) < DOS_HEADER_SIZE:
+                return False
+            if dos[:2] != b"MZ":
+                return False
+
+            pe_offset = struct.unpack_from("<I", dos, 0x3C)[0]
+            f.seek(pe_offset)
+            pe_sig = f.read(4)
+            if pe_sig != PE_SIGNATURE:
+                return False
+
+            # COFF header: NumberOfSections at offset 2
+            coff = f.read(COFF_HEADER_SIZE)
+            if len(coff) < COFF_HEADER_SIZE:
+                return False
+            num_sections = struct.unpack_from("<H", coff, 2)[0]
+            opt_hdr_size = struct.unpack_from("<H", coff, 16)[0]
+
+            # Section headers start after optional header
+            section_offset = pe_offset + 4 + COFF_HEADER_SIZE + opt_hdr_size
+            f.seek(section_offset)
+            section_data = f.read(num_sections * SECTION_HEADER_SIZE)
+            if len(section_data) < num_sections * SECTION_HEADER_SIZE:
+                return False
+
+            # Check each section name for ".hip_fat"
+            target = b".hip_fat"
+            for i in range(num_sections):
+                name = section_data[
+                    i * SECTION_HEADER_SIZE : i * SECTION_HEADER_SIZE + 8
+                ]
+                # Section names are 8 bytes, null-padded
+                if name == target or name.rstrip(b"\x00") == target:
+                    return True
+
+            return False
+
     # =========================================================================
     # Properties
     # =========================================================================

shared/kpack/python/rocm_kpack/database_handlers.py

@@ -52,16 +52,29 @@ def detect(self, path: Path, prefix_root: Path) -> Optional[str]:
         Detect if a file belongs to this database type and extract bundle key.
 
         Args:
-            path: Path to the file
+            path: Path to the file (must be under prefix_root)
             prefix_root: Root of the prefix for relative path computation
 
         Returns:
             Bundle key (e.g., 'gfx1100', 'gfx11', 'gfx12_0') if file matches,
             None otherwise. Bundle keys correspond to entries in the
             rocm-bootstrap hierarchy.
+
+        Raises:
+            ValueError: If path is not under prefix_root (caller bug).
         """
         pass
 
+    def _relative_path(self, path: Path, prefix_root: Path) -> str:
+        """Compute forward-slash relative path, raising on bad input."""
+        try:
+            return path.relative_to(prefix_root).as_posix()
+        except ValueError:
+            raise ValueError(
+                f"{self.name()} handler: path {path} is not under "
+                f"prefix_root {prefix_root}"
+            ) from None
+
     def should_move(self, path: Path) -> bool:
         """
         Determine if this file should be moved to architecture-specific artifact.
@@ -88,32 +101,26 @@ def detect(self, path: Path, prefix_root: Path) -> Optional[str]:
 
         Pattern: lib/rocblas/library/*_gfx*.{co,hsaco,dat}
         """
-        try:
-            rel_path = path.relative_to(prefix_root)
-            # Use as_posix() for consistent forward-slash paths on all platforms
-            path_str = rel_path.as_posix()
-
-            # Check if it's in rocblas/library directory
-            if "rocblas/library" not in path_str:
-                return None
-
-            # Check file extension
-            if path.suffix not in [".co", ".hsaco", ".dat"]:
-                return None
-
-            # Extract architecture from filename
-            # Look for patterns like _gfx1100, _gfx1101, gfx1102, etc.
-            match = _GFX_ARCH_PATTERN.search(path.name)
-            if match:
-                return match.group(0)
-
-            # Some .dat files don't have architecture suffix but are generic
-            # We don't move those
+        path_str = self._relative_path(path, prefix_root)
+
+        # Check if it's in rocblas/library directory
+        if "rocblas/library" not in path_str:
             return None
 
-        except (ValueError, AttributeError):
+        # Check file extension
+        if path.suffix not in [".co", ".hsaco", ".dat"]:
             return None
 
+        # Extract architecture from filename
+        # Look for patterns like _gfx1100, _gfx1101, gfx1102, etc.
+        match = _GFX_ARCH_PATTERN.search(path.name)
+        if match:
+            return match.group(0)
+
+        # Some .dat files don't have architecture suffix but are generic
+        # We don't move those
+        return None
+
 
 class HipBLASLtHandler(DatabaseHandler):
     """Handler for hipBLASLt kernel files."""
@@ -127,29 +134,23 @@ def detect(self, path: Path, prefix_root: Path) -> Optional[str]:
 
         Pattern: lib/hipblaslt/library/*_gfx*.{co,hsaco,dat}
         """
-        try:
-            rel_path = path.relative_to(prefix_root)
-            # Use as_posix() for consistent forward-slash paths on all platforms
-            path_str = rel_path.as_posix()
-
-            # Check if it's in hipblaslt/library directory
-            if "hipblaslt/library" not in path_str:
-                return None
-
-            # Check file extension
-            if path.suffix not in [".co", ".hsaco", ".dat"]:
-                return None
-
-            # Extract architecture from filename
-            match = _GFX_ARCH_PATTERN.search(path.name)
-            if match:
-                return match.group(0)
+        path_str = self._relative_path(path, prefix_root)
 
+        # Check if it's in hipblaslt/library directory
+        if "hipblaslt/library" not in path_str:
             return None
 
-        except (ValueError, AttributeError):
+        # Check file extension
+        if path.suffix not in [".co", ".hsaco", ".dat"]:
             return None
 
+        # Extract architecture from filename
+        match = _GFX_ARCH_PATTERN.search(path.name)
+        if match:
+            return match.group(0)
+
+        return None
+
 
 class HipSparseLtHandler(DatabaseHandler):
     """Handler for hipSPARSELt Tensile kernel files."""
@@ -163,25 +164,20 @@ def detect(self, path: Path, prefix_root: Path) -> Optional[str]:
 
         Pattern: lib/hipsparselt/library/*_gfx*.{co,hsaco,dat}
         """
-        try:
-            rel_path = path.relative_to(prefix_root)
-            path_str = rel_path.as_posix()
-
-            if "hipsparselt/library" not in path_str:
-                return None
-
-            if path.suffix not in [".co", ".hsaco", ".dat"]:
-                return None
-
-            match = _GFX_ARCH_PATTERN.search(path.name)
-            if match:
-                return match.group(0)
+        path_str = self._relative_path(path, prefix_root)
 
+        if "hipsparselt/library" not in path_str:
             return None
 
-        except (ValueError, AttributeError):
+        if path.suffix not in [".co", ".hsaco", ".dat"]:
             return None
 
+        match = _GFX_ARCH_PATTERN.search(path.name)
+        if match:
+            return match.group(0)
+
+        return None
+
 
 class AotritonHandler(DatabaseHandler):
     """Handler for AOTriton kernel image directories.
@@ -219,22 +215,18 @@ def detect(self, path: Path, prefix_root: Path) -> Optional[str]:
         Returns:
             Bundle key (e.g., 'gfx11', 'gfx12_0', 'gfx942') or None.
         """
-        try:
-            rel_path = path.relative_to(prefix_root)
-            path_parts = rel_path.parts
+        path_str = self._relative_path(path, prefix_root)
+        path_parts = Path(path_str).parts
 
-            for i, part in enumerate(path_parts[:-1]):
-                if part == "aotriton.images" and i + 1 < len(path_parts):
-                    arch_dir = path_parts[i + 1]
-                    if arch_dir.startswith("amd-gfx"):
-                        raw_name = arch_dir[4:]  # strip "amd-" prefix
-                        return self._BUNDLE_MAP.get(raw_name, raw_name)
-                    break
-
-            return None
+        for i, part in enumerate(path_parts[:-1]):
+            if part == "aotriton.images" and i + 1 < len(path_parts):
+                arch_dir = path_parts[i + 1]
+                if arch_dir.startswith("amd-gfx"):
+                    raw_name = arch_dir[4:]  # strip "amd-" prefix
+                    return self._BUNDLE_MAP.get(raw_name, raw_name)
+                break
 
-        except (ValueError, AttributeError, IndexError):
-            return None
+        return None
 
 
 class MIOpenHandler(DatabaseHandler):
@@ -261,34 +253,29 @@ def detect(self, path: Path, prefix_root: Path) -> Optional[str]:
 
         Pattern: share/miopen/db/gfx*.{db.txt,fdb.txt,model}
         """
-        try:
-            rel_path = path.relative_to(prefix_root)
-            path_str = rel_path.as_posix()
-
-            if "miopen/db" not in path_str:
-                return None
-
-            if not path.is_file():
-                return None
+        path_str = self._relative_path(path, prefix_root)
 
-            # Match .model, .db.txt, .fdb.txt, .OpenCL.fdb.txt, .HIP.fdb.txt
-            name = path.name
-            if not (
-                name.endswith(".model")
-                or name.endswith(".db.txt")
-                or name.endswith(".fdb.txt")
-            ):
-                return None
-
-            match = _MIOPEN_ARCH_PATTERN.search(name)
-            if match:
-                return match.group(0)
+        if "miopen/db" not in path_str:
+            return None
 
+        if not path.is_file():
             return None
 
-        except (ValueError, AttributeError):
+        # Match .model, .db.txt, .fdb.txt, .OpenCL.fdb.txt, .HIP.fdb.txt
+        name = path.name
+        if not (
+            name.endswith(".model")
+            or name.endswith(".db.txt")
+            or name.endswith(".fdb.txt")
+        ):
             return None
 
+        match = _MIOPEN_ARCH_PATTERN.search(name)
+        if match:
+            return match.group(0)
+
+        return None
+
 
 # Registry of available handlers
 AVAILABLE_HANDLERS = {

shared/kpack/python/rocm_kpack/elf/surgery.py

@@ -148,6 +148,58 @@ def load(cls, path: Path) -> "ElfSurgery":
         data = bytearray(path.read_bytes())
         return cls(data, path)
 
+    @staticmethod
+    def has_fatbin_section(file_path: Path) -> bool:
+        """Fast check for .hip_fatbin section without loading the full binary.
+
+        Reads only the ELF header, section header table, and section name
+        string table — typically a few KB total even for multi-GB binaries.
+        Returns False for non-ELF files (wrong magic, too small, 32-bit).
+        Raises on I/O errors or corrupt ELF headers.
+        """
+        with open(file_path, "rb") as f:
+            ehdr = f.read(ELF64_EHDR_SIZE)
+            if len(ehdr) < ELF64_EHDR_SIZE:
+                return False
+            if ehdr[:4] != b"\x7fELF":
+                return False
+            if ehdr[4] != 2:  # ELFCLASS64
+                return False
+
+            e_shoff = struct.unpack_from("<Q", ehdr, 40)[0]
+            e_shentsize = struct.unpack_from("<H", ehdr, 58)[0]
+            e_shnum = struct.unpack_from("<H", ehdr, 60)[0]
+            e_shstrndx = struct.unpack_from("<H", ehdr, 62)[0]
+
+            if e_shoff == 0 or e_shnum == 0 or e_shstrndx >= e_shnum:
+                return False
+
+            # Read section header table
+            f.seek(e_shoff)
+            shtab = f.read(e_shentsize * e_shnum)
+            if len(shtab) < e_shentsize * e_shnum:
+                return False
+
+            # Read .shstrtab
+            strtab_entry = e_shstrndx * e_shentsize
+            sh_offset = struct.unpack_from("<Q", shtab, strtab_entry + 24)[0]
+            sh_size = struct.unpack_from("<Q", shtab, strtab_entry + 32)[0]
+            f.seek(sh_offset)
+            shstrtab = f.read(sh_size)
+
+            # Search for ".hip_fatbin" in section names
+            target = b".hip_fatbin\x00"
+            if target not in shstrtab:
+                return False
+
+            target_offset = shstrtab.index(target)
+            for i in range(e_shnum):
+                sh_name = struct.unpack_from("<I", shtab, i * e_shentsize)[0]
+                if sh_name == target_offset:
+                    return True
+
+            return False
+
     # =========================================================================
     # Properties
     # =========================================================================

shared/kpack/python/rocm_kpack/kpack_transform.py

@@ -129,13 +129,15 @@ def read_kpack_ref_marker(binary_path: Path) -> dict | None:
 def is_fat_binary(binary_path: Path) -> bool:
     """Check if a binary contains HIP fat binary sections.
 
-    This is a quick check that doesn't require full parsing.
+    Fast-path: reads only headers (ELF section header table or PE section
+    headers), not the full binary. Works for both ELF and PE/COFF.
 
     Args:
         binary_path: Path to binary
 
     Returns:
-        True if binary appears to contain device code, False otherwise
+        True if binary contains a .hip_fatbin (ELF) or .hip_fat (COFF)
+        section, False otherwise (including non-ELF/non-PE files).
     """
     try:
         fmt = detect_binary_format(binary_path)
@@ -145,10 +147,8 @@ def is_fat_binary(binary_path: Path) -> bool:
     if fmt == "elf":
         from .elf.surgery import ElfSurgery
 
-        surgery = ElfSurgery.load(binary_path)
-        return surgery.find_section(".hip_fatbin") is not None
+        return ElfSurgery.has_fatbin_section(binary_path)
     else:
         from .coff.surgery import CoffSurgery
 
-        surgery = CoffSurgery.load(binary_path)
-        return surgery.find_section(".hip_fat") is not None
+        return CoffSurgery.has_fatbin_section(binary_path)