run.py

#!/usr/bin/env python3
"""
IntrinTrans LangGraph entry script

Usage:
  # Translate a single function (auto-locate files from VecIntrinBench)
  python run.py --func-name volk_32u_popcnt --model deepseek/deepseek-v3.2
"""

import asyncio
import argparse
import json
import os
import sys
import time
from pathlib import Path
from loguru import logger

PROJECT_ROOT = Path(__file__).parent.absolute()
sys.path.insert(0, str(PROJECT_ROOT))

from intrintrans.config import validate_config, config
from intrintrans.state import make_initial_state, TranslationPhase
from intrintrans.graph import create_graph
from intrintrans.utils.logging_utils import get_model_short_name, setup_single_mode_logging


def setup_logging(log_level: str = "INFO", log_file: str = None):
    """Basic logging setup (used for --check-config and similar scenarios)."""
    logger.remove()
    logger.add(
        sys.stderr,
        level=log_level,
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    )
    if log_file:
        logger.add(log_file, level=log_level, rotation="100 MB")


def _save_versioned_code_files(final_state: dict) -> None:
    """Save each code version that passed compile+test to an individual .cpp file:

      - {func_name}_translated.cpp  <- first translation that passed tests (pre-optimization)
      - {func_name}_opt{N}.cpp      <- N-th optimized version that passed tests (N=1,2,...)
      - {func_name}.cpp             <- final selected best version (no suffix)

    Output directory: {VECINTRINS_BASE_PATH}/src/{build_version}/
    Directory is created automatically if it does not exist.
    """
    func_name = final_state.get("func_name", "")
    build_version = final_state.get("build_version", "")
    vecintrins_base = config.vecintrins_base_path

    if not func_name:
        logger.warning("[SaveVersions] func_name is empty, skipping versioned code save")
        return
    if not vecintrins_base:
        logger.warning("[SaveVersions] VECINTRINS_BASE_PATH not set, skipping versioned code save")
        return

    output_dir = Path(vecintrins_base) / "src" / build_version
    os.makedirs(str(output_dir), exist_ok=True)

    saved_count = 0

    initial_code = final_state.get("initial_translated_code")
    if initial_code:
        translated_path = output_dir / f"{func_name}_translated.cpp"
        translated_path.write_text(initial_code, encoding="utf-8")
        logger.info(f"[SaveVersions] Saved translated version -> {translated_path}")
        saved_count += 1

    all_passed_versions = final_state.get("all_passed_versions", [])
    for i, version in enumerate(all_passed_versions, 1):
        code = version.get("code", "") if isinstance(version, dict) else ""
        if code:
            opt_path = output_dir / f"{func_name}_opt{i}.cpp"
            opt_path.write_text(code, encoding="utf-8")
            logger.info(f"[SaveVersions] Saved opt version {i} -> {opt_path}")
            saved_count += 1

    best_version = final_state.get("best_version")
    if best_version and isinstance(best_version, dict):
        best_code = best_version.get("code", "")
    else:
        best_code = final_state.get("translated_code", "")

    if best_code:
        best_path = output_dir / f"{func_name}.cpp"
        best_path.write_text(best_code, encoding="utf-8")
        logger.info(f"[SaveVersions] Saved best version -> {best_path}")
        saved_count += 1

    if saved_count > 0:
        logger.info(
            f"[SaveVersions] {saved_count} versioned file(s) saved "
            f"(translated={'yes' if initial_code else 'no'}, "
            f"opt_versions={len(all_passed_versions)}, best={'yes' if best_code else 'no'}) "
            f"-> {output_dir}"
        )
    else:
        logger.debug(f"[SaveVersions] No versioned files to save for {func_name}")


def _is_translation_successful(final_state: dict) -> bool:
    """Return True when the workflow produced a usable translated result."""
    phase = final_state.get("phase", TranslationPhase.FAILED)
    if phase != TranslationPhase.DONE:
        return False

    if final_state.get("compile_success") and final_state.get("test_passed"):
        return True

    best_version = final_state.get("best_version")
    return bool(best_version and best_version.get("passed_tests", False))


