()
| 525 | # --------------------------------------------------------------------------- |
| 526 | |
| 527 | def main() -> None: |
| 528 | parser = argparse.ArgumentParser( |
| 529 | description="KernelBench Evaluation -- Correctness + Performance", |
| 530 | ) |
| 531 | parser.add_argument("--quick", action="store_true", |
| 532 | help="Quick mode: 3 trials, 30 timed runs") |
| 533 | parser.add_argument("--correctness-only", action="store_true", |
| 534 | help="Skip performance benchmarking") |
| 535 | parser.add_argument("--n-trials", type=int, default=None, |
| 536 | help=f"Correctness trials (default: {DEFAULT_N_CORRECTNESS})") |
| 537 | parser.add_argument("--n-timed", type=int, default=None, |
| 538 | help=f"Timed iterations (default: {DEFAULT_N_TIMED})") |
| 539 | parser.add_argument("--n-warmup", type=int, default=DEFAULT_N_WARMUP, |
| 540 | help=f"Warmup iterations (default: {DEFAULT_N_WARMUP})") |
| 541 | parser.add_argument("--atol", type=float, default=DEFAULT_ATOL, |
| 542 | help=f"Absolute tolerance (default: {DEFAULT_ATOL})") |
| 543 | parser.add_argument("--rtol", type=float, default=DEFAULT_RTOL, |
| 544 | help=f"Relative tolerance (default: {DEFAULT_RTOL})") |
| 545 | parser.add_argument("--skip-stability", action="store_true", |
| 546 | help="Skip stability test") |
| 547 | parser.add_argument("--skip-determinism", action="store_true", |
| 548 | help="Skip determinism test") |
| 549 | |
| 550 | args = parser.parse_args() |
| 551 | n_trials = args.n_trials or (3 if args.quick else DEFAULT_N_CORRECTNESS) |
| 552 | n_timed = args.n_timed or (30 if args.quick else DEFAULT_N_TIMED) |
| 553 | |
| 554 | import torch |
| 555 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 556 | if device == "cpu": |
| 557 | print("WARNING: No CUDA GPU detected. Results on CPU are not meaningful.") |
| 558 | |
| 559 | meta = load_metadata() |
| 560 | uid = meta.get("uid", "unknown") |
| 561 | name = meta.get("name", "unknown") |
| 562 | level = meta.get("level", "?") |
| 563 | |
| 564 | print("=" * 65) |
| 565 | print(f"KernelBench Evaluation: {uid} -- {name}") |
| 566 | print(f"Level: {level} | Device: {device}") |
| 567 | print("=" * 65) |
| 568 | print() |
| 569 | |
| 570 | # ---- Load reference ---- |
| 571 | print("Loading reference model...") |
| 572 | Model, ref_get_inputs, ref_get_init_inputs = load_reference() |
| 573 | ref_init_args = ref_get_init_inputs() |
| 574 | model_ref = Model(*ref_init_args) |
| 575 | if hasattr(model_ref, "to"): |
| 576 | model_ref = model_ref.to(device) |
| 577 | if hasattr(model_ref, "eval"): |
| 578 | model_ref = model_ref.eval() |
| 579 | |
| 580 | # ---- Load kernel ---- |
| 581 | print("Loading ModelNew from kernel.py...") |
| 582 | ModelNew, kern_get_inputs, kern_get_init_inputs, problem_meta = load_kernel() |
| 583 | |
| 584 | if kern_get_init_inputs is not None: |
no test coverage detected