TheRock/rocm-systems/shared/kpack/python/rocm_kpack/kpack_transform.py

"""
Generic kpack transformation API.

This module provides format-agnostic versions of the kpack transformation
functions. They automatically detect whether the input binary is ELF or
PE/COFF and dispatch to the appropriate implementation.

Usage:
    from rocm_kpack import kpack_offload_binary, read_kpack_ref_marker

    # Works with both ELF and PE/COFF binaries
    result = kpack_offload_binary(input_path, output_path, ...)
    marker = read_kpack_ref_marker(binary_path)
"""

from pathlib import Path

from .format_detect import detect_binary_format, UnsupportedBinaryFormat


class NotFatBinaryError(ValueError):
    """Raised when kpack_offload_binary is called on a binary without device code.

    This is not necessarily an error condition - callers should catch this
    and skip processing for non-fat binaries.
    """

    pass


def kpack_offload_binary(
    input_path: Path,
    output_path: Path,
    kpack_search_paths: list[str] | None = None,
    kernel_name: str | None = None,
    verbose: bool = False,
) -> dict:
    """Transform a fat binary for kpack use (format auto-detected).

    This function automatically detects whether the input binary is ELF or
    PE/COFF format and delegates to the appropriate implementation.

    The transformation:
    1. Adds kpack marker section (if params provided)
    2. Updates wrapper pointers to point to marker section
    3. Rewrites HIPF->HIPK magic
    4. Zero-pages the fat binary section

    Args:
        input_path: Path to input binary
        output_path: Path for output binary
        kpack_search_paths: List of kpack file paths relative to binary.
            If None, assumes marker section already exists.
        kernel_name: Kernel identifier for TOC lookup.
            Required if kpack_search_paths is provided.
        verbose: Print progress information

    Returns:
        Dictionary with transformation results:
        - format: "elf" or "coff"
        - wrappers_found: Number of wrappers found and rewritten
        - bytes_zeroed: Number of bytes zeroed in fat binary section
        - original_size: Original file size
        - new_size: New file size

    Raises:
        NotFatBinaryError: If binary doesn't contain device code
        UnsupportedBinaryFormat: If binary is neither ELF nor PE/COFF
    """
    fmt = detect_binary_format(input_path)

    if fmt == "elf":
        from .elf.kpack_transform import kpack_offload_binary as elf_impl

        result = elf_impl(
            input_path,
            output_path,
            kpack_search_paths,
            kernel_name,
            verbose,
        )
        result["format"] = "elf"
        return result
    else:
        from .coff.kpack_transform import kpack_offload_binary as coff_impl

        result = coff_impl(
            input_path,
            output_path,
            kpack_search_paths,
            kernel_name,
            verbose,
        )
        result["format"] = "coff"
        return result


def read_kpack_ref_marker(binary_path: Path) -> dict | None:
    """Read kpack marker section from a binary (format auto-detected).

    This function automatically detects whether the binary is ELF or PE/COFF
    format and reads the marker section using the appropriate implementation.

    Args:
        binary_path: Path to binary with marker section

    Returns:
        Marker data dictionary with keys:
        - kpack_search_paths: list[str] of relative paths to kpack files
        - kernel_name: str kernel identifier for TOC lookup
        Returns None if section doesn't exist.

    Raises:
        RuntimeError: If section exists but cannot be read or parsed
        UnsupportedBinaryFormat: If binary is neither ELF nor PE/COFF
    """
    fmt = detect_binary_format(binary_path)

    if fmt == "elf":
        from .elf.kpack_transform import read_kpack_ref_marker as elf_impl

        return elf_impl(binary_path)
    else:
        from .coff.kpack_transform import read_kpack_ref_marker as coff_impl

        return coff_impl(binary_path)


def is_fat_binary(binary_path: Path) -> bool:
    """Check if a binary contains HIP fat binary sections.

    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 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)
    except UnsupportedBinaryFormat:
        return False

    if fmt == "elf":
        from .elf.surgery import ElfSurgery

        return ElfSurgery.has_fatbin_section(binary_path)
    else:
        from .coff.surgery import CoffSurgery

        return CoffSurgery.has_fatbin_section(binary_path)