runtime_builder.py

# Copyright (c) PyPTO Contributors.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# -----------------------------------------------------------------------------------------------------------
import fcntl
import json
import logging
import shutil
import subprocess
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from .environment import PROJECT_ROOT
from .platform_info import TARGETS, load_build_config, parse_platform
from .runtime_compiler import RuntimeCompiler

logger = logging.getLogger(__name__)

_GIT_COMMIT_FILE = ".git_commit"


def _get_git_head(repo_root: Path) -> str:
    """Return the current git HEAD commit hash, or empty string if unavailable."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=str(repo_root),
            capture_output=True,
            text=True,
            timeout=5,
            check=False,
        )
        return result.stdout.strip() if result.returncode == 0 else ""
    except Exception:  # noqa: BLE001
        return ""


def _invalidate_cache_if_stale(target_cache_dir: Path, current_commit: str) -> None:
    """Clear target_cache_dir if it was built from a different git commit.

    git does not update file mtimes on checkout, so cmake's incremental build
    cannot detect that source files changed. Comparing the HEAD commit stored
    at last build time against the current HEAD is a reliable signal that
    sources may have changed and a clean rebuild is needed.

    When the current commit can't be determined (no git, transient failure),
    fall through to a clean rebuild — a fresh compile is cheap relative to
    the risk of linking against stale objects.
    """
    if not current_commit:
        if target_cache_dir.is_dir():
            logger.info("git HEAD unavailable, clearing cmake cache: %s", target_cache_dir)
            shutil.rmtree(target_cache_dir)
        target_cache_dir.mkdir(parents=True, exist_ok=True)
        return
    commit_file = target_cache_dir / _GIT_COMMIT_FILE
    if commit_file.is_file():
        cached_commit = commit_file.read_text().strip()
        if cached_commit == current_commit:
            return
        logger.info(
            "git HEAD changed (%s → %s), clearing cmake cache: %s",
            cached_commit[:12],
            current_commit[:12],
            target_cache_dir,
        )
        shutil.rmtree(target_cache_dir)
    target_cache_dir.mkdir(parents=True, exist_ok=True)
    commit_file.write_text(current_commit + "\n")


@dataclass
class RuntimeBinaries:
    """Paths to the compiled runtime binaries.

    ``dispatcher_path`` points at ``libsimpler_aicpu_dispatcher.so`` and is
    required for onboard platforms (host bootstrap reads its bytes and ships
    them to the device alongside the inner SO). Sim platforms have no
    dispatcher; the field is ``None`` there. ``_lookup_binaries`` resolves
    and validates the path against the build output directory.
    """

    host_path: Path
    aicpu_path: Path
    aicore_path: Path
    simpler_log_path: Path
    sim_context_path: Optional[Path] = None
    dispatcher_path: Optional[Path] = None


class RuntimeBuilder:
    """Discovers and builds runtime implementations from src/runtime/.

    Accepts a platform selection to provide correctly configured
    RuntimeCompiler and KernelCompiler instances. Runtime and platform
    are orthogonal — the same runtime (e.g., host_build_graph) can
    be compiled for any platform (e.g., a2a3, a2a3sim).
    """

    _CACHE_DIR = PROJECT_ROOT / "build" / "cache"
    _LIB_DIR = PROJECT_ROOT / "build" / "lib"

    # Defaults for compile_commands.json placement (matches old gen_compile_commands.py).
    # Platform dirs get compdb from the most feature-rich runtime;
    # runtime dirs get compdb from the onboard (real hardware) variant.
    _COMPDB_RUNTIME = "tensormap_and_ringbuffer"
    _COMPDB_VARIANT = "onboard"

    def __init__(self, platform: str = "a2a3"):
        """
        Initialize RuntimeBuilder with platform selection.

        Args:
            platform: Target platform ("a2a3", "a2a3sim", "a5", or "a5sim")
        """
        self.platform = platform
        self._arch, self._variant = parse_platform(platform)

        runtime_root = PROJECT_ROOT
        self.runtime_root = runtime_root

        self.runtime_dir = runtime_root / "src" / self._arch / "runtime"

        # Discover available runtime implementations
        self._runtimes = {}
        if self.runtime_dir.is_dir():
            for entry in sorted(self.runtime_dir.iterdir()):
                config_path = entry / "build_config.py"
                if entry.is_dir() and config_path.is_file():
                    self._runtimes[entry.name] = config_path

        # Create platform-configured compiler
        self._runtime_compiler = RuntimeCompiler.get_instance(platform=platform)

    def list_runtimes(self) -> list:
        """Return names of discovered runtime implementations."""
        return list(self._runtimes.keys())

    def _validate_runtime(self, name: str) -> None:
        if name not in self._runtimes:
            available = ", ".join(self._runtimes.keys()) or "(none)"
            raise ValueError(
                f"Runtime '{name}' is not available for platform '{self.platform}'.\n"
                f"Available runtimes for {self.platform}: {available}\n"
                f"Note: Different platforms may support different runtimes. "
                f"Check {self.runtime_dir} for available implementations."
            )

    def _resolve_target_dirs(self, config_dir: Path, build_config: dict, target: str):
        """Resolve include and source dirs for a target from build_config."""
        cfg = build_config[target]
        include_dirs = [str((config_dir / p).resolve()) for p in cfg["include_dirs"]]
        source_dirs = [str((config_dir / p).resolve()) for p in cfg["source_dirs"]]
        return include_dirs, source_dirs

    def _lookup_binaries(self, name: str, output_dir: Path) -> RuntimeBinaries:
        """Look up pre-built binaries from output_dir.

        Resolves binary names from the compiler's target configs, then
        checks that each file exists.

        Raises:
            FileNotFoundError: If any binary is missing.
        """
        compiler = self._runtime_compiler
        paths = {}
        missing = []
        for target in TARGETS:
            target_obj = getattr(compiler, f"{target}_target")
            binary = output_dir / target_obj.get_binary_name()
            paths[target] = binary
            if not binary.is_file():
                missing.append(str(binary))

        if missing:
            raise FileNotFoundError(
                f"Pre-built runtime binaries not found for '{name}' "
                f"(platform={self.platform}):\n"
                + "\n".join(f"  {m}" for m in missing)
                + "\nRun 'pip install --no-build-isolation .' to compile them."
            )

        # Validate sim_context SO exists for sim platforms
        sim_context_path = self._resolve_sim_context_path()
        if sim_context_path is not None and not sim_context_path.is_file():
            raise FileNotFoundError(
                f"Pre-built libcpu_sim_context.so not found at {sim_context_path}.\n"
                "Run 'pip install --no-build-isolation .' to compile it."
            )

        # Validate libsimpler_log.so exists (built once per arch/variant).
        simpler_log_path = self._resolve_simpler_log_path()
        if not simpler_log_path.is_file():
            raise FileNotFoundError(
                f"Pre-built libsimpler_log.so not found at {simpler_log_path}.\n"
                "Run 'pip install --no-build-isolation .' to compile it."
            )

        # Resolve and validate libsimpler_aicpu_dispatcher.so for onboard
        # platforms. runtime_compiler stages one copy per arch into
        # <LIB_DIR>/<arch>/dispatcher/ (shared across all runtimes); sim
        # platforms have no dispatcher.
        dispatcher_path = self._resolve_dispatcher_path()
        if dispatcher_path is not None and not dispatcher_path.is_file():
            raise FileNotFoundError(
                f"Pre-built libsimpler_aicpu_dispatcher.so not found at {dispatcher_path}.\n"
                "Run 'pip install --no-build-isolation .' to compile it."
            )

        return RuntimeBinaries(
            host_path=paths["host"],
            aicpu_path=paths["aicpu"],
            aicore_path=paths["aicore"],
            simpler_log_path=simpler_log_path,
            sim_context_path=sim_context_path,
            dispatcher_path=dispatcher_path,
        )

    def get_binaries(self, name: str, build: bool = False) -> RuntimeBinaries:
        """Return paths to compiled runtime binaries.

        By default, looks up pre-built binaries from build/lib/. When
        build=True, runs cmake configure + make using persistent build
        directories under build/cache/ for incremental compilation.

        Args:
            name: Name of the runtime implementation (e.g. 'host_build_graph')
            build: If True, compile the runtime before returning paths.
                If False (default), return pre-built binary paths.

        Returns:
            RuntimeBinaries with paths to host, aicpu, and aicore binaries.

        Raises:
            FileNotFoundError: If build=False and pre-built binaries are missing.
        """
        self._validate_runtime(name)

        arch, variant = self._arch, self._variant
        output_dir = self._LIB_DIR / arch / variant / name
        # Per-arch shared destination for libsimpler_aicpu_dispatcher.so. The
        # dispatcher has no runtime-specific code, so all runtimes on a given
        # arch reuse the same SO instead of carrying a copy each (~50 KB × N).
        # None on sim — sim variants have no dispatcher.
        dispatcher_staging_dir = self._LIB_DIR / arch / "dispatcher" if variant != "sim" else None

        if not build:
            return self._lookup_binaries(name, output_dir)

        config_path = self._runtimes[name]
        config_dir = config_path.parent
        build_config = load_build_config(config_path)

        compiler = self._runtime_compiler

        current_commit = _get_git_head(PROJECT_ROOT)

        def _compile_target(target: str) -> Path:
            include_dirs, source_dirs = self._resolve_target_dirs(config_dir, build_config, target)
            # compile() adds a {target}/ subdirectory inside build_dir
            cache_dir = self._CACHE_DIR / arch / variant / name
            cache_dir.mkdir(parents=True, exist_ok=True)

            # File lock to prevent concurrent cmake runs in the same build dir.
            # Each target gets its own lock so host/aicpu/aicore build in parallel,
            # but two processes building the same target are serialized.
            lock_path = cache_dir / f".{target}.lock"
            with open(lock_path, "w") as lock_fd:
                fcntl.flock(lock_fd, fcntl.LOCK_EX)
                _invalidate_cache_if_stale(cache_dir / target, current_commit)
                return compiler.compile(  # type: ignore[return-value]
                    target,
                    include_dirs,
                    source_dirs,
                    build_dir=str(cache_dir),
                    output_dir=output_dir,
                    dispatcher_dest=dispatcher_staging_dir if target == "aicpu" else None,
                )

        logger.info("Compiling AICore, AICPU, Host in parallel...")

        # libsimpler_log.so must finish before the host runtime is built —
        # the host CMake links against it via -lsimpler_log -L<output_dir>.
        simpler_log_path = self.ensure_simpler_log(build=True)

        with ThreadPoolExecutor(max_workers=4) as executor:
            fut_host = executor.submit(_compile_target, "host")
            fut_aicpu = executor.submit(_compile_target, "aicpu")
            fut_aicore = executor.submit(_compile_target, "aicore")
            fut_sim_ctx = executor.submit(self.ensure_sim_context, build=True) if variant == "sim" else None

            host_path = fut_host.result()
            aicpu_path = fut_aicpu.result()
            aicore_path = fut_aicore.result()
            sim_context_path = fut_sim_ctx.result() if fut_sim_ctx else None

        self._place_compile_commands(name)
        logger.info("Build complete!")
        # runtime_compiler stages libsimpler_aicpu_dispatcher.so into the
        # per-arch shared directory when target=='aicpu'. Surface it through
        # RuntimeBinaries so ChipWorker.init can pass the path to
        # LoadAicpuOp::BootstrapDispatcher.
        dispatcher_path = self._resolve_dispatcher_path()
        if dispatcher_path is not None and not dispatcher_path.is_file():
            dispatcher_path = None
        return RuntimeBinaries(
            host_path=host_path,
            aicpu_path=aicpu_path,
            aicore_path=aicore_path,
            simpler_log_path=simpler_log_path,
            sim_context_path=sim_context_path,
            dispatcher_path=dispatcher_path,
        )

    def _resolve_dispatcher_path(self) -> Optional[Path]:
        """Return path to libsimpler_aicpu_dispatcher.so for onboard variants.

        Returns ``None`` for sim variants (no dispatcher needed: sim's AICPU
        runs in-process). For onboard, runtime_compiler stages one shared
        copy per arch under ``build/lib/<arch>/dispatcher/`` (the dispatcher
        has no runtime-specific code, so all onboard runtimes on a given
        arch use the same SO). Validated separately by ``_lookup_binaries``.
        """
        if self._variant == "sim":
            return None
        return self._LIB_DIR / self._arch / "dispatcher" / "libsimpler_aicpu_dispatcher.so"

    def _resolve_sim_context_path(self) -> Optional[Path]:
        """Return path to libcpu_sim_context.so for sim platforms, None for onboard.

        Like libsimpler_log.so, the library is process-global — its source has
        no arch-specific code, so one shared copy per host toolchain is enough.
        Lives at build/lib/libcpu_sim_context.so.
        """
        if self._variant != "sim":
            return None
        return self._LIB_DIR / "libcpu_sim_context.so"

    def _resolve_simpler_log_path(self) -> Path:
        """Return path to libsimpler_log.so.

        Process-global, not arch- or variant-specific — the source is plain
        C++ with no platform conditionals, so one shared copy per host
        toolchain is sufficient. Lives at build/lib/libsimpler_log.so.
        """
        return self._LIB_DIR / "libsimpler_log.so"

    def ensure_simpler_log(self, build: bool = False) -> Path:
        """Build or locate the process-global libsimpler_log.so."""
        output_dir = self._LIB_DIR
        so_path = output_dir / "libsimpler_log.so"

        if not build and so_path.is_file():
            return so_path
        if not build:
            raise FileNotFoundError(
                f"Pre-built libsimpler_log.so not found at {so_path}.\n"
                "Run 'pip install --no-build-isolation .' to compile it."
            )

        cache_dir = self._CACHE_DIR
        cache_dir.mkdir(parents=True, exist_ok=True)
        lock_path = cache_dir / ".simpler_log.lock"
        with open(lock_path, "w") as lock_fd:
            fcntl.flock(lock_fd, fcntl.LOCK_EX)
            result = self._runtime_compiler.compile_simpler_log(
                build_dir=str(cache_dir),
                output_dir=output_dir,
            )
            return Path(result)  # type: ignore[arg-type]

    def ensure_sim_context(self, build: bool = False) -> Optional[Path]:
        """Build or locate the process-global cpu_sim_context SO (sim only)."""
        if self._variant != "sim":
            return None

        output_dir = self._LIB_DIR
        so_path = output_dir / "libcpu_sim_context.so"

        if not build and so_path.is_file():
            return so_path
        if not build:
            raise FileNotFoundError(
                f"Pre-built libcpu_sim_context.so not found at {so_path}.\n"
                "Run 'pip install --no-build-isolation .' to compile it."
            )

        cache_dir = self._CACHE_DIR
        cache_dir.mkdir(parents=True, exist_ok=True)
        lock_path = cache_dir / ".sim_context.lock"
        with open(lock_path, "w") as lock_fd:
            fcntl.flock(lock_fd, fcntl.LOCK_EX)
            result = self._runtime_compiler.compile_sim_context(
                build_dir=str(cache_dir),
                output_dir=output_dir,
            )
            return Path(result)  # type: ignore[arg-type]

    def _place_compile_commands(self, runtime_name: str) -> None:
        """Merge compile_commands.json from build/cache/ targets into source dirs.

        Placement follows the old gen_compile_commands.py defaults:
        - Runtime dirs get compdb only from the preferred variant (onboard),
          so clangd sees real-hardware compile flags.
        - Platform dirs get compdb only from the preferred runtime
          (tensormap_and_ringbuffer), the most feature-rich runtime.

        When the preferred variant/runtime isn't available for this arch,
        falls back to writing unconditionally.
        """
        arch, variant = self._arch, self._variant
        entries = []
        for target in TARGETS:
            cc = self._CACHE_DIR / arch / variant / runtime_name / target / "compile_commands.json"
            if cc.exists():
                try:
                    entries.extend(json.loads(cc.read_text()))
                except (json.JSONDecodeError, OSError):
                    pass

        if not entries:
            return

        merged = json.dumps(entries, indent=2) + "\n"

        # Place in runtime source directory (prefer onboard variant, fallback if unavailable)
        write_runtime = (
            variant == self._COMPDB_VARIANT
            or not (self.runtime_root / "src" / arch / "platform" / self._COMPDB_VARIANT).is_dir()
        )
        if write_runtime:
            runtime_dir = self.runtime_root / "src" / arch / "runtime" / runtime_name
            if runtime_dir.is_dir():
                (runtime_dir / "compile_commands.json").write_text(merged)

        # Place in platform variant source directory (prefer tensormap_and_ringbuffer, fallback if unavailable)
        write_platform = runtime_name == self._COMPDB_RUNTIME or self._COMPDB_RUNTIME not in self._runtimes
        if write_platform:
            platform_dir = self.runtime_root / "src" / arch / "platform" / variant
            if platform_dir.is_dir():
                (platform_dir / "compile_commands.json").write_text(merged)