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..8069dac69 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-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.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..6205299a6 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-5 * (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