TheRock/rocm-systems/shared/kpack/tests/test_unbundling_visitor.py

"""Tests for UnbundlingVisitor with bundled binary test assets."""

import shutil
from pathlib import Path

import pytest

from rocm_kpack.artifact_scanner import (
    ArtifactPath,
    ArtifactScanner,
    ArtifactVisitor,
    RecognizerRegistry,
)
from rocm_kpack.binutils import BundledBinary, Toolchain
from rocm_kpack.elf import ElfSurgery, add_kpack_ref_section, kpack_offload_binary


class UnbundlingVisitor(ArtifactVisitor):
    """Visitor that unbundles binaries and copies other artifacts.

    For bundled binaries:
    - Extracts each architecture bundle to {binary_name}.unbundled/{arch}.co
    - Creates host-only version of the binary (without .hip_fatbin section)

    For opaque files:
    - Copies verbatim preserving directory structure

    For kernel databases:
    - Copies all database artifacts preserving structure
    """

    def __init__(self, output_root: Path, toolchain: Toolchain):
        """Initialize visitor.

        Args:
            output_root: Root directory where artifacts will be written
            toolchain: Toolchain for unbundling operations
        """
        self.output_root = output_root
        self.toolchain = toolchain
        self.visited_opaque_files: list[Path] = []
        self.visited_bundled_binaries: list[Path] = []
        self.visited_databases: list[Path] = []

    def visit_opaque_file(self, artifact_path: ArtifactPath) -> None:
        """Copy opaque file verbatim."""
        self.visited_opaque_files.append(artifact_path.relative_path)
        dest = self.output_root / artifact_path.relative_path
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(artifact_path.absolute_path, dest)

    def visit_bundled_binary(
        self, artifact_path: ArtifactPath, bundled_binary: BundledBinary
    ) -> None:
        """Unbundle binary to .unbundled/ directory and create host-only version.

        Creates:
        - {binary_name}.unbundled/{arch}.co for each architecture
        - {binary_name} (host-only version without device code)
        """
        self.visited_bundled_binaries.append(artifact_path.relative_path)

        binary_name = artifact_path.relative_path.name
        binary_parent = artifact_path.relative_path.parent
        unbundled_dir_name = f"{binary_name}.unbundled"

        # Create .unbundled directory
        unbundled_dir = self.output_root / binary_parent / unbundled_dir_name
        unbundled_dir.mkdir(parents=True, exist_ok=True)

        # Unbundle to temporary directory
        with bundled_binary.unbundle(delete_on_close=True) as contents:
            # Copy only GPU architecture bundles (.hsaco files) to .unbundled/
            for filename in contents.file_names:
                if filename.endswith(".hsaco"):
                    src = contents.dest_dir / filename
                    # Extract architecture from filename like "hipv4-amdgcn-amd-amdhsa--gfx1100.hsaco"
                    # and rename to just "{arch}.co"
                    arch = filename.split("--")[-1].replace(".hsaco", "")
                    dest = unbundled_dir / f"{arch}.co"
                    shutil.copy2(src, dest)

        # Create host-only version (without device code)
        binary_dest = self.output_root / artifact_path.relative_path
        binary_dest.parent.mkdir(parents=True, exist_ok=True)
        kpack_offload_binary(bundled_binary.file_path, binary_dest)

    def visit_kernel_database(self, artifact_path: ArtifactPath, database) -> None:
        """Copy kernel database artifacts preserving structure."""
        self.visited_databases.append(artifact_path.relative_path)

        # Copy all artifacts from the database
        for kernel in database.get_kernel_artifacts():
            src = artifact_path.absolute_path / kernel.relative_path
            dest = self.output_root / artifact_path.relative_path / kernel.relative_path
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src, dest)


@pytest.fixture
def bundled_test_tree(tmp_path: Path, toolchain: Toolchain) -> Path:
    """Create a test tree with bundled binaries and opaque files.

    Uses actual test assets from test_assets/bundled_binaries/linux/cov5/.

    Structure:
        test_root/
            config.txt (opaque file)
            bins/
                test_kernel_multi.exe (bundled binary with .rocm_kpack_ref)
                libtest_kernel_single.so (bundled binary with .rocm_kpack_ref)
                host_only.exe (opaque file - no .hip_fatbin)
            data/
                input.dat (opaque file)
    """
    root = tmp_path / "test_root"
    root.mkdir()

    # Opaque files
    (root / "config.txt").write_text("configuration data")
    (root / "data").mkdir()
    (root / "data" / "input.dat").write_text("input data")

    # Setup bundled binaries from test assets
    bins_dir = root / "bins"
    bins_dir.mkdir()

    assets_dir = Path("test_assets/bundled_binaries/linux/cov5")

    # Add .rocm_kpack_ref sections to bundled binaries
    # (required for create_host_only to work)

    # Add marker to test_kernel_multi.exe
    surgery = ElfSurgery.load(assets_dir / "test_kernel_multi.exe")
    add_kpack_ref_section(
        surgery,
        kpack_search_paths=["test.kpack"],
        kernel_name="test_kernel",
    )
    surgery.save(bins_dir / "test_kernel_multi.exe")

    # Add marker to libtest_kernel_single.so
    surgery = ElfSurgery.load(assets_dir / "libtest_kernel_single.so")
    add_kpack_ref_section(
        surgery,
        kpack_search_paths=["test.kpack"],
        kernel_name="test_kernel",
    )
    surgery.save(bins_dir / "libtest_kernel_single.so")

    # Copy host-only binary (should be treated as opaque)
    shutil.copy2(assets_dir / "host_only.exe", bins_dir / "host_only.exe")

    return root


