diff --git a/problems/linalg.yaml b/problems/linalg.yaml index 93f243a3..74b52f16 100644 --- a/problems/linalg.yaml +++ b/problems/linalg.yaml @@ -20,3 +20,8 @@ problems: deadline: "2026-07-15" gpus: - B200 + - directory: linalg/eigh_v2 + name: eigh_v2 + deadline: "2026-07-15" + gpus: + - B200 diff --git a/problems/linalg/eigh_v2/eval.py b/problems/linalg/eigh_v2/eval.py new file mode 100644 index 00000000..c811e98a --- /dev/null +++ b/problems/linalg/eigh_v2/eval.py @@ -0,0 +1,350 @@ +import dataclasses +import math +import multiprocessing +import os +import re +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +from torch.cuda.nvtx import range as nvtx_range + +from reference import check_implementation, generate_input +from utils import clear_l2_cache, set_seed + +try: + from task import TestSpec +except ImportError: + TestSpec = dict + + +MAX_ITERATIONS_PER_BENCHMARK = 50 +BENCHMARK_INPUT_BYTES_TARGET = 256 * 1024 * 1024 + + +_ROOFLINE_PEAK_BW_BYTES_PER_S = 8.0e12 +_ROOFLINE_PEAK_FLOP_PER_S = 9.0e15 +_ROOFLINE_LOOSE_FRACTION = 1.0e-3 + + +def _roofline_floor_ns(batch: int, n: int) -> float: + if batch <= 0 or n <= 0: + return 0.0 + bytes_min = (2.0 * batch * n * n + batch * n) * 4.0 + t_bw_ns = bytes_min / _ROOFLINE_PEAK_BW_BYTES_PER_S * 1e9 + flops_min = float(batch) * (float(n) ** 3) + t_flop_ns = flops_min / _ROOFLINE_PEAK_FLOP_PER_S * 1e9 + return _ROOFLINE_LOOSE_FRACTION * max(t_bw_ns, t_flop_ns) + + +class PopcornOutput: + def __init__(self, fd: int): + self.file = os.fdopen(fd, "w") + os.set_inheritable(fd, False) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.file.close() + + def print(self, *args, **kwargs): + print(*args, **kwargs, file=self.file, flush=True) + + def log(self, key, value): + self.print(f"{key}: {value}") + + +@dataclasses.dataclass +class TestCase: + args: dict + spec: str + + +@dataclasses.dataclass +class Stats: + runs: int + mean: float + std: float + err: float + best: float + worst: float + + +def _combine(a: int, b: int) -> int: + return int(a + (a + b) * (a + b + 1) // 2) + + +def get_test_cases(file_name: str, seed: Optional[int]) -> list[TestCase]: + try: + content = Path(file_name).read_text() + except Exception as exc: + print(f"Could not open test file `{file_name}`: {exc}", file=sys.stderr) + exit(113) + + tests = [] + match = r"\s*([a-zA-Z_][a-zA-Z0-9_]*):\s*([a-zA-Z_][a-zA-Z0-9_]*|[+-]?[0-9]+)\s*" + for line in content.splitlines(): + case = {} + for part in line.split(";"): + matched = re.match(match, part) + if not re.fullmatch(match, part): + print(f"invalid test case: '{line}': '{part}'", file=sys.stderr) + exit(113) + key = matched[1] + val = matched[2] + try: + val = int(val) + except ValueError: + pass + case[key] = val + tests.append(TestCase(spec=line, args=case)) + + if seed is not None: + for test in tests: + if "seed" in test.args: + test.args["seed"] = _combine(test.args["seed"], seed) + return tests + + +def calculate_stats(durations: list[float]) -> Stats: + runs = len(durations) + total = sum(durations) + avg = total / runs + variance = sum((x - avg) ** 2 for x in durations) + std = math.sqrt(variance / (runs - 1)) if runs > 1 else 0.0 + err = std / math.sqrt(runs) if runs > 0 else 0.0 + return Stats( + runs=runs, + mean=avg, + std=std, + err=err, + best=float(min(durations)), + worst=float(max(durations)), + ) + + +def _clone_data(data): + if isinstance(data, tuple): + return tuple(_clone_data(x) for x in data) + if isinstance(data, list): + return [_clone_data(x) for x in data] + if isinstance(data, dict): + return {k: _clone_data(v) for k, v in data.items()} + if isinstance(data, torch.Tensor): + return data.clone() + return data + + +def _run_single_test(test: TestCase): + from submission import custom_kernel + + data = generate_input(**test.args) + torch.cuda.synchronize() + output = custom_kernel(_clone_data(data)) + torch.cuda.synchronize() + return check_implementation(data, output) + + +def run_single_test(pool: multiprocessing.Pool, test: TestCase): + return pool.apply(_run_single_test, (test,)) + + +def run_testing(logger: PopcornOutput, pool: multiprocessing.Pool, tests: list[TestCase]): + passed = True + logger.log("test-count", len(tests)) + for idx, test in enumerate(tests): + logger.log(f"test.{idx}.spec", test.spec) + good, message = run_single_test(pool, test) + if good: + logger.log(f"test.{idx}.status", "pass") + if message: + logger.log(f"test.{idx}.message", message) + else: + logger.log(f"test.{idx}.status", "fail") + logger.log(f"test.{idx}.error", message) + passed = False + logger.log("check", "pass" if passed else "fail") + return 0 if passed else 112 + + +def _make_data_batch(test: TestCase, count: int, seed_offset: int = 0): + args = dict(test.args) + if "seed" in args: + args["seed"] += seed_offset + data_list = [] + for _ in range(count): + if "seed" in args: + args["seed"] += 42 + data_list.append(generate_input(**args)) + return data_list + + +def _benchmark_batch_count(test: TestCase) -> int: + batch = int(test.args.get("batch", 1)) + n = int(test.args.get("n", 1)) + bytes_per_input = (batch * n * n) * 4 + if bytes_per_input <= 0: + return 1 + return max(1, min(MAX_ITERATIONS_PER_BENCHMARK, BENCHMARK_INPUT_BYTES_TARGET // bytes_per_input)) + + +def _run_single_benchmark( + test: TestCase, + recheck: bool, + max_repeats: int, + max_time_ns: float, +) -> Stats | Any: + from submission import custom_kernel + + data_list = _make_data_batch(test, _benchmark_batch_count(test)) + check_copy = _clone_data(data_list) + + outputs = [custom_kernel(_clone_data(data)) for data in data_list] + for reference_data, output in zip(check_copy, outputs): + good, message = check_implementation(reference_data, output) + if not good: + return message + + durations = [] + bm_start_time = time.perf_counter_ns() + for i in range(max_repeats): + if recheck: + data_list = _make_data_batch(test, _benchmark_batch_count(test), seed_offset=(i + 1) * 13) + check_copy = _clone_data(data_list) + torch.cuda.synchronize() + clear_l2_cache() + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + outputs = [custom_kernel(data) for data in data_list] + end_event.record() + torch.cuda.synchronize() + durations.append(start_event.elapsed_time(end_event) * 1e6 / len(data_list)) + + if recheck: + for reference_data, output in zip(check_copy, outputs): + good, message = check_implementation(reference_data, output) + if not good: + return message + + total_bm_duration = time.perf_counter_ns() - bm_start_time + if i > 1 and total_bm_duration > 1e8: + stats = calculate_stats(durations) + relative_err = float("inf") if stats.mean == 0 else stats.err / stats.mean + if ( + relative_err < 0.001 + or stats.mean * stats.runs > max_time_ns + or total_bm_duration > 120e9 + ): + break + + stats = calculate_stats(durations) + batch = int(test.args.get("batch", 1)) + n = int(test.args.get("n", 1)) + floor_ns = _roofline_floor_ns(batch, n) + if stats.mean < floor_ns: + return ( + "reported time below physical roofline floor: " + f"batch={batch}, n={n}, reported_mean={stats.mean:.6g} ns " + f"({stats.mean / 1000.0:.6g} us), floor={floor_ns:.6g} ns " + f"({floor_ns / 1000.0:.6g} us)" + ) + + return stats + + +def run_single_benchmark( + pool: multiprocessing.Pool, + test: TestCase, + recheck: bool, + max_repeats: int, + max_time_ns: float, +): + return pool.apply(_run_single_benchmark, (test, recheck, max_repeats, max_time_ns)) + + +def run_benchmarking(logger: PopcornOutput, pool: multiprocessing.Pool, tests: list[TestCase]): + run_single_benchmark(pool, tests[0], False, 200, 10e7) + + passed = True + logger.log("benchmark-count", len(tests)) + for idx, test in enumerate(tests): + logger.log(f"benchmark.{idx}.spec", test.spec) + result = run_single_benchmark(pool, test, True, 200, 10e9) + if isinstance(result, Stats): + for field in dataclasses.fields(Stats): + logger.log(f"benchmark.{idx}.{field.name}", getattr(result, field.name)) + else: + logger.log(f"benchmark.{idx}.status", "fail") + logger.log(f"benchmark.{idx}.error", result) + passed = False + logger.log("check", "pass" if passed else "fail") + return 0 if passed else 112 + + +def _run_single_profile(test: TestCase): + from submission import custom_kernel + + with nvtx_range("generate input"): + data = generate_input(**test.args) + torch.cuda.synchronize() + + cloned = _clone_data(data) + with nvtx_range("custom_kernel"): + output = custom_kernel(cloned) + torch.cuda.synchronize() + + return check_implementation(data, output) + + +def run_single_profile(pool: multiprocessing.Pool, test: TestCase): + return pool.apply(_run_single_profile, (test,)) + + +def run_profiling(logger: PopcornOutput, pool: multiprocessing.Pool, tests: list[TestCase]): + logger.log("benchmark-count", len(tests)) + test = tests[0] + logger.log("benchmark.0.spec", test.spec) + good, message = run_single_profile(pool, test) + if not good: + logger.log("benchmark.0.status", "fail") + logger.log("benchmark.0.error", message) + logger.log("check", "fail") + return 112 + logger.log("check", "pass") + return 0 + + +def main(): + fd = os.getenv("POPCORN_FD") + if not fd: + return 111 + if len(sys.argv) < 3: + return 2 + + mode = sys.argv[1] + seed = os.getenv("POPCORN_SEED") + os.unsetenv("POPCORN_SEED") + seed = int(seed) if seed else None + set_seed(seed or 42) + tests = get_test_cases(sys.argv[2], seed) + + with PopcornOutput(int(fd)) as logger: + mp_context = multiprocessing.get_context("spawn") + with mp_context.Pool(1) as pool: + if mode == "test": + return run_testing(logger, pool, tests) + if mode == "benchmark": + return run_benchmarking(logger, pool, tests) + if mode == "leaderboard": + return run_benchmarking(logger, pool, tests) + if mode == "profile": + return run_profiling(logger, pool, tests) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/problems/linalg/eigh_v2/reference.py b/problems/linalg/eigh_v2/reference.py new file mode 100644 index 00000000..08ec56c8 --- /dev/null +++ b/problems/linalg/eigh_v2/reference.py @@ -0,0 +1,432 @@ +import math + +import torch +from task import input_t, output_t + + +# Intentionally broad, dimension-scaled residual gates. Eigh has sign and +# eigenspace non-uniqueness, and we want to admit reasonable approximate or +# low-bit internal strategies without comparing against reference eigenvectors. +_EIGEN_RTOL_FACTOR = 200.0 +_EIGVAL_RTOL_FACTOR = 200.0 +_ORTH_RTOL_FACTOR = 100.0 +_SORT_RTOL_FACTOR = 100.0 + + +def _as_plain_fp64(value: torch.Tensor) -> torch.Tensor: + plain = torch.Tensor.as_subclass(torch.Tensor.detach(value), torch.Tensor) + return torch.Tensor.double(plain) + + +def _matrix_l1_norm(value: torch.Tensor) -> torch.Tensor: + return torch.linalg.matrix_norm(_as_plain_fp64(value), ord=1, dim=(-2, -1)) + + +def _property_rtol(n: int, factor: float) -> float: + eps = torch.finfo(torch.float32).eps + return factor * max(n, 1) * eps + + +def _scaled_residual( + residual: torch.Tensor, + scale: torch.Tensor, + n: int, +) -> torch.Tensor: + eps = torch.finfo(torch.float32).eps + return residual / (eps * max(n, 1) * scale.clamp_min(1e-30)) + + +def _band_mask(n: int, bandwidth: int, device: torch.device) -> torch.Tensor: + idx = torch.arange(n, device=device) + return (idx[:, None] - idx[None, :]).abs() <= bandwidth + + +def _symmetrize(a: torch.Tensor) -> torch.Tensor: + return 0.5 * (a + a.transpose(-1, -2)) + + +def _signed_logspace(batch: int, n: int, cond: int, device: torch.device) -> torch.Tensor: + span = max(cond, 1) + magnitudes = torch.logspace(-float(span), 0.0, n, device=device, dtype=torch.float32) + signs = torch.ones((n,), device=device, dtype=torch.float32) + signs[::2] = -1.0 + values = magnitudes * signs + return values.expand(batch, n).contiguous() + + +def _random_orthogonal(batch: int, n: int, gen: torch.Generator, device: torch.device) -> torch.Tensor: + x = torch.randn((batch, n, n), device=device, dtype=torch.float32, generator=gen) + q, r = torch.linalg.qr(x) + signs = torch.sign(torch.diagonal(r, dim1=-2, dim2=-1)).clamp(min=0.0).mul(2.0).sub(1.0) + return q * signs.unsqueeze(-2) + + +def _make_from_spectrum(values: torch.Tensor, gen: torch.Generator) -> torch.Tensor: + batch, n = values.shape + q = _random_orthogonal(batch, n, gen, values.device) + a = (q * values.unsqueeze(-2)) @ q.transpose(-1, -2) + return _symmetrize(a).contiguous() + + +def _lapack_scale(itype: int) -> float: + if itype in (6, 11, 14, 17): + return float(torch.finfo(torch.float32).max**0.5) + if itype in (7, 12, 15, 18): + return float(torch.finfo(torch.float32).tiny**0.5) + return 1.0 + + +def _lapack_signed_values( + batch: int, + n: int, + mode: str, + gen: torch.Generator, + device: torch.device, +) -> torch.Tensor: + ulp = 2.0 * torch.finfo(torch.float32).eps + if n == 1: + values = torch.ones((1,), device=device, dtype=torch.float32) + elif mode == "even": + values = torch.linspace(1.0, ulp, n, device=device, dtype=torch.float32) + elif mode == "geometric": + values = torch.logspace(0.0, math.log10(ulp), n, device=device, dtype=torch.float32) + elif mode == "clustered": + values = torch.full((n,), ulp, device=device, dtype=torch.float32) + values[0] = 1.0 + else: + raise ValueError(f"unknown LAPACK spectrum mode: {mode}") + + signs = torch.randint(0, 2, (batch, n), device=device, generator=gen, dtype=torch.int64) + return values.expand(batch, n) * signs.to(torch.float32).mul_(2.0).sub_(1.0) + + +_LAPACK_CASE_TYPES = { + "lapack_zero": 1, + "lapack_identity": 2, + "lapack_diag_even_spectrum": 3, + "lapack_diag_geometric_spectrum": 4, + "lapack_diag_clustered_spectrum": 5, + "lapack_diag_geometric_high_magnitude": 6, + "lapack_diag_geometric_low_magnitude": 7, + "lapack_dense_even_spectrum": 8, + "lapack_dense_geometric_spectrum": 9, + "lapack_dense_clustered_spectrum": 10, + "lapack_dense_even_high_magnitude": 11, + "lapack_dense_even_low_magnitude": 12, + "lapack_random_symmetric": 13, + "lapack_random_symmetric_high_magnitude": 14, + "lapack_random_symmetric_low_magnitude": 15, + "lapack_band_even_spectrum": 16, + "lapack_band_even_high_magnitude": 17, + "lapack_band_even_low_magnitude": 18, +} + + +def _generate_lapack(batch: int, n: int, itype: int, gen: torch.Generator, device: torch.device) -> torch.Tensor: + assert 1 <= itype <= 18, "LAPACK itype must be in [1, 18]" + scale = _lapack_scale(itype) + if itype == 1: + return torch.zeros((batch, n, n), device=device, dtype=torch.float32) + if itype == 2: + return torch.eye(n, device=device, dtype=torch.float32).expand(batch, n, n).clone() * scale + if itype == 3: + return torch.diag_embed(_lapack_signed_values(batch, n, "even", gen, device) * scale) + if itype in (4, 6, 7): + return torch.diag_embed(_lapack_signed_values(batch, n, "geometric", gen, device) * scale) + if itype == 5: + return torch.diag_embed(_lapack_signed_values(batch, n, "clustered", gen, device) * scale) + if itype in (8, 11, 12): + return _make_from_spectrum(_lapack_signed_values(batch, n, "even", gen, device) * scale, gen) + if itype == 9: + return _make_from_spectrum(_lapack_signed_values(batch, n, "geometric", gen, device), gen) + if itype == 10: + return _make_from_spectrum(_lapack_signed_values(batch, n, "clustered", gen, device), gen) + if itype in (13, 14, 15): + a = torch.empty((batch, n, n), device=device, dtype=torch.float32).uniform_(-1.0, 1.0, generator=gen) + return _symmetrize(a) * scale + if itype in (16, 17, 18): + a = _make_from_spectrum(_lapack_signed_values(batch, n, "even", gen, device) * scale, gen) + # DDRVST specifies a symmetric band matrix with eigenvalues. This + # generator bands a planted-spectrum dense matrix, so the final banded + # matrix's spectrum is perturbed; the checker validates the returned + # eigendecomposition of the final FP32 input. + bandwidth = torch.randint(0, n, (batch,), device=device, generator=gen) + idx = torch.arange(n, device=device) + mask = (idx[None, :, None] - idx[None, None, :]).abs() <= bandwidth[:, None, None] + return (a * mask).contiguous() + raise ValueError(f"unknown LAPACK matrix type: {itype}") + + +def _apply_case(a: torch.Tensor, case: str, cond: int, gen: torch.Generator) -> torch.Tensor: + batch, n, _ = a.shape + device = a.device + + if case == "dense": + a = _symmetrize(a) + if cond: + scales = torch.logspace(0.0, -float(cond), n, device=device, dtype=torch.float32) + a = scales.reshape(1, n, 1) * a * scales.reshape(1, 1, n) + elif case == "spectrum": + values = _signed_logspace(batch, n, cond, device) + a = _make_from_spectrum(values, gen) + elif case == "psd": + scales = torch.logspace(0.0, -float(max(cond, 1)), n, device=device, dtype=torch.float32) + g = a * scales.reshape(1, 1, n) + a = (g @ g.transpose(-1, -2)) / float(n) + elif case == "rankdef": + rank = max(1, (3 * n) // 4) + values = torch.zeros((batch, n), device=device, dtype=torch.float32) + values[:, -rank:] = torch.logspace( + -float(max(cond, 1)), 0.0, rank, device=device, dtype=torch.float32 + ) + a = _make_from_spectrum(values, gen) + elif case == "nearrank": + rank = max(1, (3 * n) // 4) + values = torch.empty((batch, n), device=device, dtype=torch.float32) + values[:, : n - rank] = 1.0e-6 * torch.logspace( + -2.0, 0.0, n - rank, device=device, dtype=torch.float32 + ) + values[:, n - rank :] = torch.logspace( + -float(max(cond, 1)), 0.0, rank, device=device, dtype=torch.float32 + ) + a = _make_from_spectrum(values, gen) + elif case == "repeated": + groups = max(1, min(16, n // 8)) + base = torch.linspace(-1.0, 1.0, groups, device=device, dtype=torch.float32) + values = base.repeat_interleave((n + groups - 1) // groups)[:n] + values = values.expand(batch, n).contiguous() + a = _make_from_spectrum(values, gen) + elif case == "clustered": + center = torch.linspace(-1.0, 1.0, n, device=device, dtype=torch.float32) + jitter = torch.linspace(-1.0, 1.0, n, device=device, dtype=torch.float32) + values = center.sign().clamp(min=0.0).mul(2.0).sub(1.0) + 1.0e-5 * jitter + values[n // 3 : 2 * n // 3] = 1.0 + 1.0e-6 * jitter[n // 3 : 2 * n // 3] + values = values.sort().values.expand(batch, n).contiguous() + a = _make_from_spectrum(values, gen) + elif case == "diagonal": + values = _signed_logspace(batch, n, cond, device) + a = torch.diag_embed(values) + elif case == "band": + bandwidth = max(2, min(32, n // 32)) + a = _symmetrize(a) * _band_mask(n, bandwidth, device) + diag_boost = torch.linspace(-1.0, 1.0, n, device=device, dtype=torch.float32) + a.diagonal(dim1=-2, dim2=-1).add_(diag_boost) + elif case == "rowscale": + row_cond = max(cond, 4) + scales = torch.logspace(0.0, -float(row_cond), n, device=device, dtype=torch.float32) + a = scales.reshape(1, n, 1) * _symmetrize(a) * scales.reshape(1, 1, n) + elif case == "blockdiag": + out = torch.zeros_like(a) + profiles = ("dense", "spectrum", "psd", "repeated", "clustered") + max_block = max(1, min(128, n // 2 if n > 1 else 1)) + for b in range(batch): + start = 0 + while start < n: + remaining = n - start + if remaining <= max_block: + block = remaining + else: + block = int(torch.randint(1, max_block + 1, (1,), device=device, generator=gen).item()) + profile = profiles[ + int(torch.randint(0, len(profiles), (1,), device=device, generator=gen).item()) + ] + block_a = a[b : b + 1, start : start + block, start : start + block] + out[b : b + 1, start : start + block, start : start + block] = _apply_case( + block_a, + profile, + cond, + gen, + ) + start += block + a = out + else: + raise ValueError(f"unknown eigh test case: {case}") + + return _symmetrize(a).contiguous() + + +_MIXED_PROFILES = ( + "dense", + "spectrum", + "psd", + "rankdef", + "nearrank", + "repeated", + "clustered", + "band", + "rowscale", +) +_MIXED_WEIGHTS = (6.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0) + + +def _generate_mixed(a: torch.Tensor, cond: int, gen: torch.Generator) -> torch.Tensor: + batch = a.shape[0] + device = a.device + weights = torch.tensor(_MIXED_WEIGHTS, dtype=torch.float32, device=device) + labels = torch.multinomial(weights, batch, replacement=True, generator=gen) + + if batch >= 2: + is_dense = labels == 0 + if not bool(is_dense.any()): + labels[int(torch.randint(0, batch, (1,), device=device, generator=gen))] = 0 + elif bool(is_dense.all()): + pos = int(torch.randint(0, batch, (1,), device=device, generator=gen)) + labels[pos] = int(torch.randint(1, len(_MIXED_PROFILES), (1,), device=device, generator=gen)) + + for k, prof in enumerate(_MIXED_PROFILES): + mask = labels == k + if bool(mask.any()): + a[mask] = _apply_case(a[mask], prof, cond, gen) + return a + + +def generate_input(batch: int, n: int, cond: int, seed: int, case: str = "dense") -> input_t: + assert batch > 0, "batch must be positive" + assert n > 0, "n must be positive" + assert cond >= 0, "cond must be non-negative" + + device = "cuda" if torch.cuda.is_available() else "cpu" + gen = torch.Generator(device=device) + gen.manual_seed(seed) + + case = case.lower() + if case in _LAPACK_CASE_TYPES: + return _generate_lapack(batch, n, _LAPACK_CASE_TYPES[case], gen, torch.device(device)).contiguous() + + a = torch.randn((batch, n, n), device=device, dtype=torch.float32, generator=gen) + if case == "mixed": + return _generate_mixed(a, cond, gen).contiguous() + return _apply_case(a, case, cond, gen).contiguous() + + +def ref_kernel(data: input_t) -> output_t: + values, vectors = torch.linalg.eigh(data) + return vectors, values + + +def _check_tensor(name: str, value: torch.Tensor, shape: tuple[int, ...], device: torch.device) -> str | None: + if type(value) is not torch.Tensor: + if not isinstance(value, torch.Tensor): + return f"{name} must be a torch.Tensor" + return ( + f"{name} must be a plain torch.Tensor, got subclass " + f"{type(value).__module__}.{type(value).__qualname__}" + ) + if value.shape != shape: + return f"{name} shape must be {shape}, got {tuple(value.shape)}" + if value.dtype != torch.float32: + return f"{name} dtype must be torch.float32, got {value.dtype}" + if value.device != device: + return f"{name} must be on {device}, got {value.device}" + if not torch.isfinite(value).all().item(): + return f"{name} contains NaN or Inf" + return None + + +def _check_ascending(values: torch.Tensor, n: int) -> tuple[bool, str]: + if values.shape[-1] <= 1: + return True, "" + diffs = values[..., 1:] - values[..., :-1] + scale = values.abs().amax(dim=-1, keepdim=True).clamp_min(1.0) + allowed = _property_rtol(n, _SORT_RTOL_FACTOR) * scale + failed = diffs < -allowed + if bool(failed.any().item()): + matrix, col = torch.nonzero(failed, as_tuple=False)[0].tolist() + return False, ( + "eigenvalues must be sorted in ascending order: " + f"matrix={matrix}, index={col}, " + f"left={values[matrix, col].item():.3g}, right={values[matrix, col + 1].item():.3g}" + ) + return True, "" + + +def check_implementation(data: input_t, output: output_t) -> tuple[bool, str]: + a = data + batch, n, _ = a.shape + eigen_rtol = _property_rtol(n, _EIGEN_RTOL_FACTOR) + eigval_rtol = _property_rtol(n, _EIGVAL_RTOL_FACTOR) + orth_rtol = _property_rtol(n, _ORTH_RTOL_FACTOR) + + if not isinstance(output, tuple) or len(output) != 2: + return False, "output must be a tuple `(Q, L)`" + + q, values = output + error = _check_tensor("Q", q, (batch, n, n), a.device) + if error is not None: + return False, error + error = _check_tensor("L", values, (batch, n), a.device) + if error is not None: + return False, error + + good, message = _check_ascending(values, n) + if not good: + return False, message + + a_check = _as_plain_fp64(a) + q_check = _as_plain_fp64(q) + values_check = _as_plain_fp64(values) + aq = a_check @ q_check + ql = q_check * values_check.unsqueeze(-2) + if not torch.isfinite(aq).all().item() or not torch.isfinite(ql).all().item(): + return False, "A @ Q or Q @ diag(L) contains NaN or Inf" + + eigen_residual = _matrix_l1_norm(aq - ql) + eigen_scale = _matrix_l1_norm(a_check) + eigen_allowed = eigen_rtol * eigen_scale + eigen_scaled = _scaled_residual(eigen_residual, eigen_scale, n) + if not torch.isfinite(eigen_scaled).all().item(): + return False, "A @ Q - Q @ diag(L) residual produced NaN or Inf" + eigen_failed = eigen_residual > eigen_allowed + if bool(eigen_failed.any().item()): + worst = int(eigen_scaled.argmax().item()) + return False, ( + "A @ Q - Q @ diag(L) is too large: " + f"matrix={worst}, residual={eigen_residual[worst].item():.3g}, " + f"allowed={eigen_allowed[worst].item():.3g}, " + f"scaled={eigen_scaled[worst].item():.3g}" + ) + + eigval_ref = torch.linalg.eigvalsh(a).double() + eigval_residual = torch.linalg.vector_norm(values_check - eigval_ref, ord=float("inf"), dim=-1) + eigval_scale = torch.maximum( + torch.linalg.vector_norm(eigval_ref, ord=float("inf"), dim=-1), + (eigen_scale / max(n, 1)), + ).clamp_min(1.0) + eigval_allowed = eigval_rtol * eigval_scale + eigval_scaled = _scaled_residual(eigval_residual, eigval_scale, n) + if not torch.isfinite(eigval_scaled).all().item(): + return False, "eigenvalue error produced NaN or Inf" + eigval_failed = eigval_residual > eigval_allowed + if bool(eigval_failed.any().item()): + worst = int(eigval_scaled.argmax().item()) + return False, ( + "eigenvalues differ too much from torch.linalg.eigvalsh(A): " + f"matrix={worst}, residual={eigval_residual[worst].item():.3g}, " + f"allowed={eigval_allowed[worst].item():.3g}, " + f"scaled={eigval_scaled[worst].item():.3g}" + ) + + eye = torch.eye(n, device=a.device, dtype=torch.float64).expand(batch, n, n) + qtq = q_check.transpose(-1, -2) @ q_check + if not torch.isfinite(qtq).all().item(): + return False, "Q.T @ Q contains NaN or Inf" + orth_residual = _matrix_l1_norm(qtq - eye).amax() + orth_scale = _matrix_l1_norm(eye).amax() + orth_allowed = orth_rtol * orth_scale + orth_scaled = _scaled_residual(orth_residual, orth_scale, n) + if orth_residual.item() > orth_allowed.item(): + return False, ( + "Q is not orthogonal enough: " + f"residual={orth_residual.item():.3g}, allowed={orth_allowed.item():.3g}, " + f"scaled={orth_scaled.item():.3g}" + ) + + return True, ( + f"eigen_rtol={eigen_rtol:.3g}; " + f"eigval_rtol={eigval_rtol:.3g}; " + f"orth_rtol={orth_rtol:.3g}; " + f"scaled_eigen_residual={eigen_scaled.amax().item():.3g}; " + f"scaled_eigenvalue_residual={eigval_scaled.amax().item():.3g}; " + f"scaled_orthogonality_residual={orth_scaled.item():.3g}; " + f"batch={batch}; n={n}" + ) diff --git a/problems/linalg/eigh_v2/submission.py b/problems/linalg/eigh_v2/submission.py new file mode 100644 index 00000000..19465b3e --- /dev/null +++ b/problems/linalg/eigh_v2/submission.py @@ -0,0 +1,7 @@ +import torch +from task import input_t, output_t + + +def custom_kernel(data: input_t) -> output_t: + values, vectors = torch.linalg.eigh(data) + return vectors, values diff --git a/problems/linalg/eigh_v2/submissions/README.md b/problems/linalg/eigh_v2/submissions/README.md new file mode 100644 index 00000000..d3ca7778 --- /dev/null +++ b/problems/linalg/eigh_v2/submissions/README.md @@ -0,0 +1,18 @@ +# Eigh v2 Submission Attempts + +These are durable smoke-test submissions for the `eigh_v2` problem. They are not +part of the public starter template; they exist so evaluator changes can be +checked against a simple baseline and the fastest structured attempt found so +far. + +- `torch_eigh.py`: direct dense eigensolver baseline. +- `triton_diagonal_fast_path.py`: exact diagonal fast path that uses Triton to + materialize the permuted eigenvector basis, with dense fallback. + +## Local KernelBot Measurements + +Regenerate measurements against the `eigh_v2` leaderboard before quoting +speedups. The retained Triton submission remains a structured smoke test: it +validates that diagonal fast paths can pass correctness, while dense, mixed, +rank-deficient, clustered, row-scaled, block-diagonal, and LAPACK dense spectrum +cases use the dense fallback path. diff --git a/problems/linalg/eigh_v2/submissions/torch_eigh.py b/problems/linalg/eigh_v2/submissions/torch_eigh.py new file mode 100644 index 00000000..19465b3e --- /dev/null +++ b/problems/linalg/eigh_v2/submissions/torch_eigh.py @@ -0,0 +1,7 @@ +import torch +from task import input_t, output_t + + +def custom_kernel(data: input_t) -> output_t: + values, vectors = torch.linalg.eigh(data) + return vectors, values diff --git a/problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py b/problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py new file mode 100644 index 00000000..7c769533 --- /dev/null +++ b/problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py @@ -0,0 +1,48 @@ +import torch +import triton +import triton.language as tl +from task import input_t, output_t + + +@triton.jit +def _write_permuted_eye_kernel( + vectors, + perm, + total: tl.constexpr, + n: tl.constexpr, + block_size: tl.constexpr, +): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + mask = offsets < total + matrix_size: tl.constexpr = n * n + batch = offsets // matrix_size + rem = offsets - batch * matrix_size + row = rem // n + col = rem - row * n + source_row = tl.load(perm + batch * n + col, mask=mask, other=0) + values = row == source_row + tl.store(vectors + offsets, values, mask=mask) + + +def _is_exact_diagonal(data: torch.Tensor) -> bool: + batch, n, _ = data.shape + return bool(torch.count_nonzero(data).item() == batch * n) + + +def _diagonal_eigh(data: torch.Tensor) -> output_t: + values, perm = torch.diagonal(data, dim1=-2, dim2=-1).sort(dim=-1) + batch, n = values.shape + vectors = torch.empty((batch, n, n), device=data.device, dtype=torch.float32) + total = vectors.numel() + block_size = 256 + grid = (triton.cdiv(total, block_size),) + _write_permuted_eye_kernel[grid](vectors, perm, total, n, block_size) + return vectors, values.contiguous() + + +def custom_kernel(data: input_t) -> output_t: + if _is_exact_diagonal(data): + return _diagonal_eigh(data) + + values, vectors = torch.linalg.eigh(data) + return vectors, values diff --git a/problems/linalg/eigh_v2/task.py b/problems/linalg/eigh_v2/task.py new file mode 100644 index 00000000..e0547dcc --- /dev/null +++ b/problems/linalg/eigh_v2/task.py @@ -0,0 +1,13 @@ +import torch +from typing import NotRequired, TypeVar, TypedDict + +input_t = TypeVar("input_t", bound=torch.Tensor) +output_t = TypeVar("output_t", bound=tuple[torch.Tensor, torch.Tensor]) + + +class TestSpec(TypedDict): + batch: int + n: int + cond: int + seed: int + case: NotRequired[str] diff --git a/problems/linalg/eigh_v2/task.yml b/problems/linalg/eigh_v2/task.yml new file mode 100644 index 00000000..64b9bec0 --- /dev/null +++ b/problems/linalg/eigh_v2/task.yml @@ -0,0 +1,142 @@ +# name: eigh_v2 + +files: + - {"name": "submission.py", "source": "@SUBMISSION@"} + - {"name": "task.py", "source": "task.py"} + - {"name": "utils.py", "source": "../../pmpp_v2/utils.py"} + - {"name": "reference.py", "source": "reference.py"} + - {"name": "eval.py", "source": "eval.py"} + +lang: "py" + +description: | + Implement batched real symmetric eigendecomposition. + + Input is `A`, a `batch x n x n` CUDA tensor in `torch.float32`. Every input + matrix is symmetric up to FP32 roundoff. + + Return `(Q, L)` in the same eigenvector convention as `torch.linalg.eigh(A)`: + `Q` is a `batch x n x n` FP32 tensor whose columns are orthonormal + eigenvectors, and `L` is a `batch x n` FP32 tensor of eigenvalues sorted in + ascending order. The checker validates the invariant `A @ Q = Q @ diag(L)`, + checks `L` against `torch.linalg.eigvalsh(A)`, and checks orthogonality of + `Q`. + + Eigenvectors are not unique. Individual signs may flip, and repeated or + tightly clustered eigenvalues may rotate within their eigenspaces. Correctness + is therefore based on matrix identities, not elementwise comparison against a + reference eigensolver. + + This shape set mirrors the optimizer-statistics motivation in `qr_v2`: square + matrices produced from gradient views and second-moment style statistics. + Batched `512 x 512` is the central target, while `1024`, `2048`, and `4096` + cover larger square factors. + + Test and benchmark specs include a `cond` field. In this task `cond` is a + deterministic dynamic-range knob, not an exact requested condition number. + Some cases create spectra spanning `10^-cond` to `1`; others apply row/column + scaling or generate covariance-like positive semidefinite matrices. Stress + cases include rank-deficient, near-rank-deficient, repeated-eigenvalue, + clustered-eigenvalue, diagonal, banded, row-scaled, block-diagonal, and mixed + inputs. + Correctness tests also include LAPACK DDRVST-inspired matrix types from + `TESTING/EIG/ddrvst.f`: zero, identity, diagonal spectra, dense + planted-spectrum matrices, random symmetric matrices, and high- and + low-magnitude banded variants. + + `eigh_v2` keeps the same API as `eigh`, but isolates stricter correctness and + benchmark-integrity checks so the original leaderboard remains unchanged. + + The `mixed` case builds a heterogeneous batch: each matrix is independently + assigned a conditioning profile at a random position in the batch. This + mirrors real optimizer-statistics batches, where per-layer or per-block + factors do not share one numerical structure. The benchmark set includes + multiple distributions at the same important `512 x 512` shape, so choosing + precision solely from public shape IDs is intentionally less effective than + inspecting numerical quality. + + Correctness is residual-gated against the original FP32 input. Low-bit FP16, + FP8, or NVFP4 work is allowed as an internal implementation strategy: + returned factors must still be FP32 and must represent a numerically + meaningful eigendecomposition. Residuals are measured in FP64 to reduce + checker noise, and the gates are intentionally dimension-scaled and + invariant-based rather than elementwise reference comparisons. This leaves + room for approximate low-bit solutions while still rejecting non-orthogonal + factors, unsorted eigenvalues, and outputs that do not represent the input. + The hard gates are the eigen-equation residual `A @ Q - Q @ diag(L)`, + eigenvalue error against `torch.linalg.eigvalsh(A)`, and orthogonality + residual `Q.T @ Q - I`, each applied with dimension-scaled FP32 tolerances. + For a square orthonormal `Q`, the reconstruction identity follows from the + eigen-equation, so v2 skips the redundant reconstruction matmul in the fast + correctness path. + + Among passing submissions, ranking is by runtime using the geometric mean of + benchmark cases. + +config: + main: "eval.py" + +templates: + Python: "submission.py" + +test_timeout: 240 +benchmark_timeout: 420 +ranked_timeout: 720 +ranking_by: "geom" +gpus: + - B200 + +tests: + - {"batch": 20, "n": 32, "cond": 1, "seed": 53124} + - {"batch": 40, "n": 176, "cond": 1, "seed": 3321} + - {"batch": 40, "n": 352, "cond": 1, "seed": 1200} + - {"batch": 16, "n": 512, "cond": 2, "seed": 32523} + - {"batch": 4, "n": 1024, "cond": 2, "seed": 4327} + - {"batch": 2, "n": 2048, "cond": 1, "seed": 224466} + - {"batch": 16, "n": 512, "cond": 4, "seed": 32524, "case": "spectrum"} + - {"batch": 16, "n": 512, "cond": 2, "seed": 32525, "case": "psd"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32526, "case": "rankdef"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32527, "case": "nearrank"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32528, "case": "repeated"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32529, "case": "clustered"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32530, "case": "band"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32531, "case": "rowscale"} + - {"batch": 8, "n": 512, "cond": 2, "seed": 32533, "case": "blockdiag"} + - {"batch": 4, "n": 1024, "cond": 4, "seed": 4328, "case": "spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 4329, "case": "rankdef"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 4330, "case": "clustered"} + - {"batch": 2, "n": 1024, "cond": 2, "seed": 4332, "case": "blockdiag"} + - {"batch": 2, "n": 2048, "cond": 2, "seed": 224467, "case": "mixed"} + - {"batch": 1, "n": 4096, "cond": 1, "seed": 75343, "case": "diagonal"} + - {"batch": 16, "n": 512, "cond": 2, "seed": 32532, "case": "mixed"} + - {"batch": 4, "n": 1024, "cond": 2, "seed": 4331, "case": "mixed"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920001, "case": "lapack_zero"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920002, "case": "lapack_identity"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920003, "case": "lapack_diag_even_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920004, "case": "lapack_diag_geometric_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920005, "case": "lapack_diag_clustered_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920006, "case": "lapack_diag_geometric_high_magnitude"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920007, "case": "lapack_diag_geometric_low_magnitude"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920008, "case": "lapack_dense_even_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920009, "case": "lapack_dense_geometric_spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920010, "case": "lapack_dense_clustered_spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920011, "case": "lapack_dense_even_high_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920012, "case": "lapack_dense_even_low_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920013, "case": "lapack_random_symmetric"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920014, "case": "lapack_random_symmetric_high_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920015, "case": "lapack_random_symmetric_low_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920016, "case": "lapack_band_even_spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920017, "case": "lapack_band_even_high_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920018, "case": "lapack_band_even_low_magnitude"} + +benchmarks: + - {"batch": 20, "n": 32, "cond": 1, "seed": 43214} + - {"batch": 40, "n": 176, "cond": 1, "seed": 423011} + - {"batch": 40, "n": 352, "cond": 1, "seed": 123456} + - {"batch": 640, "n": 512, "cond": 2, "seed": 1029} + - {"batch": 640, "n": 512, "cond": 2, "seed": 770001, "case": "mixed"} + - {"batch": 640, "n": 512, "cond": 0, "seed": 770003, "case": "rankdef"} + - {"batch": 640, "n": 512, "cond": 0, "seed": 770004, "case": "clustered"} + - {"batch": 640, "n": 512, "cond": 4, "seed": 770006, "case": "rowscale"} + - {"batch": 60, "n": 1024, "cond": 2, "seed": 770002, "case": "mixed"} + - {"batch": 60, "n": 1024, "cond": 0, "seed": 780007, "case": "lapack_dense_geometric_spectrum"}