main.py

import os
import argparse
import sys
from utils.logger import logger, setup_logger
from utils.llm_client import LLMClient
from env.benchmark import BenchmarkTask
from env.runner import KernelRunner
from env.compiler import check_compilation_environment
from core.memory.tree import Tree
from core.search.policy import EpsilonGreedyPolicy, BestOfKPolicy
from core.agent.plan_agent import PlanAgent
from core.agent.code_agent import CodeAgent
from core.agent.debug_agent import DebugAgent
from core.agent.single_agent import SingleAgent

def parse_args():
    parser = argparse.ArgumentParser(description="STARK: Strategic Team of Agents for Refining Kernels")
    parser.add_argument("--ref_path", type=str, required=True, help="Path to reference PyTorch model file.")
    parser.add_argument("-B", "--budget", type=int, default=30, help="Maximum number of optimization attempts.")
    parser.add_argument("--provider", type=str, default="openai", help="LLM Provider (openai, anthropic, or custom).")
    parser.add_argument("--model", type=str, default="gpt-4o", help="Model name to use for LLMs.")
    parser.add_argument("--api_base", type=str, default=None, help="Custom API Base URL.")
    parser.add_argument("--epsilon", type=float, default=0.3, help="Exploration rate for selection policy.")
    parser.add_argument("--n_root", type=int, default=5, help="Root throttling number.")
    parser.add_argument("--n_child", type=int, default=3, help="Dead-branch pruning number.")
    parser.add_argument("-r", "--leaderboard_size", type=int, default=2, help="Leaderboard size.")
    parser.add_argument("--log_file", type=str, default="stark_run.log", help="Path to save log file.")
    parser.add_argument("--ablation", type=str, choices=["none", "search_agent", "ma_only"], default="none", 
                        help="Ablation mode: 'none' (STARK), 'search_agent' (Single-Agent with search), 'ma_only' (MA with Best-of-K).")
    return parser.parse_args()

