From 386a1a4eb5f1839efe6a5d269dd3993dbdd4c9ce Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 5 Nov 2026 02:06:04 -0800 Subject: [PATCH] Reject deferred Tensor subclasses in PMPP v2 results --- problems/pmpp_v2/conv2d_py/task.yml | 3 +++ problems/pmpp_v2/grayscale_py/task.yml | 3 +++ problems/pmpp_v2/histogram_py/task.yml | 3 +++ problems/pmpp_v2/matmul_py/task.yml | 3 +++ problems/pmpp_v2/prefixsum_py/task.yml | 3 +++ problems/pmpp_v2/sort_py/task.yml | 3 +++ problems/pmpp_v2/utils.py | 6 ++++-- problems/pmpp_v2/vectoradd_py/task.yml | 3 +++ problems/pmpp_v2/vectorsum_py/task.yml | 3 +++ tests/pmpp_v2/test_checkers.py | 18 ++++++++++++++++++ 10 files changed, 46 insertions(+), 2 deletions(-) diff --git a/problems/pmpp_v2/conv2d_py/task.yml b/problems/pmpp_v2/conv2d_py/task.yml index 9b44b2b72..26ba48233 100644 --- a/problems/pmpp_v2/conv2d_py/task.yml +++ b/problems/pmpp_v2/conv2d_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + Implement a 2D convolution kernel that matches the reference implementation. The kernel should perform 2D convolution with the given specifications We will benchmark different sizes, kernel sizes, channels and batch sizes but they will all diff --git a/problems/pmpp_v2/grayscale_py/task.yml b/problems/pmpp_v2/grayscale_py/task.yml index d1cbb30ab..842e42ba7 100644 --- a/problems/pmpp_v2/grayscale_py/task.yml +++ b/problems/pmpp_v2/grayscale_py/task.yml @@ -8,6 +8,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + Implement an RGB to grayscale conversion kernel that matches the reference implementation. The kernel should convert square RGB images with even sizes to grayscale using the standard coefficients: Y = 0.2989 R + 0.5870 G + 0.1140 B diff --git a/problems/pmpp_v2/histogram_py/task.yml b/problems/pmpp_v2/histogram_py/task.yml index 3b71d819e..192509ab5 100644 --- a/problems/pmpp_v2/histogram_py/task.yml +++ b/problems/pmpp_v2/histogram_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + Implement a histogram kernel that counts the number of elements falling into each bin across the specified range. 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 diff --git a/problems/pmpp_v2/matmul_py/task.yml b/problems/pmpp_v2/matmul_py/task.yml index 864ba1713..8fa1837c6 100644 --- a/problems/pmpp_v2/matmul_py/task.yml +++ b/problems/pmpp_v2/matmul_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + Implement a custom matmul function that matches the reference implementation. The function should handle a tuple of input tensors and apply matmul The shapes of all outer and inner dimensions of tensors are multiples of 16 diff --git a/problems/pmpp_v2/prefixsum_py/task.yml b/problems/pmpp_v2/prefixsum_py/task.yml index 6205299a6..1c086ce02 100644 --- a/problems/pmpp_v2/prefixsum_py/task.yml +++ b/problems/pmpp_v2/prefixsum_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + 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. Input and output use float32. Correctness is checked against float64 accumulation diff --git a/problems/pmpp_v2/sort_py/task.yml b/problems/pmpp_v2/sort_py/task.yml index d438b9237..029fb86ea 100644 --- a/problems/pmpp_v2/sort_py/task.yml +++ b/problems/pmpp_v2/sort_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + 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. diff --git a/problems/pmpp_v2/utils.py b/problems/pmpp_v2/utils.py index fcaa59454..f163287d9 100644 --- a/problems/pmpp_v2/utils.py +++ b/problems/pmpp_v2/utils.py @@ -28,8 +28,10 @@ def get_device(use_cuda: bool = True) -> torch.device: 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"] + # Subclass hooks can fabricate comparisons or defer the kernel until the + # checker touches its result, moving computation outside the timed region. + if type(received) is not torch.Tensor: + return ["TYPE MISMATCH: output must be a plain, materialized torch.Tensor"] if received.shape != expected.shape: return [ f"SIZE MISMATCH: got {tuple(received.shape)}, expected {tuple(expected.shape)}" diff --git a/problems/pmpp_v2/vectoradd_py/task.yml b/problems/pmpp_v2/vectoradd_py/task.yml index f662f16a1..6d3f0ff6b 100644 --- a/problems/pmpp_v2/vectoradd_py/task.yml +++ b/problems/pmpp_v2/vectoradd_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + Implement a float16 vector addition kernel. Input: tuple(torch.Tensor, torch.Tensor) with tensors of shape (N, N) and type torch.float16. These tensors are from diff --git a/problems/pmpp_v2/vectorsum_py/task.yml b/problems/pmpp_v2/vectorsum_py/task.yml index fc752a83e..107828f7b 100644 --- a/problems/pmpp_v2/vectorsum_py/task.yml +++ b/problems/pmpp_v2/vectorsum_py/task.yml @@ -10,6 +10,9 @@ files: lang: "py" description: | + Return a plain, materialized torch.Tensor. Tensor subclasses and deferred + computation during checking are not allowed; compute the result inside custom_kernel. + Implement a vector sum reduction kernel. This kernel computes the sum of all elements in the input tensor. Input: A tensor of shape `(N,)` with values from a normal distribution with mean 0 and variance 1. diff --git a/tests/pmpp_v2/test_checkers.py b/tests/pmpp_v2/test_checkers.py index 7b5cc2c42..ca6cf484c 100644 --- a/tests/pmpp_v2/test_checkers.py +++ b/tests/pmpp_v2/test_checkers.py @@ -15,6 +15,24 @@ def test_reference_result_passes(problem, load_problem, examples): assert passed, message +@pytest.mark.parametrize("problem", PROBLEMS) +def test_tensor_subclasses_are_rejected_without_dispatch( + problem, load_problem, examples +): + class InterceptedOutput(torch.Tensor): + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + raise AssertionError("checker invoked an untrusted output hook") + + ref = load_problem(problem) + data = examples[problem] + base = ref.ref_kernel(data).clone() + output = torch.Tensor._make_subclass(InterceptedOutput, base, False) + passed, message = ref.check_implementation(data, output) + assert not passed + assert "TYPE MISMATCH" in 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):