()
| 552 | # --------------------------------------------------------------------------- |
| 553 | |
| 554 | def main() -> None: |
| 555 | parser = argparse.ArgumentParser( |
| 556 | description="KernelBench Bridge -- Load and manage KernelBench problems", |
| 557 | ) |
| 558 | sub = parser.add_subparsers(dest="command", help="Command") |
| 559 | |
| 560 | # -- fetch -- |
| 561 | fetch_p = sub.add_parser("fetch", help="Download problems into local cache") |
| 562 | fetch_p.add_argument("--source", choices=["hf", "local", "file"], default="hf") |
| 563 | fetch_p.add_argument("--level", type=int, default=None) |
| 564 | fetch_p.add_argument("--problem", type=int, default=None) |
| 565 | fetch_p.add_argument("--repo-path", type=str, default=None) |
| 566 | fetch_p.add_argument("--file-path", type=str, default=None) |
| 567 | |
| 568 | # -- list -- |
| 569 | list_p = sub.add_parser("list", help="List cached problems") |
| 570 | list_p.add_argument("--level", type=int, default=None) |
| 571 | |
| 572 | # -- info -- |
| 573 | info_p = sub.add_parser("info", help="Show detailed problem info") |
| 574 | info_p.add_argument("--level", type=int, required=True) |
| 575 | info_p.add_argument("--problem", type=int, required=True) |
| 576 | |
| 577 | # -- setup -- |
| 578 | setup_p = sub.add_parser("setup", help="Set up workspace for a problem") |
| 579 | setup_p.add_argument("--level", type=int, required=True) |
| 580 | setup_p.add_argument("--problem", type=int, required=True) |
| 581 | setup_p.add_argument("--backend", choices=["cuda", "triton"], default="cuda") |
| 582 | setup_p.add_argument( |
| 583 | "--source", choices=["hf", "local", "file"], default=None, |
| 584 | help="Auto-fetch from this source if problem not in cache", |
| 585 | ) |
| 586 | setup_p.add_argument("--repo-path", type=str, default=None) |
| 587 | setup_p.add_argument("--file-path", type=str, default=None) |
| 588 | |
| 589 | args = parser.parse_args() |
| 590 | |
| 591 | if args.command == "fetch": |
| 592 | if args.source == "hf": |
| 593 | load_from_huggingface(level=args.level, problem_id=args.problem) |
| 594 | elif args.source == "local": |
| 595 | if not args.repo_path: |
| 596 | print("ERROR: --repo-path required for local source") |
| 597 | sys.exit(1) |
| 598 | load_from_local_repo(args.repo_path, level=args.level, problem_id=args.problem) |
| 599 | elif args.source == "file": |
| 600 | if not args.file_path: |
| 601 | print("ERROR: --file-path required for file source") |
| 602 | sys.exit(1) |
| 603 | load_from_file(args.file_path, level=args.level or 1, problem_id=args.problem or 0) |
| 604 | |
| 605 | elif args.command == "list": |
| 606 | cached = list_cached(level=args.level) |
| 607 | if not cached: |
| 608 | print("No cached problems. Run 'fetch' first:") |
| 609 | print(" uv run kernelbench/bridge.py fetch --source hf --level 1") |
| 610 | return |
| 611 | print(f"{'Level':<7} {'ID':<6} {'Name'}") |
no test coverage detected