async def translate_single(
    func_name: str,
    source_code: str,
    test_code: str,
    model: str,
    source_version: str = "arm",
    max_translation_attempts: int = None,
    max_optimization_attempts: int = None,
    perf_test_file_path: str = None,
    enable_checkpointing: bool = False,
) -> dict:
    """Translate a single function.

    Args:
        func_name: Function name
        source_code: Original NEON code
        test_code: Functional test code
        model: LLM model name
        source_version: Source architecture version (default "arm")
        max_translation_attempts: Max translation attempts
        max_optimization_attempts: Max optimization rounds
        perf_test_file_path: Performance test file path
        enable_checkpointing: Whether to enable LangGraph MemorySaver checkpoints

    Returns:
        Dict containing translation results (with final_state key for build_result_json)
    """
    effective_max_opt = max_optimization_attempts or config.max_optimization_attempts

    initial_state = make_initial_state(
        source_code=source_code,
        test_code=test_code,
        func_name=func_name,
        model=model,
        source_version=source_version,
        max_translation_attempts=max_translation_attempts or config.max_translation_attempts,
        max_optimization_attempts=effective_max_opt,
        max_fix_attempts=config.max_fix_attempts,
        perf_test_file_path=perf_test_file_path,
    )

    initial_state["task_start_time"] = time.time()

    graph_factory = create_graph

    if enable_checkpointing:
        try:
            from langgraph.checkpoint.memory import MemorySaver
            checkpointer = MemorySaver()
            graph = graph_factory(checkpointer=checkpointer)
            config_dict = {"configurable": {"thread_id": f"{func_name}-{model}"}}
        except ImportError:
            logger.warning("MemorySaver not available, running without checkpointing")
            graph = graph_factory()
            config_dict = {}
    else:
        graph = graph_factory()
        config_dict = {}

    logger.info(f"Starting translation: func={func_name} model={model}")

    try:
        if config_dict:
            final_state = await graph.ainvoke(initial_state, config=config_dict)
        else:
            final_state = await graph.ainvoke(initial_state)
    except Exception as e:
        logger.error(f"Graph execution failed: {e}")
        return {
            "func_name": func_name,
            "model": model,
            "success": False,
            "error": str(e),
            "phase": "failed",
            "final_state": dict(initial_state),
        }

    phase = final_state.get("phase", TranslationPhase.FAILED)
    success = _is_translation_successful(final_state)

    result = {
        "func_name": func_name,
        "model": model,
        "success": success,
        "phase": str(phase),
        "translation_attempts": final_state.get("translation_attempts", 0),
        "optimization_attempts": final_state.get("optimization_attempts", 0),
        "optimization_iterations": len(final_state.get("all_passed_versions", [])),
    }

    if success:
        best_version = final_state.get("best_version")
        if best_version:
            result["best_code"] = best_version.get("code", "")
            result["speedup"] = best_version.get("speedup")
        else:
            result["best_code"] = final_state.get("translated_code", "")

        logger.info(
            f"Translation succeeded: {func_name} "
            f"(attempts={result['translation_attempts']}, "
            f"opt_versions={result['optimization_iterations']})"
        )
    else:
        compile_success = final_state.get("compile_success", False)
        test_passed = final_state.get("test_passed", False)
        error_reason = final_state.get("error_reason")
        if not error_reason:
            if not compile_success:
                compile_error = final_state.get("compile_error")
                error_reason = compile_error or "Compilation failed (no details)"
            elif not test_passed:
                test_error_msg = final_state.get("test_error")
                error_reason = test_error_msg or "Test failed (no details)"
            else:
                error_reason = "Translation completed but success conditions not met"
        result["error_reason"] = error_reason
        logger.warning(
            f"Translation failed: {func_name} "
            f"(phase={phase}, reason={error_reason[:100]})"
        )

    result["final_state"] = dict(final_state)

    _save_versioned_code_files(dict(final_state))

    return result


def build_result_json(final_state: dict, model: str, func_name: str) -> dict:
    """Build a complete result JSON from the final state."""
    steps = final_state.get("steps", [])

    total_prompt = sum(s.get("prompt_tokens", 0) for s in steps)
    total_completion = sum(s.get("completion_tokens", 0) for s in steps)
    total_reasoning = sum(s.get("reasoning_tokens", 0) for s in steps)
    total_tokens = sum(s.get("total_tokens", 0) for s in steps)
    total_cost = sum(s.get("cost", 0.0) for s in steps)
    total_duration = time.time() - final_state.get("task_start_time", time.time())

    phase = final_state.get("phase", "UNKNOWN")
    if hasattr(phase, "value"):
        phase = phase.value

    success = _is_translation_successful(final_state)

    best_version = final_state.get("best_version")
    final_code = None
    if best_version:
        final_code = best_version.get("code") if isinstance(best_version, dict) else None
    if not final_code:
        final_code = final_state.get("translated_code")

    result = {
        "metadata": {
            "model": model,
            "func_name": func_name,
            "source_version": final_state.get("source_version", "arm"),
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        },
        "steps": steps,
        "summary": {
            "success": success,
            "final_phase": str(phase),
            "compile_success": final_state.get("compile_success", False),
            "test_passed": final_state.get("test_passed", False),
            "error_reason": final_state.get("error_reason"),
            "translation_attempts": final_state.get("translation_attempts", 0),
            "optimization_attempts": final_state.get("optimization_attempts", 0),
            "passed_versions_count": len(final_state.get("all_passed_versions", [])),
            "total_steps": len(steps),
            "total_duration_seconds": round(total_duration, 2),
            "token_usage": {
                "total_prompt_tokens": total_prompt,
                "total_completion_tokens": total_completion,
                "total_reasoning_tokens": total_reasoning,
                "total_tokens": total_tokens,
            },
            "total_cost": round(total_cost, 6),
        },
        "final_code": final_code,
    }

    return result


