import argparse from pathlib import Path import shutil import subprocess import tempfile from enum import Enum from typing import Any import msgpack from rocm_kpack.ccob_parser import ExtractedCodeObject, extract_code_objects_from_fatbin from rocm_kpack.coff.surgery import CoffSurgery from rocm_kpack.format_detect import detect_binary_format, UnsupportedBinaryFormat class BinaryType(Enum): """Type of bundled binary file.""" STANDALONE = "standalone" # .co files - directly in bundler format BUNDLED = ( "bundled" # Executables/libraries with device code section (ELF or PE/COFF) ) class Toolchain: """Manages configuration of various toolchain locations. Tools are lazily found and cached on first access, so construction never fails. Only when a specific tool is accessed will it be searched for and validated. """ def __init__( self, *, clang_offload_bundler: Path | None = None, objcopy: Path | None = None, objdump: Path | None = None, readelf: Path | None = None, ): # Store explicit paths (may be None) self._clang_offload_bundler_path = clang_offload_bundler self._objcopy_path = objcopy self._objdump_path = objdump self._readelf_path = readelf # Cached resolved paths self._clang_offload_bundler_cached: Path | None = None self._objcopy_cached: Path | None = None self._objdump_cached: Path | None = None self._readelf_cached: Path | None = None @staticmethod def configure_argparse(p: argparse.ArgumentParser): p.add_argument( "--clang-offload-bundler", type=Path, help="Path to clang-offload-bundler" ) @staticmethod def from_args(args: argparse.Namespace) -> "Toolchain": clang_offload_bundler: Path | None = args.clang_offload_bundler return Toolchain(clang_offload_bundler=clang_offload_bundler) def _validate_or_find( self, tool_file_name: str, explicit_path: Path | None ) -> Path: if explicit_path is None: found_path = shutil.which(tool_file_name) if found_path is None: raise OSError( f"Could not find tool '{tool_file_name}' on system path. " f"Set KPACK_LLVM_BIN to your LLVM bin directory." ) explicit_path = Path(found_path) if not explicit_path.exists(): raise OSError( f"Tool '{tool_file_name}' at path {explicit_path} does not exist" ) return explicit_path def _validate_or_find_with_fallback( self, tool_file_name: str, fallback_name: str, explicit_path: Path | None, ) -> Path: """Find a tool, trying fallback name if primary not found. Args: tool_file_name: Primary tool name (e.g., "objdump") fallback_name: Fallback tool name (e.g., "llvm-objdump") explicit_path: Explicit path if provided by user Returns: Path to found tool Raises: OSError: If neither tool can be found """ if explicit_path is not None: if not explicit_path.exists(): raise OSError( f"Tool '{tool_file_name}' at path {explicit_path} does not exist" ) return explicit_path # Try primary name first found_path = shutil.which(tool_file_name) if found_path is not None: return Path(found_path) # Try fallback name found_path = shutil.which(fallback_name) if found_path is not None: return Path(found_path) raise OSError( f"Could not find tool '{tool_file_name}' or '{fallback_name}' on system path. " f"Set KPACK_LLVM_BIN to your LLVM bin directory." ) @property def clang_offload_bundler(self) -> Path: """Get clang-offload-bundler path (lazy, cached).""" if self._clang_offload_bundler_cached is None: self._clang_offload_bundler_cached = self._validate_or_find( "clang-offload-bundler", self._clang_offload_bundler_path ) return self._clang_offload_bundler_cached @property def objcopy(self) -> Path: """Get objcopy path (lazy, cached).""" if self._objcopy_cached is None: self._objcopy_cached = self._validate_or_find("objcopy", self._objcopy_path) return self._objcopy_cached @property def readelf(self) -> Path: """Get readelf path (lazy, cached).""" if self._readelf_cached is None: self._readelf_cached = self._validate_or_find("readelf", self._readelf_path) return self._readelf_cached @property def objdump(self) -> Path: """Get objdump path (lazy, cached). Tries 'objdump' first, falls back to 'llvm-objdump' if not found. This allows the same code to work on Linux (GNU objdump) and Windows (LLVM objdump from ROCm). """ if self._objdump_cached is None: self._objdump_cached = self._validate_or_find_with_fallback( "objdump", "llvm-objdump", self._objdump_path ) return self._objdump_cached def exec_capture_text(self, args: list[str | Path]): return subprocess.check_output( [str(a) for a in args], stderr=subprocess.STDOUT ).decode() def exec(self, args: list[str | Path]): # Use check_output to capture stderr in exceptions (discarding the output) subprocess.check_output([str(a) for a in args], stderr=subprocess.STDOUT) class UnbundledContents: """Represents a directory of unbundled contents. This is a context manager that will optionally delete the contents on close. """ def __init__( self, source_binary: "BundledBinary", dest_dir: Path, delete_on_close: bool, target_list: list[tuple[str, str]], ): self.source_binary = source_binary self.dest_dir = dest_dir self.delete_on_close = delete_on_close self.target_list = target_list def __enter__(self) -> "UnbundledContents": return self def __exit__(self, exc_type, exc_value, traceback): self.close() @property def file_names(self) -> list[str]: return [kv[1] for kv in self.target_list] def close(self): if self.delete_on_close and self.dest_dir.exists(): shutil.rmtree(self.dest_dir) def __repr__(self): return f"UnbundledContents(dest_dir={self.dest_dir}, target_list={self.target_list})" class BundledBinary: """Represents a bundled binary at some path. Supports two types of bundled binaries: 1. STANDALONE - Files directly in clang-offload-bundler format (e.g., .co files) 2. BUNDLED - ELF binaries containing .hip_fatbin section (executables, shared libraries) For BUNDLED binaries, the .hip_fatbin section is extracted and treated as a bundler-format input. """ def __init__(self, file_path: Path, *, toolchain: Toolchain | None = None): # Initialize _temp_dir first to ensure cleanup works even if init fails self._temp_dir: Path | None = None # For extracted .hip_fatbin sections self.toolchain = toolchain or Toolchain() self.file_path = file_path self.binary_type = self._detect_binary_type() def unbundle( self, *, dest_dir: Path | None = None, delete_on_close: bool = True ) -> UnbundledContents: """Unbundles the binary, returning a context manager which can be used to hold the unbundled files open for as long as needed. """ if dest_dir is None: dest_dir = Path(tempfile.mkdtemp()) # Extract code objects once, then derive both the target list and # write the files from the same extraction result. code_objects = self._extract_code_objects() target_list = self._build_target_list(code_objects) contents = UnbundledContents( self, dest_dir, delete_on_close=delete_on_close, target_list=target_list ) try: dest_dir.mkdir(parents=True, exist_ok=True) outputs = [dest_dir / kv[1] for kv in target_list] self._write_code_objects(code_objects, outputs) except: contents.close() raise return contents def _detect_binary_type(self) -> BinaryType: """Detect if this is a standalone bundler file or a bundled binary. For ELF, checks for .hip_fatbin section via readelf. For PE/COFF, checks for .hip_fat section via CoffSurgery. Files without device code sections are STANDALONE (.co files in bundler format). Returns: BinaryType indicating the file type Raises: RuntimeError: For unexpected errors during detection """ try: fmt = detect_binary_format(self.file_path) except UnsupportedBinaryFormat: return BinaryType.STANDALONE if fmt == "elf": try: result = subprocess.run( [str(self.toolchain.readelf), "-S", str(self.file_path.resolve())], capture_output=True, text=True, check=True, ) if ".hip_fatbin" in result.stdout: return BinaryType.BUNDLED else: return BinaryType.STANDALONE except subprocess.CalledProcessError: return BinaryType.STANDALONE except Exception as e: raise RuntimeError( f"Unexpected error detecting binary type for {self.file_path}: {e}" ) from e else: # PE/COFF try: surgery = CoffSurgery.load(self.file_path) if surgery.find_section(".hip_fat") is not None: return BinaryType.BUNDLED else: return BinaryType.STANDALONE except Exception as e: raise RuntimeError( f"Unexpected error detecting binary type for {self.file_path}: {e}" ) from e def _get_bundler_input(self) -> Path: """Get the file path to use as input to clang-offload-bundler. For STANDALONE files, returns the file path directly. For BUNDLED ELF binaries, extracts .hip_fatbin section via objcopy. For BUNDLED PE/COFF binaries, extracts .hip_fat section via CoffSurgery. Returns: Path to file in bundler format """ if self.binary_type == BinaryType.STANDALONE: return self.file_path if self._temp_dir is None: self._temp_dir = Path(tempfile.mkdtemp()) fatbin_path = self._temp_dir / "fatbin.o" fmt = detect_binary_format(self.file_path) if fmt == "coff": surgery = CoffSurgery.load(self.file_path) section = surgery.find_section(".hip_fat") if section is None: raise RuntimeError( f"PE binary {self.file_path} has no .hip_fat section" ) content = surgery.get_section_content(section) content = content[: section.virtual_size] fatbin_path.write_bytes(content) else: # ELF: extract .hip_fatbin section via objcopy abs_file_path = self.file_path.resolve() abs_fatbin_path = fatbin_path.resolve() try: self.toolchain.exec( [ self.toolchain.objcopy, "--dump-section", f".hip_fatbin={abs_fatbin_path}", abs_file_path, ] ) except subprocess.CalledProcessError as e: error_output = e.output.decode() if e.output else "(no output)" raise RuntimeError( f"Failed to extract .hip_fatbin section from {self.file_path}. " f"objcopy exit code: {e.returncode}. Output: {error_output}" ) from e return fatbin_path def _extract_code_objects(self) -> list[ExtractedCodeObject]: """Extract code objects from the fatbin data. Reads the bundler input once and parses all code objects. """ bundler_input = self._get_bundler_input() data = bundler_input.read_bytes() return extract_code_objects_from_fatbin(data) def _build_target_list( self, code_objects: list[ExtractedCodeObject] ) -> list[tuple[str, str]]: """Build a list of (target_name, file_name) from extracted code objects. For concatenated bundles (RDC), filenames are indexed to distinguish multiple code objects with the same target triple. """ targets = [obj.target for obj in code_objects] needs_indexing = len(targets) != len(set(targets)) result = [] for i, obj in enumerate(code_objects): base_ext = ".hsaco" if obj.target.startswith("hip") else ".elf" if needs_indexing: filename = f"{obj.target}_{i}{base_ext}" else: filename = f"{obj.target}{base_ext}" result.append((obj.target, filename)) return result def _write_code_objects( self, code_objects: list[ExtractedCodeObject], outputs: list[Path] ) -> None: """Write extracted code objects to output files.""" if len(code_objects) != len(outputs): raise ValueError( f"Output count mismatch: {len(code_objects)} code objects " f"but {len(outputs)} output paths" ) for obj, output_path in zip(code_objects, outputs): output_path.write_bytes(obj.data) def _list_bundled_targets(self, file_path: Path) -> list[tuple[str, str]]: """Returns a list of (target_name, file_name) for all bundles. Uses our own bundle parser for all formats: CCOB-compressed, single uncompressed, and concatenated uncompressed (RDC case). """ code_objects = self._extract_code_objects() return self._build_target_list(code_objects) def list_bundles(self) -> list[str]: """List all architecture bundles in the binary. Returns: List of architecture strings (e.g., ['gfx1100', 'gfx1101']) Only returns GPU architectures, not host bundles. """ target_list = self._list_bundled_targets(self.file_path) architectures = [] for target_name, _ in target_list: # Extract architecture from target names like: # "hipv4-amdgcn-amd-amdhsa--gfx1100" -> "gfx1100" if target_name.startswith("hip"): parts = target_name.split("--") if len(parts) >= 2: architectures.append(parts[-1]) return architectures def remove_section_simple(self, output_path: Path, section_name: str) -> None: """Remove a section using objcopy (simple header removal, no space reclaimed). This is a simple section removal that only updates ELF headers without reclaiming disk space. Useful for benchmarking or comparison purposes. For production use with .hip_fatbin removal, use kpack_offload_binary() from rocm_kpack.elf instead, which properly reclaims disk space. Args: output_path: Path where modified binary will be written section_name: Name of section to remove (e.g., ".hip_fatbin") Raises: RuntimeError: If objcopy fails """ try: self.toolchain.exec( [ self.toolchain.objcopy, "--remove-section", section_name, self.file_path, output_path, ] ) except subprocess.CalledProcessError as e: raise RuntimeError( f"Failed to remove section {section_name} from {self.file_path}: {e}" ) def cleanup(self) -> None: """Clean up temporary files created during operations.""" if self._temp_dir and self._temp_dir.exists(): shutil.rmtree(self._temp_dir) self._temp_dir = None def __del__(self): """Cleanup on deletion.""" self.cleanup() def get_section_vaddr( toolchain: Toolchain, binary_path: Path, section_name: str ) -> int | None: """ Get the virtual address of a section in an ELF binary. Args: toolchain: Toolchain instance providing readelf binary_path: Path to ELF binary section_name: Name of section (e.g., ".custom_data", ".rocm_kpack_ref") Returns: Virtual address (sh_addr) of the section if it exists and has ALLOC flag, None otherwise. Note: Only returns addresses for sections with the ALLOC flag (A), which indicates they are mapped to memory at load time (part of a PT_LOAD segment). """ try: # Run readelf to get section headers result = subprocess.run( [str(toolchain.readelf), "-S", str(binary_path)], capture_output=True, text=True, check=True, ) except subprocess.CalledProcessError: return None # Parse section headers # Format (two-line entries): # Line 1: [Nr] Name Type Address Offset # Line 2: Size EntSize Flags Link Info Align lines = result.stdout.split("\n") for i, line in enumerate(lines): if section_name in line: parts = line.split() # Check if this is a section header line (starts with [Nr]) if len(parts) >= 5 and parts[0].startswith("["): try: # Address column is at index 3 vaddr = int(parts[3], 16) # Check flags on the next line if i + 1 < len(lines): next_parts = lines[i + 1].split() if len(next_parts) >= 3: flags = next_parts[2] # Only return address if section has ALLOC flag (A) if "A" in flags: return vaddr except (ValueError, IndexError): continue return None def get_section_size( toolchain: Toolchain, binary_path: Path, section_name: str ) -> int | None: """ Get the size of a section in an ELF binary. Args: toolchain: Toolchain instance providing readelf binary_path: Path to ELF binary section_name: Name of section (e.g., ".hipFatBinSegment") Returns: Size (sh_size) of the section if found, None otherwise. """ try: result = subprocess.run( [str(toolchain.readelf), "-S", str(binary_path)], capture_output=True, text=True, check=True, ) except subprocess.CalledProcessError: return None # Parse section headers # Format (two-line entries): # Line 1: [Nr] Name Type Address Offset # Line 2: Size EntSize Flags Link Info Align lines = result.stdout.split("\n") for i, line in enumerate(lines): if section_name in line: parts = line.split() # Check if this is a section header line (starts with [Nr]) if len(parts) >= 5 and parts[0].startswith("["): try: # Size is on the next line, first column if i + 1 < len(lines): next_parts = lines[i + 1].split() if len(next_parts) >= 1: size = int(next_parts[0], 16) return size except (ValueError, IndexError): continue return None def has_section( binary_path: Path, section_name: str, *, toolchain: Toolchain | None = None, ) -> bool: """Check if a binary has a specific section. Args: binary_path: Path to binary section_name: Name of section to check for (e.g., ".hip_fatbin", ".rocm_kpack_ref") toolchain: Toolchain instance (created if not provided) Returns: True if section exists, False otherwise Note: This function abstracts binary format tooling (readelf for ELF, etc.) to support cross-platform binary analysis. """ if toolchain is None: toolchain = Toolchain() try: output = toolchain.exec_capture_text( [toolchain.readelf, "-S", str(binary_path)] ) return section_name in output except Exception: return False def get_section_type( binary_path: Path, section_name: str, *, toolchain: Toolchain | None = None, ) -> str | None: """Get the type of a section in a binary (e.g., PROGBITS, NOBITS). Args: binary_path: Path to binary section_name: Name of section (e.g., ".hip_fatbin") toolchain: Toolchain instance (created if not provided) Returns: Section type string (e.g., "PROGBITS", "NOBITS"), or None if section doesn't exist Note: This function abstracts binary format tooling (readelf for ELF, etc.) to support cross-platform binary analysis. """ if toolchain is None: toolchain = Toolchain() try: output = toolchain.exec_capture_text( [toolchain.readelf, "-S", str(binary_path)] ) # Parse section headers to find the type # Format: [Nr] Name Type Address Offset for line in output.splitlines(): if section_name in line: parts = line.split() # Check if this is a section header line (starts with [Nr]) if len(parts) >= 3 and parts[0].startswith("["): # Type is at index 2 return parts[2] return None except Exception: return None