From 46710f0f30f36d526694997321eed009253430bb Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 5 Nov 2026 01:30:46 -0800 Subject: [PATCH 1/2] Harden PMPP v2 correctness checks against analytic outputs --- .github/workflows/pmpp-v2-checkers.yml | 30 ++++ problems/pmpp_v2/histogram_py/task.yml | 7 +- problems/pmpp_v2/prefixsum_py/reference.py | 47 +++--- problems/pmpp_v2/prefixsum_py/task.py | 4 +- problems/pmpp_v2/prefixsum_py/task.yml | 8 +- problems/pmpp_v2/sort_py/reference.py | 78 +++++----- problems/pmpp_v2/sort_py/task.py | 7 +- problems/pmpp_v2/sort_py/task.yml | 17 ++- problems/pmpp_v2/utils.py | 166 +++++++++++++-------- problems/pmpp_v2/vectoradd_py/reference.py | 4 +- tests/pmpp_v2/conftest.py | 75 ++++++++++ tests/pmpp_v2/test_checkers.py | 160 ++++++++++++++++++++ 12 files changed, 471 insertions(+), 132 deletions(-) create mode 100644 .github/workflows/pmpp-v2-checkers.yml create mode 100644 tests/pmpp_v2/conftest.py create mode 100644 tests/pmpp_v2/test_checkers.py diff --git a/.github/workflows/pmpp-v2-checkers.yml b/.github/workflows/pmpp-v2-checkers.yml new file mode 100644 index 000000000..b965164bf --- /dev/null +++ b/.github/workflows/pmpp-v2-checkers.yml @@ -0,0 +1,30 @@ +name: PMPP v2 checker regression tests + +on: + pull_request: + paths: + - "problems/pmpp_v2/**" + - "tests/pmpp_v2/**" + - ".github/workflows/pmpp-v2-checkers.yml" + push: + branches: [main] + paths: + - "problems/pmpp_v2/**" + - "tests/pmpp_v2/**" + - ".github/workflows/pmpp-v2-checkers.yml" + +permissions: + contents: read + +jobs: + checkers: + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - run: python -m pip install numpy pytest pyyaml + - run: python -m pytest -q tests/pmpp_v2 diff --git a/problems/pmpp_v2/histogram_py/task.yml b/problems/pmpp_v2/histogram_py/task.yml index 419529aba..3b71d819e 100644 --- a/problems/pmpp_v2/histogram_py/task.yml +++ b/problems/pmpp_v2/histogram_py/task.yml @@ -11,8 +11,9 @@ lang: "py" description: | Implement a histogram kernel that counts the number of elements falling into each bin across the specified range. - The minimum and maximum values of the range are fixed to 0 and 100 respectively. - All sizes are multiples of 16 and the number of bins is set to the size of the input tensor divided by 16. + Input values are uint8 integers in [0, 255]. All sizes are multiples of 16. + The output must be a CUDA int64 tensor of shape (256,), containing the exact count + for every bin. Broadcasting a scalar or returning floating-point counts is invalid. Input: - data: a tensor of shape (size,) @@ -29,6 +30,8 @@ tests: - {"size": 30080, "seed": 9999, "contention": 10} - {"size": 30080, "seed": 4254, "contention": 90} - {"size": 100000, "seed": 1212, "contention": 10} + - {"size": 5120, "seed": 3917, "contention": 0} + - {"size": 5120, "seed": 3918, "contention": 100} benchmarks: - {"size": 1310720, "seed": 6252, "contention": 10} diff --git a/problems/pmpp_v2/prefixsum_py/reference.py b/problems/pmpp_v2/prefixsum_py/reference.py index cc25c10a9..150e666bc 100644 --- a/problems/pmpp_v2/prefixsum_py/reference.py +++ b/problems/pmpp_v2/prefixsum_py/reference.py @@ -1,6 +1,6 @@ -from utils import match_reference, DeterministicContext import torch from task import input_t, output_t +from utils import DeterministicContext, match_reference def ref_kernel(data: input_t) -> output_t: @@ -13,33 +13,46 @@ def ref_kernel(data: input_t) -> output_t: """ with DeterministicContext(): data, output = data - output = torch.cumsum(data.to(torch.float64), dim=0).to(torch.float64) + output = torch.cumsum(data.to(torch.float64), dim=0).to(torch.float32) return output -def generate_input(size: int, seed: int) -> input_t: +def generate_input( + size: int, seed: int, distribution: str = "normal", *, device="cuda" +) -> input_t: """ Generates random input tensor. Returns: Tensor to compute prefix sum on """ - gen = torch.Generator(device="cuda") + gen = torch.Generator(device=device) gen.manual_seed(seed) x = torch.randn( - size, device="cuda", dtype=torch.float32, generator=gen + size, device=device, dtype=torch.float32, generator=gen ).contiguous() - y = torch.empty(size, device="cuda", dtype=torch.float32).contiguous() + if distribution == "positive": + x = x.abs() + 0.25 + elif distribution == "negative": + x = -x.abs() - 0.25 + elif distribution == "alternating": + x[1::2] = -x[::2][: x[1::2].numel()] + elif distribution == "sparse": + x.zero_() + if size: + indices = torch.randperm(size, device=device, generator=gen)[ + : min(size, 16) + ] + x[indices] = torch.randn(indices.numel(), device=device, generator=gen) + elif distribution != "normal": + raise ValueError(f"Unknown prefix sum distribution: {distribution}") + y = torch.empty(size, device=device, dtype=torch.float32).contiguous() return x, y -# This algorithm is very sensitive to the tolerance and the error is magnified by the input size -# The tolerance is scaled by the square root of the input size -def check_implementation(data: input_t, output: output_t) -> str: - # Then get the size for scaling the tolerance - n = data[0].numel() - - scale_factor = n ** 0.5 # Square root of input size - rtol = 1e-5 * scale_factor - atol = 1e-5 * scale_factor - - return match_reference(data, output, reference=ref_kernel, rtol=rtol, atol=atol) +def check_implementation(data: input_t, output: output_t) -> tuple[bool, str]: + # Parallel reduction orders need an absolute roundoff allowance. Scale it + # with input energy, not answer magnitude: sqrt(n)-scaled rtol allowed a + # 16% error on the largest benchmark. + input_norm = torch.linalg.vector_norm(data[0].to(torch.float64)).item() + atol = 1e-6 * (1 + input_norm) + return match_reference(data, output, reference=ref_kernel, rtol=1e-5, atol=atol) diff --git a/problems/pmpp_v2/prefixsum_py/task.py b/problems/pmpp_v2/prefixsum_py/task.py index 79a29e868..af911c6da 100644 --- a/problems/pmpp_v2/prefixsum_py/task.py +++ b/problems/pmpp_v2/prefixsum_py/task.py @@ -1,4 +1,5 @@ -from typing import TypedDict, TypeVar +from typing import NotRequired, TypedDict, TypeVar + import torch input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor]) @@ -8,3 +9,4 @@ class TestSpec(TypedDict): size: int seed: int + distribution: NotRequired[str] diff --git a/problems/pmpp_v2/prefixsum_py/task.yml b/problems/pmpp_v2/prefixsum_py/task.yml index 734546d3f..6aa100be2 100644 --- a/problems/pmpp_v2/prefixsum_py/task.yml +++ b/problems/pmpp_v2/prefixsum_py/task.yml @@ -12,7 +12,9 @@ lang: "py" description: | Implement an inclusive prefix sum (scan) kernel that matches the reference implementation. The kernel should compute the cumulative sum of all elements up to each position. - Because of numerical instability, the tolerance is scaled by the square root of the input size. + Input and output use float32. Correctness is checked against float64 accumulation + rounded to float32, with fixed relative tolerance 1e-5 and an absolute roundoff + allowance of 1e-6 * (1 + L2 norm of the input). Relative error does not grow with size. Input: - `data`: A 1D tensor of size `n` @@ -31,6 +33,10 @@ tests: - {"size": 1025, "seed": 1001} - {"size": 2048, "seed": 5531} - {"size": 4096, "seed": 9173} + - {"size": 4097, "seed": 4001, "distribution": "positive"} + - {"size": 4097, "seed": 4002, "distribution": "negative"} + - {"size": 4097, "seed": 4003, "distribution": "alternating"} + - {"size": 4097, "seed": 4004, "distribution": "sparse"} benchmarks: # - {"size": 8192, "seed": 54352} diff --git a/problems/pmpp_v2/sort_py/reference.py b/problems/pmpp_v2/sort_py/reference.py index ca1ab273c..2d6e39ab9 100644 --- a/problems/pmpp_v2/sort_py/reference.py +++ b/problems/pmpp_v2/sort_py/reference.py @@ -1,6 +1,6 @@ -from utils import make_match_reference, DeterministicContext import torch from task import input_t, output_t +from utils import DeterministicContext, make_match_reference_exact def ref_kernel(data: input_t) -> output_t: @@ -17,43 +17,45 @@ def ref_kernel(data: input_t) -> output_t: return output -def generate_input(size: int, seed: int) -> torch.Tensor: - """ - Generates random input tensor where elements are drawn from different distributions. - - Args: - size: Total size of the final 1D tensor - seed: Base seed for random generation +def generate_input( + size: int, seed: int, distribution: str = "normal", *, device="cuda" +) -> input_t: + """Seeds select random values, never their magnitude or location. - Returns: - 1D tensor of size `size` containing flattened values from different distributions + Private seeds can be very large. Adding them to float32 data erases the + random values and makes approximate, input-independent outputs pass. """ - # Calculate dimensions for a roughly square 2D matrix - rows = int(size**0.5) # Square root for roughly square shape - cols = ( - size + rows - 1 - ) // rows # Ceiling division to ensure total size >= requested size - - gen = torch.Generator(device="cuda") - result = torch.empty((rows, cols), device="cuda", dtype=torch.float32) - - # Different seed for each row! - for i in range(rows): - row_seed = seed + i - gen.manual_seed(row_seed) - - # Generate values for this row with mean=row_seed - result[i, :] = ( - torch.randn(cols, device="cuda", dtype=torch.float32, generator=gen) - + row_seed + if size < 0: + raise ValueError("size must be non-negative") + gen = torch.Generator(device=device).manual_seed(seed) + if distribution in ("normal", "sorted", "reverse", "outliers"): + values = torch.randn(size, device=device, dtype=torch.float32, generator=gen) + if distribution == "sorted": + values = values.sort().values + elif distribution == "reverse": + values = values.sort(descending=True).values + elif distribution == "outliers" and size: + indices = torch.randperm(size, device=device, generator=gen)[ + : max(1, size // 32) + ] + values[indices] *= 1024 + elif distribution == "uniform": + values = ( + torch.rand(size, device=device, dtype=torch.float32, generator=gen) * 16 - 8 ) - - # Flatten and trim to exact size requested - input_tensor = result.flatten()[:size].contiguous() - output_tensor = torch.empty_like( - input_tensor, device="cuda", dtype=torch.float32 - ).contiguous() - return input_tensor, output_tensor - - -check_implementation = make_match_reference(ref_kernel) + elif distribution == "duplicates": + values = torch.randint(-8, 9, (size,), device=device, generator=gen).to( + torch.float32 + ) + elif distribution == "constant": + value = torch.randint(-8, 9, (), device=device, generator=gen).item() + values = torch.full((size,), value, device=device, dtype=torch.float32) + else: + raise ValueError(f"Unknown sort distribution: {distribution}") + values = values.contiguous() + return values, torch.empty_like(values) + + +# Sorting preserves values exactly. Tolerance accepts fabricated values and +# lost duplicates even when the output is increasing and looks plausible. +check_implementation = make_match_reference_exact(ref_kernel) diff --git a/problems/pmpp_v2/sort_py/task.py b/problems/pmpp_v2/sort_py/task.py index 495e681c7..43b8a8942 100644 --- a/problems/pmpp_v2/sort_py/task.py +++ b/problems/pmpp_v2/sort_py/task.py @@ -1,9 +1,12 @@ -from typing import TypedDict, TypeVar +from typing import NotRequired, TypedDict, TypeVar + import torch input_t = TypeVar("input_t", bound=[torch.Tensor, torch.Tensor]) output_t = TypeVar("output_t", bound=torch.Tensor) + class TestSpec(TypedDict): size: int - seed: int + seed: int + distribution: NotRequired[str] diff --git a/problems/pmpp_v2/sort_py/task.yml b/problems/pmpp_v2/sort_py/task.yml index 7e78a156c..d438b9237 100644 --- a/problems/pmpp_v2/sort_py/task.yml +++ b/problems/pmpp_v2/sort_py/task.yml @@ -13,9 +13,12 @@ description: | Implement a sort kernel that matches the reference implementation. The kernel should sort the input array in ascending order using a sort algorithm of your choice. - Input arrays are generated as random floating-point numbers, where each row of a roughly square matrix - is drawn from a normal distribution with a different mean value per row based on the seed and then flattened into a 1D array. - + Input and output are contiguous 1D CUDA float32 tensors of the same length. + Inputs contain finite values from normal, uniform, duplicate-heavy, constant, + already sorted, reverse-sorted, and outlier distributions. Seeds control randomness + only; their numeric values are never added to the data. + The output must equal the sorted input exactly, including duplicate counts. + Generating a similar distribution or approximating sorted values is incorrect. config: main: "eval.py" @@ -23,11 +26,19 @@ config: templates: Python: "../template.py" tests: + - {"size": 1, "seed": 7123} + - {"size": 17, "seed": 9173, "distribution": "duplicates"} - {"size": 1023, "seed": 4242} - {"size": 1024, "seed": 5236} - {"size": 1025, "seed": 1001} - {"size": 2048, "seed": 5531} - {"size": 4096, "seed": 9173} + - {"size": 4097, "seed": 4001, "distribution": "uniform"} + - {"size": 4097, "seed": 4002, "distribution": "duplicates"} + - {"size": 4097, "seed": 4003, "distribution": "constant"} + - {"size": 4097, "seed": 4004, "distribution": "sorted"} + - {"size": 4097, "seed": 4005, "distribution": "reverse"} + - {"size": 4097, "seed": 4006, "distribution": "outliers"} benchmarks: - {"size": 100000, "seed": 54352} diff --git a/problems/pmpp_v2/utils.py b/problems/pmpp_v2/utils.py index e8a9082f5..fcaa59454 100644 --- a/problems/pmpp_v2/utils.py +++ b/problems/pmpp_v2/utils.py @@ -1,5 +1,6 @@ import os import random + import numpy as np import torch @@ -25,14 +26,57 @@ def get_device(use_cuda: bool = True) -> torch.device: return torch.device("cpu") +def _check_tensor_contract(received, expected) -> list[str]: + """Reject malformed results before comparisons can broadcast or promote them.""" + if not isinstance(received, torch.Tensor): + return ["TYPE MISMATCH: output must be a torch.Tensor"] + if received.shape != expected.shape: + return [ + f"SIZE MISMATCH: got {tuple(received.shape)}, expected {tuple(expected.shape)}" + ] + if received.dtype != expected.dtype: + return [f"DTYPE MISMATCH: got {received.dtype}, expected {expected.dtype}"] + if received.device != expected.device: + return [f"DEVICE MISMATCH: got {received.device}, expected {expected.device}"] + if expected.is_floating_point(): + if not torch.isfinite(expected).all().item(): + return ["NON-FINITE REFERENCE: PMPP v2 requires finite results"] + if not torch.isfinite(received).all().item(): + return ["NON-FINITE OUTPUT"] + return [] + + +def _mismatch_details(received, expected, mismatched, max_print): + num_mismatched = mismatched.count_nonzero().item() + if not num_mismatched: + return [] + details = [f"Number of mismatched elements: {num_mismatched}"] + # A fabricated large result can mismatch everywhere. Never allocate an + # index tensor proportional to all failures just to report five examples. + flat = mismatched.reshape(-1) + printed = 0 + for start in range(0, flat.numel(), 65536): + indices = flat[start : start + 65536].nonzero().flatten()[: max_print - printed] + for offset in indices.tolist(): + position = start + offset + index = [] + for size in reversed(expected.shape): + index.append(position % size) + position //= size + index = tuple(reversed(index)) + details.append(f"ERROR AT {index}: {received[index]} {expected[index]}") + printed += 1 + if printed >= max_print: + break + if num_mismatched > printed: + details.append(f"... and {num_mismatched - printed} more mismatched elements.") + return details + + # Adapted from https://github.com/linkedin/Liger-Kernel/blob/main/test/utils.py @torch.no_grad() def verbose_allclose( - received: torch.Tensor, - expected: torch.Tensor, - rtol=1e-05, - atol=1e-08, - max_print=5 + received: torch.Tensor, expected: torch.Tensor, rtol=1e-05, atol=1e-08, max_print=5 ) -> list[str]: """ Assert that two tensors are element-wise equal within a tolerance, providing detailed information about mismatches. @@ -47,9 +91,9 @@ def verbose_allclose( Raises: AssertionError: If the tensors are not all close within the given tolerance. """ - # Check if the shapes of the tensors match - if received.shape != expected.shape: - return ["SIZE MISMATCH"] + reasons = _check_tensor_contract(received, expected) + if reasons: + return reasons # Calculate the difference between the tensors diff = torch.abs(received - expected) @@ -57,44 +101,17 @@ def verbose_allclose( # Determine the tolerance tolerance = atol + rtol * torch.abs(expected) - # Find tolerance mismatched elements - tol_mismatched = diff > tolerance - - # Find nan mismatched elements - nan_mismatched = torch.logical_xor(torch.isnan(received), torch.isnan(expected)) - - # Find +inf mismatched elements - posinf_mismatched = torch.logical_xor(torch.isposinf(received), torch.isposinf(expected)) - # Find -inf mismatched elements - neginf_mismatched = torch.logical_xor(torch.isneginf(received), torch.isneginf(expected)) - - # Find all mismatched elements - mismatched = torch.logical_or( - torch.logical_or(tol_mismatched, nan_mismatched), - torch.logical_or(posinf_mismatched, neginf_mismatched), - ) - - mismatched_indices = torch.nonzero(mismatched) - - # Count the number of mismatched elements - num_mismatched = mismatched.count_nonzero().item() - - # Generate detailed information if there are mismatches - if num_mismatched >= 1: - mismatch_details = [f"Number of mismatched elements: {num_mismatched}"] + # Non-finite values were rejected above. Keep only one mismatch mask for + # benchmark-sized scans instead of several full-size NaN/Inf masks. + mismatched = diff > tolerance - for index in mismatched_indices[:max_print]: - i = tuple(index.tolist()) - mismatch_details.append(f"ERROR AT {i}: {received[i]} {expected[i]}") - if num_mismatched > max_print: - mismatch_details.append(f"... and {num_mismatched - max_print} more mismatched elements.") - return mismatch_details - - return [] + return _mismatch_details(received, expected, mismatched, max_print) @torch.no_grad() -def verbose_allequal(received: torch.Tensor, expected: torch.Tensor, max_print: int=5): +def verbose_allequal( + received: torch.Tensor, expected: torch.Tensor, max_print: int = 5 +): """ Assert that two tensors are element-wise perfectly equal, providing detailed information about mismatches. @@ -106,27 +123,16 @@ def verbose_allequal(received: torch.Tensor, expected: torch.Tensor, max_print: Returns: Empty string if tensors are equal, otherwise detailed error information """ + reasons = _check_tensor_contract(received, expected) + if reasons: + return reasons mismatched = torch.not_equal(received, expected) - mismatched_indices = torch.nonzero(mismatched) - - # Count the number of mismatched elements - num_mismatched = mismatched.count_nonzero().item() - - # Generate detailed information if there are mismatches - if num_mismatched >= 1: - mismatch_details = [f"Number of mismatched elements: {num_mismatched}"] - - for index in mismatched_indices[:max_print]: - i = tuple(index.tolist()) - mismatch_details.append(f"ERROR AT {i}: {received[i]} {expected[i]}") - if num_mismatched > max_print: - mismatch_details.append(f"... and {num_mismatched - max_print} more mismatched elements.") - return mismatch_details - - return [] + return _mismatch_details(received, expected, mismatched, max_print) -def match_reference(data, output, reference: callable, rtol=1e-05, atol=1e-08) -> tuple[bool, str]: +def match_reference( + data, output, reference: callable, rtol=1e-05, atol=1e-08 +) -> tuple[bool, str]: """ Convenient "default" implementation for tasks' `check_implementation` function. """ @@ -134,14 +140,33 @@ def match_reference(data, output, reference: callable, rtol=1e-05, atol=1e-08) - reasons = verbose_allclose(output, expected, rtol=rtol, atol=atol) if len(reasons) > 0: - return False, "mismatch found! custom implementation doesn't match reference: " + " ".join(reasons) + return ( + False, + "mismatch found! custom implementation doesn't match reference: " + + " ".join(reasons), + ) - return True, '' + return True, "" def make_match_reference(reference: callable, **kwargs): def wrapped(data, output): return match_reference(data, output, reference=reference, **kwargs) + + return wrapped + + +def make_match_reference_exact(reference: callable): + """Use for operations such as sorting that must preserve exact values.""" + + @torch.no_grad() + def wrapped(data, output): + expected = reference(data) + reasons = verbose_allequal(output, expected) + if reasons: + return False, "exact reference mismatch: " + " ".join(reasons) + return True, "" + return wrapped @@ -150,21 +175,30 @@ def __init__(self): self.allow_tf32 = None self.deterministic = None self.cublas = None + self.algorithms = None + self.warn_only = None + self.matmul_tf32 = None def __enter__(self): - self.cublas = os.environ.get('CUBLAS_WORKSPACE_CONFIG', '') + self.cublas = os.environ.get("CUBLAS_WORKSPACE_CONFIG", "") self.allow_tf32 = torch.backends.cudnn.allow_tf32 self.deterministic = torch.backends.cudnn.deterministic + self.algorithms = torch.are_deterministic_algorithms_enabled() + self.warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + self.matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 torch.backends.cudnn.allow_tf32 = False torch.backends.cudnn.deterministic = True + torch.backends.cuda.matmul.allow_tf32 = False torch.use_deterministic_algorithms(True) return self def __exit__(self, exc_type, exc_value, traceback): torch.backends.cudnn.allow_tf32 = self.allow_tf32 torch.backends.cudnn.deterministic = self.deterministic - torch.use_deterministic_algorithms(False) - os.environ['CUBLAS_WORKSPACE_CONFIG'] = self.cublas + torch.backends.cuda.matmul.allow_tf32 = self.matmul_tf32 + torch.use_deterministic_algorithms(self.algorithms, warn_only=self.warn_only) + os.environ["CUBLAS_WORKSPACE_CONFIG"] = self.cublas + def clear_l2_cache(): # import cupy as cp @@ -173,4 +207,4 @@ def clear_l2_cache(): dummy = torch.empty((32, 1024, 1024), dtype=torch.int64, device="cuda") # write stuff to dummy.fill_(42) - del dummy \ No newline at end of file + del dummy diff --git a/problems/pmpp_v2/vectoradd_py/reference.py b/problems/pmpp_v2/vectoradd_py/reference.py index 9789711b3..3f200545c 100644 --- a/problems/pmpp_v2/vectoradd_py/reference.py +++ b/problems/pmpp_v2/vectoradd_py/reference.py @@ -1,6 +1,6 @@ -from utils import make_match_reference, DeterministicContext import torch from task import input_t, output_t +from utils import DeterministicContext, make_match_reference_exact def ref_kernel(data: input_t) -> output_t: @@ -35,4 +35,4 @@ def generate_input(size: int, seed: int) -> input_t: return A, B, C -check_implementation = make_match_reference(ref_kernel) +check_implementation = make_match_reference_exact(ref_kernel) diff --git a/tests/pmpp_v2/conftest.py b/tests/pmpp_v2/conftest.py new file mode 100644 index 000000000..d4b90aa50 --- /dev/null +++ b/tests/pmpp_v2/conftest.py @@ -0,0 +1,75 @@ +import importlib.util +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch + +ROOT = Path(__file__).resolve().parents[2] / "problems" / "pmpp_v2" +PROBLEMS = ( + "sort", + "prefixsum", + "histogram", + "vectoradd", + "vectorsum", + "grayscale", + "matmul", + "conv2d", +) + + +def load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def load_problem(): + utils = load_module("pmpp_test_utils", ROOT / "utils.py") + + def load(name): + folder = ROOT / f"{name}_py" + task = load_module(f"pmpp_test_{name}_task", folder / "task.py") + with patch.dict(sys.modules, {"task": task, "utils": utils}): + return load_module(f"pmpp_test_{name}_reference", folder / "reference.py") + + return load + + +@pytest.fixture +def examples(): + gen = torch.Generator().manual_seed(72) + return { + "sort": (torch.tensor([3.0, -2.0, 1.0, 1.0, 0.0]), torch.empty(5)), + "prefixsum": (torch.tensor([1e6, 1.0, -1e6, 2.0, -3.0, 4.0]), torch.empty(6)), + "histogram": ( + torch.arange(256, dtype=torch.uint8).repeat(4), + torch.empty(256, dtype=torch.int64), + ), + "vectoradd": ( + torch.randn(4, 4, dtype=torch.float16, generator=gen), + torch.randn(4, 4, dtype=torch.float16, generator=gen), + torch.empty(4, 4, dtype=torch.float16), + ), + "vectorsum": (torch.tensor([1e6, 1.0, -1e6, 2.0]), torch.empty(1)), + "grayscale": ( + torch.tensor( + [[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], [[0.0, 0.0, 1.0], [1.0, 1.0, 1.0]]] + ), + torch.empty(2, 2), + ), + "matmul": ( + torch.randn(16, 16, dtype=torch.float16, generator=gen), + torch.randn(16, 16, dtype=torch.float16, generator=gen), + torch.empty(16, 16, dtype=torch.float16), + ), + "conv2d": ( + torch.randn(1, 2, 8, 8, generator=gen), + torch.randn(2, 2, 4, 4, generator=gen), + torch.empty(1, 2, 5, 5), + ), + } diff --git a/tests/pmpp_v2/test_checkers.py b/tests/pmpp_v2/test_checkers.py new file mode 100644 index 000000000..7b5cc2c42 --- /dev/null +++ b/tests/pmpp_v2/test_checkers.py @@ -0,0 +1,160 @@ +import math + +import pytest +import torch +import yaml +from conftest import PROBLEMS, ROOT + + +@pytest.mark.parametrize("problem", PROBLEMS) +def test_reference_result_passes(problem, load_problem, examples): + ref = load_problem(problem) + data = examples[problem] + output = ref.ref_kernel(data).clone() + passed, message = ref.check_implementation(data, output) + assert passed, message + + +@pytest.mark.parametrize("problem", PROBLEMS) +@pytest.mark.parametrize("bad_result", ("shape", "dtype", "type", "device")) +def test_malformed_results_are_rejected(problem, bad_result, load_problem, examples): + ref = load_problem(problem) + data = examples[problem] + output = ref.ref_kernel(data).clone() + if bad_result == "shape": + output = output.unsqueeze(0) + elif bad_result == "dtype": + output = output.to( + torch.float64 if output.dtype != torch.float64 else torch.float32 + ) + elif bad_result == "type": + output = output.tolist() + else: + output = torch.empty_like(output, device="meta") + passed, message = ref.check_implementation(data, output) + assert not passed + assert "MISMATCH" in message + + +@pytest.mark.parametrize("problem", tuple(p for p in PROBLEMS if p != "histogram")) +@pytest.mark.parametrize("value", (float("nan"), float("inf"), -float("inf"))) +def test_nonfinite_results_are_rejected(problem, value, load_problem, examples): + ref = load_problem(problem) + data = examples[problem] + output = torch.full_like(ref.ref_kernel(data), value) + assert not ref.check_implementation(data, output)[0] + + +def test_histogram_does_not_accept_broadcast_scalar(load_problem, examples): + # Every bin is four. The old comparator broadcast this scalar and accepted it. + ref = load_problem("histogram") + assert not ref.check_implementation(examples["histogram"], torch.tensor(4))[0] + + +@pytest.mark.parametrize("problem", PROBLEMS) +def test_input_independent_constant_is_rejected(problem, load_problem, examples): + ref = load_problem(problem) + data = examples[problem] + output = torch.full_like(ref.ref_kernel(data), 17) + assert not ref.check_implementation(data, output)[0] + + +def test_sort_rejects_large_seed_analytic_sequence(load_problem): + # Old seed-as-mean generation collapses float32 input to one value. Its rtol + # then accepts a monotone sequence of invented values millions apart. + ref = load_problem("sort") + seed, size = 10**12, 4096 + original = torch.full((size,), seed, dtype=torch.float32) + fabricated = torch.linspace(seed - 4_000_000, seed + 4_000_000, size) + assert not torch.equal(fabricated, original) + assert not ref.check_implementation( + (original, torch.empty_like(original)), fabricated + )[0] + + +def test_sort_rejects_one_ulp_changes_and_lost_duplicates(load_problem): + ref = load_problem("sort") + values = torch.tensor([-1.0, 0.0, 1.0, 1.0, 2.0]) + expected = values.sort().values + changed = torch.nextafter(expected, torch.full_like(expected, math.inf)) + assert not ref.check_implementation((values, torch.empty_like(values)), changed)[0] + lost_duplicate = torch.tensor([-1.0, 0.0, 1.0, 2.0, 2.0]) + assert not ref.check_implementation( + (values, torch.empty_like(values)), lost_duplicate + )[0] + + +@pytest.mark.parametrize( + "distribution", + ("normal", "uniform", "duplicates", "constant", "sorted", "reverse", "outliers"), +) +@pytest.mark.parametrize("seed", (42, 10**12, 2**62 + 7123)) +def test_sort_generators_keep_entropy_and_sort_correctly( + distribution, seed, load_problem +): + ref = load_problem("sort") + data = ref.generate_input(1025, seed, distribution, device="cpu") + repeated = ref.generate_input(1025, seed, distribution, device="cpu") + assert torch.equal(data[0], repeated[0]) + assert data[0].dtype == torch.float32 + assert torch.isfinite(data[0]).all() + assert data[0].abs().max() < 10_000 + if distribution != "constant": + assert data[0].unique().numel() > 1 + expected = torch.sort(data[0]).values + assert ref.check_implementation(data, expected)[0] + + +@pytest.mark.parametrize( + "distribution", ("normal", "positive", "negative", "alternating", "sparse") +) +@pytest.mark.parametrize("size", (1, 1023, 4097, 262144)) +def test_parallel_float32_prefix_sum_passes(distribution, size, load_problem): + ref = load_problem("prefixsum") + data = ref.generate_input(size, 7821, distribution, device="cpu") + assert ref.check_implementation(data, data[0].cumsum(0))[0] + + +def test_prefix_sum_rejects_percentage_error(load_problem): + ref = load_problem("prefixsum") + values = torch.ones(262144) + forged = values.cumsum(0) * 1.001 + assert not ref.check_implementation((values, torch.empty_like(values)), forged)[0] + + +def test_prefix_sum_tolerance_tracks_input_energy(load_problem): + ref = load_problem("prefixsum") + values = torch.zeros(262144) + forged = torch.full_like(values, 0.001) + assert not ref.check_implementation((values, torch.empty_like(values)), forged)[0] + + +@pytest.mark.parametrize("problem", ("sort", "prefixsum")) +def test_all_configured_cases_generate_and_reference_passes(problem, load_problem): + ref = load_problem(problem) + task = yaml.safe_load((ROOT / f"{problem}_py" / "task.yml").read_text()) + for case in task["tests"]: + data = ref.generate_input(**case, device="cpu") + expected = ref.ref_kernel(data).clone() + assert ref.check_implementation(data, expected)[0], case + + +def test_determinism_settings_are_restored_on_exception(load_problem): + ref = load_problem("sort") + before = ( + torch.are_deterministic_algorithms_enabled(), + torch.is_deterministic_algorithms_warn_only_enabled(), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.deterministic, + torch.backends.cudnn.allow_tf32, + ) + with pytest.raises(RuntimeError), ref.DeterministicContext(): + raise RuntimeError("deliberate failure") + after = ( + torch.are_deterministic_algorithms_enabled(), + torch.is_deterministic_algorithms_warn_only_enabled(), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.deterministic, + torch.backends.cudnn.allow_tf32, + ) + assert after == before From 8c72634af7e8d8ae330890699f4d8582d14a904e Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 5 Nov 2026 01:37:46 -0800 Subject: [PATCH 2/2] Allow observed A100 scan roundoff without growing relative tolerance --- problems/pmpp_v2/prefixsum_py/reference.py | 2 +- problems/pmpp_v2/prefixsum_py/task.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/problems/pmpp_v2/prefixsum_py/reference.py b/problems/pmpp_v2/prefixsum_py/reference.py index 150e666bc..8069dac69 100644 --- a/problems/pmpp_v2/prefixsum_py/reference.py +++ b/problems/pmpp_v2/prefixsum_py/reference.py @@ -54,5 +54,5 @@ def check_implementation(data: input_t, output: output_t) -> tuple[bool, str]: # with input energy, not answer magnitude: sqrt(n)-scaled rtol allowed a # 16% error on the largest benchmark. input_norm = torch.linalg.vector_norm(data[0].to(torch.float64)).item() - atol = 1e-6 * (1 + input_norm) + atol = 1e-5 * (1 + input_norm) return match_reference(data, output, reference=ref_kernel, rtol=1e-5, atol=atol) diff --git a/problems/pmpp_v2/prefixsum_py/task.yml b/problems/pmpp_v2/prefixsum_py/task.yml index 6aa100be2..6205299a6 100644 --- a/problems/pmpp_v2/prefixsum_py/task.yml +++ b/problems/pmpp_v2/prefixsum_py/task.yml @@ -14,7 +14,7 @@ description: | The kernel should compute the cumulative sum of all elements up to each position. Input and output use float32. Correctness is checked against float64 accumulation rounded to float32, with fixed relative tolerance 1e-5 and an absolute roundoff - allowance of 1e-6 * (1 + L2 norm of the input). Relative error does not grow with size. + allowance of 1e-5 * (1 + L2 norm of the input). Relative error does not grow with size. Input: - `data`: A 1D tensor of size `n`