def main():
    args = parse_args()
    
    # Setup logging to both stdout and a file
    setup_logger(log_file=args.log_file)
    logger.info("Initializing STARK Optimization framework...")
    logger.info(f"Ablation Mode: {args.ablation}")
    
    # 1. Run compiler environment check
    diagnostics = check_compilation_environment()
    logger.info(f"Compilation environment diagnostics: {diagnostics}")
    
    # 2. Load the benchmark task
    try:
        task = BenchmarkTask(args.ref_path)
    except Exception as e:
        logger.error(f"Failed to load task: {e}")
        sys.exit(1)
        
    # 3. Benchmark the reference PyTorch implementation
    runner = KernelRunner(args.ref_path)
    logger.info("Benchmarking reference implementation to get baseline performance...")
    ref_eval = runner.evaluate(task.get_source_code())
    
    if not ref_eval["ok"] or not ref_eval["correct"]:
        logger.error(f"Failed to run reference model: {ref_eval['runtime_log']}")
        sys.exit(1)
        
    ref_runtime = ref_eval["runtime"]
    logger.info(f"Reference Baseline Runtime: {ref_runtime:.6f} ms")
    
    # 4. Initialize search tree and leaderboard
    tree = Tree(root_kernel_code=task.get_source_code(), root_runtime=ref_runtime)
    
    # Select Policy based on ablation mode
    if args.ablation == "ma_only":
        policy = BestOfKPolicy()
    else:
        policy = EpsilonGreedyPolicy(epsilon=args.epsilon, n_root=args.n_root, n_child=args.n_child)
    
    # 5. Initialize LLM Clients and Agents
    client = LLMClient(provider=args.provider, api_base=args.api_base)
    plan_agent = PlanAgent(client, args.model, temperature=0.8)
    code_agent = CodeAgent(client, args.model, temperature=0.1)
    debug_agent = DebugAgent(client, args.model, temperature=0.1)
    single_agent = SingleAgent(client, args.model, temperature=0.8)
    
    logger.info("Starting STARK main optimization loop...")
    
    # 6. Main Optimization loop (Algorithm 1)
    for t in range(1, args.budget + 1):
        logger.info(f"\n==================== ATTEMPT {t}/{args.budget} ====================")
        
        # Step 4: Select node to refine
        selected_node = policy.select(tree)
        logger.info(f"Selected Node {selected_node.node_id} for refinement (score={selected_node.score})")
        
        if args.ablation == "search_agent":
            # SINGLE AGENT ABLATION (Search Agent)
            # Uses tree memory and epsilon-greedy search, but monolithic single-step generation
            optimized_code = single_agent.optimize(selected_node, tree, r=args.leaderboard_size)
            eval_res = runner.evaluate(optimized_code)
            child_node = tree.add_child(
                parent_id=selected_node.node_id,
                kernel_code=optimized_code,
                plan="SingleAgent Search Ablation",
                anchors=None,
                is_compiled=eval_res["ok"],
                is_correct=eval_res["correct"],
                runtime=eval_res["runtime"],
                compile_log=eval_res["compile_log"],
                runtime_log=eval_res["runtime_log"]
            )
        else:
            # MA-ONLY OR FULL STARK WORKFLOW
            # Step 5: Check if selected node is buggy (compile/correctness failure)
            if selected_node.has_bug:
                # Step 6 & 7: Run DebugAgent
                fixed_code = debug_agent.debug_code(selected_node, tree)
                
                # Step 8: Evaluate
                eval_res = runner.evaluate(fixed_code)
                
                # Step 17: Add child node
                child_node = tree.add_child(
                    parent_id=selected_node.node_id,
                    kernel_code=fixed_code,
                    plan=selected_node.plan,  # inherit plan from parent
                    anchors=selected_node.anchors, # inherit anchors
                    is_compiled=eval_res["ok"],
                    is_correct=eval_res["correct"],
                    runtime=eval_res["runtime"],
                    compile_log=eval_res["compile_log"],
                    runtime_log=eval_res["runtime_log"]
                )
            else:
                # Step 11 & 12: Run PlanAgent to propose optimization
                plan, anchors, ok = plan_agent.propose_plan(selected_node, tree, r=args.leaderboard_size)
                
                if not ok:
                    logger.warning(f"PlanAgent failed to find valid anchors on Node {selected_node.node_id}. Skipping code generation.")
                    continue
                    
                # Step 13 & 14: Run CodeAgent to realize the plan
                optimized_code = code_agent.generate_code(selected_node, tree, plan, anchors)
                
                # Step 15: Evaluate
                eval_res = runner.evaluate(optimized_code)
                
                # Step 17: Add child node
                child_node = tree.add_child(
                    parent_id=selected_node.node_id,
                    kernel_code=optimized_code,
                    plan=plan,
                    anchors=anchors,
                    is_compiled=eval_res["ok"],
                    is_correct=eval_res["correct"],
                    runtime=eval_res["runtime"],
                    compile_log=eval_res["compile_log"],
                    runtime_log=eval_res["runtime_log"]
                )
            
        # Step 18: Update leaderboard C
        tree.update_leaderboard(child_node, r=args.leaderboard_size)
        
        # Log progress
        best_node = tree.get_best_node()
        if best_node.node_id != 0:
            speedup = ref_runtime / best_node.runtime
            logger.info(f"Current best kernel node: {best_node.node_id} (Runtime: {best_node.runtime:.6f} ms, Speedup: {speedup:.2f}x)")
        else:
            logger.info("No customized kernel has outperformed reference PyTorch model yet.")
            
    # 7. Print final results
    logger.info("\n==================== OPTIMIZATION COMPLETE ====================")
    best_node = tree.get_best_node()
    if best_node.node_id == 0:
        logger.warning("STARK did not find any customized kernel that outperforms the baseline or is correct.")
        sys.exit(0)
        
    speedup = ref_runtime / best_node.runtime
    logger.info(f"Optimized Kernel found! Node ID: {best_node.node_id}")
    logger.info(f"Reference Runtime: {ref_runtime:.6f} ms")
    logger.info(f"Optimized Runtime: {best_node.runtime:.6f} ms")
    logger.info(f"Speedup: {speedup:.2f}x")
    
    # Save optimized code
    output_filename = f"optimized_{task.task_name}.py"
    try:
        with open(output_filename, "w", encoding="utf-8") as f:
            f.write(best_node.kernel_code)
        logger.info(f"Saved optimized model to: {os.path.abspath(output_filename)}")
    except Exception as e:
        logger.error(f"Failed to save optimized code: {e}")

if __name__ == "__main__":
    main()