Fix PMPP v2 sort reward hack and harden correctness checks

.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

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}

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)

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]

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}

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)

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]

problems/pmpp_v2/sort_py/task.yml

@@ -13,21 +13,32 @@ 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" 
 
 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}