pmpp v2

problems/pmpp/conv2d_py/reference.py

@@ -3,6 +3,7 @@
 import torch.nn.functional as F
 from task import input_t, output_t
 
+
 class DisableCuDNNTF32:
     def __init__(self):
         self.allow_tf32 = torch.backends.cudnn.allow_tf32
@@ -28,45 +29,55 @@ def ref_kernel(data: input_t) -> output_t:
         Output tensor after convolution
     """
     with DisableCuDNNTF32():
-        input_tensor, kernel = data
+        input_tensor, kernel, output = data
         return F.conv2d(
-            input_tensor, 
+            input_tensor,
             kernel,
-
             # No padding and no striding
             # TODO: Can revisit this in future problems
             stride=1,
-            padding=0
+            padding=0,
         )
 
 
-def generate_input(size: int, kernelsize: int, channels: int, batch: int, seed: int) -> input_t:
+def generate_input(
+    size: int, kernelsize: int, channels: int, batch: int, seed: int
+) -> input_t:
     """
     Generates random input and kernel tensors.
     Returns:
         Tuple of (input tensor, kernel tensor)
     """
-    gen = torch.Generator(device='cuda')
+    gen = torch.Generator(device="cuda")
     gen.manual_seed(seed)
-    
+
     # Generate input tensor: [batch, in_channels, height, width]
     input_tensor = torch.randn(
-        batch, channels, size, size,
-        device='cuda', 
-        dtype=torch.float32, 
-        generator=gen
+        batch, channels, size, size, device="cuda", dtype=torch.float32, generator=gen
     ).contiguous()
-    
+
     # Generate kernel tensor: [out_channels, in_channels, kernel_height, kernel_width]
     # Here we use same number of output channels as input channels for simplicity
     kernel = torch.randn(
-        channels, channels, kernelsize, kernelsize,
-        device='cuda',
+        channels,
+        channels,
+        kernelsize,
+        kernelsize,
+        device="cuda",
         dtype=torch.float32,
-        generator=gen
+        generator=gen,
     ).contiguous()
-
-    return (input_tensor, kernel)
+
+    output_tensor = torch.empty(
+        batch,
+        channels,
+        size - kernelsize + 1,
+        size - kernelsize + 1,
+        device="cuda",
+        dtype=torch.float32,
+    )
+
+    return (input_tensor, kernel, output_tensor)
 
 
-check_implementation = make_match_reference(ref_kernel, rtol=1e-3, atol=1e-3)
+check_implementation = make_match_reference(ref_kernel, rtol=1e-3, atol=1e-3)

problems/pmpp/conv2d_py/solutions/correct/ref.py

@@ -4,10 +4,6 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
-    input_tensor, kernel = data
-    return F.conv2d(
-        input_tensor,
-        kernel,
-        stride=1,
-        padding=0
-    )
+    input_tensor, kernel, output = data
+    output = F.conv2d(input_tensor, kernel, stride=1, padding=0)
+    return output

problems/pmpp/conv2d_py/solutions/wrong/empty.py

@@ -5,7 +5,15 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
-    input_tensor, kernel = data
-    return torch.empty((input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2]-kernel.shape[3]+1, input_tensor.shape[3]-kernel.shape[3]+1),
-                       device=kernel.device, dtype=kernel.dtype
+    input_tensor, kernel, output = data
+    output = torch.empty(
+        (
+            input_tensor.shape[0],
+            input_tensor.shape[1],
+            input_tensor.shape[2] - kernel.shape[3] + 1,
+            input_tensor.shape[3] - kernel.shape[3] + 1,
+        ),
+        device=kernel.device,
+        dtype=kernel.dtype,
     )
+    return output

problems/pmpp/conv2d_py/submission.py

@@ -12,13 +12,6 @@ def custom_kernel(data: input_t) -> output_t:
     Returns:
         Output tensor after convolution
     """
-
-    torch.backends.cudnn.allow_tf32 = False
-    torch.backends.cudnn.deterministic = True
-    input_tensor, kernel = data
-    return F.conv2d(
-        input_tensor, 
-        kernel,
-        stride=1,
-        padding=0
-    )
+    input_tensor, kernel, output = data
+    output = F.conv2d(input_tensor, kernel, stride=1, padding=0)
+    return output

problems/pmpp/conv2d_py/task.py

@@ -1,7 +1,7 @@
 from typing import TypedDict, TypeVar, Tuple
 import torch
 
-input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, torch.Tensor])
+input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, torch.Tensor, torch.Tensor])
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
 
@@ -10,4 +10,4 @@ class TestSpec(TypedDict):
     kernelsize: int
     channels: int
     batch: int