def test_unbundling_visitor_unbundles_binaries(
    bundled_test_tree: Path, tmp_path: Path, toolchain: Toolchain
):
    """Test that UnbundlingVisitor correctly unbundles bundled binaries."""
    output_dir = tmp_path / "output"
    output_dir.mkdir()

    registry = RecognizerRegistry()
    scanner = ArtifactScanner(registry, toolchain=toolchain)
    visitor = UnbundlingVisitor(output_dir, toolchain)

    # Scan and process the tree
    scanner.scan_tree(bundled_test_tree, visitor)

    # Verify opaque files were copied
    assert (output_dir / "config.txt").exists()
    assert (output_dir / "config.txt").read_text() == "configuration data"
    assert (output_dir / "data" / "input.dat").exists()
    assert (output_dir / "data" / "input.dat").read_text() == "input data"

    # Verify host-only binary was copied as opaque file
    assert (output_dir / "bins" / "host_only.exe").exists()

    # Verify bundled binaries were unbundled
    # Check test_kernel_multi.exe
    multi_unbundled = output_dir / "bins" / "test_kernel_multi.exe.unbundled"
    assert multi_unbundled.exists()
    assert multi_unbundled.is_dir()
    assert (multi_unbundled / "gfx1100.co").exists()
    assert (multi_unbundled / "gfx1101.co").exists()

    # Check libtest_kernel_single.so
    single_unbundled = output_dir / "bins" / "libtest_kernel_single.so.unbundled"
    assert single_unbundled.exists()
    assert single_unbundled.is_dir()
    assert (single_unbundled / "gfx1100.co").exists()

    # Verify host-only binaries were created
    assert (output_dir / "bins" / "test_kernel_multi.exe").exists()
    assert (output_dir / "bins" / "libtest_kernel_single.so").exists()

    # Note: executables get their PHDR table normalized to p_vaddr == p_offset
    # for old-kernel (EL8 4.18) direct exec; on a tiny fully-zero-paged fixture
    # that overhead can exceed the device-code savings. Real binaries shrink.
    host_only_multi_size = (
        (output_dir / "bins" / "test_kernel_multi.exe").stat().st_size
    )
    assert host_only_multi_size > 0, "Host-only binary should exist and be valid"

    host_only_single_size = (
        (output_dir / "bins" / "libtest_kernel_single.so").stat().st_size
    )
    # Note: Small binaries may grow due to ELF structural overhead.
    # The important assertion is that multi-arch binaries (with large .hip_fatbin)
    # shrink significantly. Small test binaries may have overhead that exceeds savings.
    # TODO(#3): Optimize overhead of section add.
    assert host_only_single_size > 0, "Host-only library should exist and be valid"


def test_unbundling_visitor_counts(
    bundled_test_tree: Path, tmp_path: Path, toolchain: Toolchain
):
    """Test that UnbundlingVisitor tracks visited artifacts correctly."""
    output_dir = tmp_path / "output"
    output_dir.mkdir()

    registry = RecognizerRegistry()
    scanner = ArtifactScanner(registry, toolchain=toolchain)
    visitor = UnbundlingVisitor(output_dir, toolchain)

    scanner.scan_tree(bundled_test_tree, visitor)

    # Verify counts
    # Opaque files: config.txt, data/input.dat, bins/host_only.exe
    assert len(visitor.visited_opaque_files) == 3

    # Bundled binaries: test_kernel_multi.exe, libtest_kernel_single.so
    assert len(visitor.visited_bundled_binaries) == 2

    # No kernel databases in this tree
    assert len(visitor.visited_databases) == 0

    # Verify specific paths
    assert Path("config.txt") in visitor.visited_opaque_files
    assert Path("data/input.dat") in visitor.visited_opaque_files
    assert Path("bins/host_only.exe") in visitor.visited_opaque_files
    assert Path("bins/test_kernel_multi.exe") in visitor.visited_bundled_binaries
    assert Path("bins/libtest_kernel_single.so") in visitor.visited_bundled_binaries


def test_unbundling_visitor_architecture_extraction(
    bundled_test_tree: Path, tmp_path: Path, toolchain: Toolchain
):
    """Test that unbundled .co files have correct architecture names."""
    output_dir = tmp_path / "output"
    output_dir.mkdir()

    registry = RecognizerRegistry()
    scanner = ArtifactScanner(registry, toolchain=toolchain)
    visitor = UnbundlingVisitor(output_dir, toolchain)

    scanner.scan_tree(bundled_test_tree, visitor)

    # Check that .co files exist and are non-empty
    gfx1100_co = output_dir / "bins" / "test_kernel_multi.exe.unbundled" / "gfx1100.co"
    gfx1101_co = output_dir / "bins" / "test_kernel_multi.exe.unbundled" / "gfx1101.co"

    assert gfx1100_co.exists()
    assert gfx1101_co.exists()

    # Verify they are non-empty (contain actual code)
    assert gfx1100_co.stat().st_size > 0
    assert gfx1101_co.stat().st_size > 0

    # Verify they are different (different architectures)
    assert gfx1100_co.read_bytes() != gfx1101_co.read_bytes()