Run all correctness stages. Returns dict with results.
(kernel_fn: Callable, config: dict, quick: bool = False)
| 639 | |
| 640 | |
| 641 | def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> dict: |
| 642 | """Run all correctness stages. Returns dict with results.""" |
| 643 | device = "cuda" |
| 644 | results = { |
| 645 | "smoke_test": "SKIP", |
| 646 | "shape_sweep": "SKIP", |
| 647 | "numerical_stability": "SKIP", |
| 648 | "determinism": "SKIP", |
| 649 | "edge_cases": "SKIP", |
| 650 | "correctness": "FAIL", |
| 651 | } |
| 652 | details = [] |
| 653 | all_pass = True |
| 654 | |
| 655 | gen_fn = config["input_generator"] |
| 656 | ref_fn = config["reference_fn"] |
| 657 | sizes = config["test_sizes"] |
| 658 | dtypes = config["test_dtypes"] |
| 659 | tols = config["tolerances"] |
| 660 | |
| 661 | # ------------------------------------------------------------------ |
| 662 | # Stage 1: SMOKE TEST -- tiny input, tight tolerance |
| 663 | # ------------------------------------------------------------------ |
| 664 | print("\n--- Stage 1: Smoke Test ---") |
| 665 | try: |
| 666 | tiny_label, tiny_size = sizes[0] |
| 667 | # Use first dtype |
| 668 | dtype0 = dtypes[0] |
| 669 | inputs = gen_fn(tiny_size, dtype0, device, seed=42) |
| 670 | expected = ref_fn(inputs) |
| 671 | with _Timeout(30): |
| 672 | output = kernel_fn(**inputs) |
| 673 | |
| 674 | if _has_nan_inf(output): |
| 675 | results["smoke_test"] = "FAIL" |
| 676 | details.append(f" smoke: NaN/Inf in output") |
| 677 | all_pass = False |
| 678 | print(f" FAIL: NaN/Inf in output") |
| 679 | else: |
| 680 | tol = tols.get(dtype0, {"atol": 1e-2, "rtol": 1e-2}) |
| 681 | cmp = _compare(output, expected, **tol) |
| 682 | if cmp["match"]: |
| 683 | results["smoke_test"] = "PASS" |
| 684 | print(f" PASS (max_abs_error={cmp['max_abs_error']:.6e})") |
| 685 | else: |
| 686 | results["smoke_test"] = "FAIL" |
| 687 | details.append(f" smoke: {cmp['reason']}") |
| 688 | all_pass = False |
| 689 | print(f" FAIL: {cmp['reason']}") |
| 690 | except BenchTimeoutError: |
| 691 | results["smoke_test"] = "FAIL" |
| 692 | details.append(" smoke: TIMEOUT") |
| 693 | all_pass = False |
| 694 | print(" FAIL: TIMEOUT") |
| 695 | except torch.cuda.OutOfMemoryError: |
| 696 | results["smoke_test"] = "FAIL" |
| 697 | details.append(" smoke: OOM") |
| 698 | all_pass = False |
no test coverage detected