-    seed: int   
+    seed: int

problems/pmpp/grayscale_py/reference.py

@@ -7,17 +7,19 @@ def ref_kernel(data: input_t) -> output_t:
     """
     Reference implementation of RGB to grayscale conversion using PyTorch.
     Uses the standard coefficients: Y = 0.2989 R + 0.5870 G + 0.1140 B
-    
+
     Args:
         data: RGB tensor of shape (H, W, 3) with values in [0, 1]
     Returns:
         Grayscale tensor of shape (H, W) with values in [0, 1]
     """
+    data, output = data
     # Standard RGB to Grayscale coefficients
-    weights = torch.tensor([0.2989, 0.5870, 0.1140], 
-                         device=data.device, 
-                         dtype=data.dtype)
-    return torch.sum(data * weights, dim=-1)
+    weights = torch.tensor(
+        [0.2989, 0.5870, 0.1140], device=data.device, dtype=data.dtype
+    )
+    output = torch.sum(data * weights, dim=-1)
+    return output
 
 
 def generate_input(size: int, seed: int) -> input_t:
@@ -26,12 +28,16 @@ def generate_input(size: int, seed: int) -> input_t:
     Returns:
         Tensor of shape (size, size, 3) with values in [0, 1]
     """
-    gen = torch.Generator(device='cuda')
+    gen = torch.Generator(device="cuda")
     gen.manual_seed(seed)
-    return torch.rand(size, size, 3, 
-                     device='cuda', 
-                     dtype=torch.float32, 
-                     generator=gen).contiguous()
+
+    x = torch.rand(
+        size, size, 3, device="cuda", dtype=torch.float32, generator=gen
+    ).contiguous()
+
+    y = torch.empty(size, size, device="cuda", dtype=torch.float32).contiguous()
+
+    return (x, y)
 
 
 check_implementation = make_match_reference(ref_kernel, rtol=1e-4, atol=1e-4)

problems/pmpp/grayscale_py/solutions/correct/ref.py

@@ -3,7 +3,9 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
-    weights = torch.tensor([0.2989, 0.5870, 0.1140], 
-                         device=data.device, 
-                         dtype=data.dtype)
-    return torch.sum(data * weights, dim=-1)
+    data, output = data
+    weights = torch.tensor(
+        [0.2989, 0.5870, 0.1140], device=data.device, dtype=data.dtype
+    )
+    output = torch.sum(data * weights, dim=-1)
+    return output

problems/pmpp/grayscale_py/solutions/wrong/empty.py

@@ -4,4 +4,5 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
+    data, _output = data
     return torch.empty(size=(data.shape[0], data.shape[1]), device=data.device, dtype=data.dtype)

problems/pmpp/grayscale_py/submission.py

@@ -2,7 +2,9 @@
 import torch
 
 def custom_kernel(data: input_t) -> output_t:
+    data, output = data
     weights = torch.tensor([0.2989, 0.5870, 0.1140], 
                          device=data.device, 
                          dtype=data.dtype)
-    return torch.sum(data * weights, dim=-1)
+    output = torch.sum(data * weights, dim=-1)
+    return output

problems/pmpp/grayscale_py/task.py

@@ -1,9 +1,14 @@
 from typing import TypedDict, TypeVar
 import torch
 
-input_t = TypeVar("input_t", bound=torch.Tensor)  # Input will be (H, W, 3) RGB tensor
-output_t = TypeVar("output_t", bound=torch.Tensor)  # Output will be (H, W) grayscale tensor
+input_t = TypeVar(
+    "input_t", bound=tuple[torch.Tensor, torch.Tensor]
+)  # Input is a pair of tensors (input, output) where input is (H, W, 3) RGB tensor and output is (H, W) grayscale tensor
+output_t = TypeVar(
+    "output_t", bound=torch.Tensor
+)  # Output will be (H, W) grayscale tensor
+
 
 class TestSpec(TypedDict):
     size: int  # Size of the square image (H=W)
-    seed: int 
+    seed: int

problems/pmpp/histogram_py/reference.py

@@ -11,8 +11,10 @@ def ref_kernel(data: input_t) -> output_t:
     Returns:
         Tensor containing bin counts
     """
+    data, output = data
     # Count values in each bin
-    return torch.bincount(data, minlength=256)
+    output = torch.bincount(data, minlength=256)
+    return output
 
 
 def generate_input(size: int, contention: float, seed: int) -> input_t:
@@ -37,7 +39,9 @@ def generate_input(size: int, contention: float, seed: int) -> input_t:
     evil_loc = torch.rand((size,), device='cuda', dtype=torch.float32, generator=gen) < (contention / 100.0)
     data[evil_loc] = evil_value
 
