run_agent.py

#!/usr/bin/env python3

from __future__ import annotations

from pathlib import Path

from rocm_agent import ROCmKernelAgent


def main() -> None:
    repo_root = Path(__file__).resolve().parent
    workdir = repo_root / "agent_workdir"

    agent = ROCmKernelAgent(
        workdir=str(workdir),
        config_ids=[0, 1, 2],
        exploration=1.0,
        warmup=15,
        iters=40,
    )

    summary = agent.run(steps=9)

    print("=== ROCm Agent Run Summary ===")
    print(f"Device: {summary.device.device_name} ({summary.device.gcn_arch_name})")
    print(f"Best Config (fastest): {summary.best_config}")
    print(f"Best Config (by reward): {summary.best_config_by_reward}")
    print(f"Best Reward: {summary.best_reward:.4f}")

    for idx, trial in enumerate(summary.trials, start=1):
        print(
            f"Trial {idx:02d} | arm={trial.config_id} cfg={trial.config} "
            f"| reward={trial.reward:.4f} | "
            f"ok={trial.correctness_passed} | us={trial.candidate_us:.2f} | "
            f"vs_compile={trial.speedup_vs_compile:.3f}x"
        )

    print("Saved config and traces to: agent_workdir/best_config.json")


if __name__ == "__main__":
    main()