def save_result_json(result: dict, model: str, func_name: str, result_dir: str = "result"):
    """Save result JSON to result/<model>/<model>-<casename>.json."""
    model_short = get_model_short_name(model)
    result_path = Path(result_dir) / model_short / f"{model_short}-{func_name}.json"
    result_path.parent.mkdir(parents=True, exist_ok=True)

    with open(result_path, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False, default=str)

    logger.info(f"Result saved to: {result_path}")
    return str(result_path)


async def main():
    parser = argparse.ArgumentParser(
        description="IntrinTrans: ARM NEON to RISC-V Vector intrinsic code migration tool (LangGraph edition)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Translate a single function (auto-locate files from VecIntrinBench)
  python run.py --func-name volk_32u_popcnt --model deepseek/deepseek-v3.2
        """,
    )

    parser.add_argument(
        "--func-name", type=str, required=False,
        help="Function name (casename). Files are auto-located from VECINTRINS_BASE_PATH."
    )

    parser.add_argument("--source-version", type=str, default="arm",
                        help="Source architecture version (default: arm)")

    parser.add_argument("--model", type=str, default=None,
                        help="LLM model name (default: read from .env OPENAI_MODEL)")

    parser.add_argument("--max-translation", type=int, default=None,
                        help="Max translation attempts (default: from config)")
    parser.add_argument("--max-optimization", type=int, default=None,
                        help="Max optimization iterations (default: from config)")

    parser.add_argument("--enable-checkpointing", action="store_true",
                        help="Enable LangGraph MemorySaver checkpoints (resumable execution)")

    parser.add_argument("--log-level", type=str, default="INFO",
                        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
                        help="Log level (default: INFO)")

    parser.add_argument("--check-config", action="store_true",
                        help="Validate configuration only, do not run translation")

    args = parser.parse_args()

    setup_logging(args.log_level)

    if args.check_config:
        config.print_config_summary()
        is_valid = validate_config()
        sys.exit(0 if is_valid else 1)

    if not args.func_name:
        parser.error("--func-name is required (or use --check-config to validate configuration)")

    if not validate_config():
        logger.error("Configuration validation failed. Please check your .env file.")
        sys.exit(1)

    model = args.model or config.openai_model
    if not model:
        logger.error("No model specified. Use --model or set OPENAI_MODEL in .env")
        sys.exit(1)

    func_name = args.func_name
    base_path = config.vecintrins_base_path
    if not base_path:
        logger.error("VECINTRINS_BASE_PATH not set in .env; cannot auto-locate source file")
        sys.exit(1)

    source_path = Path(base_path) / "src" / "arm" / f"{func_name}.cpp"
    if not source_path.exists():
        logger.error(f"Source file not found: {source_path}")
        sys.exit(1)

    test_path = Path(base_path) / "test" / f"{func_name}.cpp"
    if not test_path.exists():
        logger.error(f"Test file not found: {test_path}")
        sys.exit(1)

    perf_candidate = Path(base_path) / "perf" / f"{func_name}.cpp"
    perf_file_path = str(perf_candidate) if perf_candidate.exists() else None

    source_code = source_path.read_text(encoding="utf-8")
    test_code = test_path.read_text(encoding="utf-8")

    setup_single_mode_logging(model=model, casename=func_name)

    result = await translate_single(
        func_name=func_name,
        source_code=source_code,
        test_code=test_code,
        model=model,
        source_version=args.source_version,
        max_translation_attempts=args.max_translation,
        max_optimization_attempts=args.max_optimization,
        perf_test_file_path=perf_file_path,
        enable_checkpointing=args.enable_checkpointing,
    )

    final_state = result.get("final_state", {})
    result_json = build_result_json(final_state, model, func_name)
    save_result_json(result_json, model, func_name)

    summary = result_json.get("summary", {})
    token_usage = summary.get("token_usage", {})
    print("\n" + "=" * 60)
    print(f"Translation Result: {func_name}")
    print("=" * 60)
    print(f"  Status:    {'SUCCESS' if result['success'] else 'FAILED'}")
    print(f"  Phase:     {result['phase']}")
    print(f"  Trans Attempts: {result.get('translation_attempts', 'N/A')}")
    print(f"  Opt Iterations: {result.get('optimization_iterations', 0)}")
    print(f"  Duration:  {summary.get('total_duration_seconds', 'N/A')}s")
    print(f"  Tokens:    {token_usage.get('total_tokens', 0)} "
          f"(prompt={token_usage.get('total_prompt_tokens', 0)}, "
          f"completion={token_usage.get('total_completion_tokens', 0)})")
    print(f"  Cost:      ${summary.get('total_cost', 0):.6f}")
    if not result["success"]:
        print(f"  Reason:    {(result.get('error_reason') or 'Unknown')[:200]}")
    print("=" * 60)

    sys.exit(0 if result["success"] else 1)


if __name__ == "__main__":
    asyncio.run(main())