-    return data.contiguous()
+    output = torch.empty(256, device='cuda', dtype=torch.int64).contiguous()
+
+    return data.contiguous(), output
 
 
 def check_implementation(data, output):

problems/pmpp/histogram_py/solutions/correct/ref.py

@@ -3,4 +3,6 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
-    return torch.bincount(data, minlength=256)
+    data, output = data
+    output = torch.bincount(data, minlength=256)
+    return output

problems/pmpp/histogram_py/solutions/wrong/empty.py

@@ -4,4 +4,5 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
+    data, _output = data
     return torch.empty(size=(256,), device=data.device, dtype=data.dtype)

problems/pmpp/histogram_py/submission.py

@@ -1,6 +1,7 @@
 import torch
 from task import input_t, output_t
 
+
 def custom_kernel(data: input_t) -> output_t:
     """
     Reference implementation of histogram using PyTorch.
@@ -9,4 +10,7 @@ def custom_kernel(data: input_t) -> output_t:
     Returns:
         Tensor containing bin counts
     """
-    return torch.bincount(data, minlength=256)
+    data, output = data
+    # Compute histogram with 256 bins
+    output = torch.bincount(data, minlength=256)
+    return output

problems/pmpp/histogram_py/task.py

@@ -1,11 +1,11 @@
 from typing import TypedDict, TypeVar
 import torch
 
-input_t = TypeVar("input_t", bound=torch.Tensor)
+input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor])
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
+
 class TestSpec(TypedDict):
     size: int
     seed: int
     contention: int
-

problems/pmpp/matmul_py/reference.py

@@ -10,7 +10,8 @@ def generate_input(m: int, n: int, k: int, seed: int) -> input_t:
     a.uniform_(0, 1, generator=gen)
     b = torch.empty(k, n, device='cuda', dtype=torch.float16)
     b.uniform_(0, 1, generator=gen)
-    return (a, b)
+    c = torch.empty(m, n, device='cuda', dtype=torch.float16)
+    return (a, b, c)
 
 
 def ref_kernel(data: input_t) -> output_t:

problems/pmpp/matmul_py/solutions/correct/ref.py

@@ -3,6 +3,6 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
-    a, b = data
-    return a @ b
-
+    a, b, c = data
+    c = a @ b
+    return c

problems/pmpp/matmul_py/solutions/wrong/low-precision.py

@@ -3,5 +3,6 @@
 
 
 def custom_kernel(data: input_t) -> output_t:
-    a, b = data
-    return (a.to(torch.bfloat16) @ b.to(torch.bfloat16)).to(a.dtype)
+    a, b, c = data
+    c = (a.to(torch.bfloat16) @ b.to(torch.bfloat16)).to(a.dtype)
+    return c

problems/pmpp/matmul_py/submission.py

@@ -1,5 +1,6 @@
 from task import input_t, output_t
 
 def custom_kernel(data: input_t) -> output_t:
-    a, b = data
-    return a @ b
+    a, b, c = data
+    c = a @ b
+    return c

problems/pmpp/matmul_py/task.py

@@ -1,7 +1,7 @@
 import torch
 from typing import TypeVar, TypedDict
 
-input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor])
+input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor, torch.Tensor])
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
 class TestSpec(TypedDict):

problems/pmpp/prefixsum_py/reference.py

@@ -11,7 +11,9 @@ def ref_kernel(data: input_t) -> output_t:
     Returns:
         Tensor containing the inclusive prefix sum
     """
-    return torch.cumsum(data.to(torch.float64), dim=0).to(torch.float64)
+    data, output = data
+    output = torch.cumsum(data.to(torch.float64), dim=0).to(torch.float64)
+    return output
 
 
 def generate_input(size: int, seed: int) -> input_t:
@@ -20,18 +22,22 @@ def generate_input(size: int, seed: int) -> input_t:
     Returns:
         Tensor to compute prefix sum on
     """
-    gen = torch.Generator(device='cuda')
+    gen = torch.Generator(device="cuda")
     gen.manual_seed(seed)
-    return torch.randn(size, device='cuda', dtype=torch.float32, generator=gen).contiguous()
+    x = torch.randn(
+        size, device="cuda", dtype=torch.float32, generator=gen
+    ).contiguous()
+    y = torch.empty(size, device="cuda", 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.numel()
-    
-    scale_factor = n ** 0.5  # Square root of input size
+
+    scale_factor = n**0.5  # Square root of input size
     rtol = 1e-5 * scale_factor
     atol = 1e-5 * scale_factor