Updates to Problems

problems/pmpp/conv2d_py/reference.py

@@ -1,26 +1,28 @@
 from utils import verbose_allclose
 import torch
 import torch.nn.functional as F
-from task import input_t, output_t, KernelSpec
+from task import input_t, output_t
 
-def ref_kernel(data: input_t, spec: KernelSpec) -> output_t:
+def ref_kernel(data: input_t) -> output_t:
     """
     Reference implementation of 2D convolution using PyTorch.
     Args:
         data: Tuple of (input tensor, kernel tensor)
-        spec: Convolution specifications (stride, padding)
     Returns:
         Output tensor after convolution
     """
     input_tensor, kernel = data
     return F.conv2d(
         input_tensor, 
         kernel,
-        stride=spec.stride,
-        padding=spec.padding
+
+        # No padding and no striding
+        # TODO: Can revisit this in future problems
+        stride=1,
+        padding=0
     )
 
-def generate_input(size: int, kernel_size: 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:
@@ -40,7 +42,7 @@ def generate_input(size: int, kernel_size: int, channels: int, batch: int, seed:
     # 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, kernel_size, kernel_size,
+        channels, channels, kernelsize, kernelsize,
         device='cuda',
         dtype=torch.float32,
         generator=gen
@@ -50,10 +52,9 @@ def generate_input(size: int, kernel_size: int, channels: int, batch: int, seed:
 
 def check_implementation(
     data: input_t,
-    spec: KernelSpec,
     output: output_t,
 ) -> str:
-    expected = ref_kernel(data, spec)
+    expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected, rtol=1e-3, atol=1e-3)
 
     if len(reasons) > 0:

problems/pmpp/conv2d_py/submission.py

@@ -0,0 +1,21 @@
+from task import input_t, output_t
+import torch
+import torch.nn.functional as F
+
+
+def custom_kernel(data: input_t) -> output_t:
+    """
+    Implementation of 2D convolution using PyTorch with no padding and no striding.
+    Args:
+        data: Tuple of (input tensor, kernel tensor)
+        spec: Convolution specifications
+    Returns:
+        Output tensor after convolution
+    """
+    input_tensor, kernel = data
+    return F.conv2d(
+        input_tensor, 
+        kernel,
+        stride=1,
+        padding=0
+    )

problems/pmpp/conv2d_py/task.py

@@ -1,18 +1,13 @@
 from typing import TypedDict, TypeVar, Tuple
 import torch
-from dataclasses import dataclass
 
 input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, torch.Tensor])
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
-@dataclass
-class KernelSpec:
-    stride: int
-    padding: int
 
 class TestSpec(TypedDict):
     size: int
-    kernel_size: int
+    kernelsize: int
     channels: int
     batch: int
-    seed: int 
+    seed: int   

problems/pmpp/conv2d_py/task.yml

@@ -11,21 +11,30 @@ lang: "py"
 
 description: |
   Implement a 2D convolution kernel that matches the reference implementation.
-  The kernel should perform 2D convolution with the given specifications (stride and padding).
+  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 
+  be even numbers with the exception of batch size which can sometimes be 1
+  We assume no padding and striding and instead vary the size of the input and kernel,
+  number of channels, and batch size.
+
+  Input: Tuple of (input_tensor, kernel)
+    - input_tensor: 4D tensor of shape (batch, channels, height, width) with arbitrary values
+    - kernel: 4D tensor of shape (channels, channels, kernelsize, kernelsize) with arbitrary values
+  Output: 4D tensor of shape (batch, channels, height-kernelsize+1, width-kernelsize+1) with convolved values
 
 config:
   main: "eval.py" 
 
 tests:
-  - {"size": 32, "kernel_size": 3, "channels": 16, "batch": 1, "seed": 4242}
-  - {"size": 32, "kernel_size": 5, "channels": 16, "batch": 2, "seed": 5236}
-  - {"size": 64, "kernel_size": 3, "channels": 32, "batch": 1, "seed": 1001}
-  - {"size": 64, "kernel_size": 5, "channels": 32, "batch": 2, "seed": 5531}
-  - {"size": 128, "kernel_size": 3, "channels": 64, "batch": 1, "seed": 9173}
+  - {"size": 32, "kernelsize": 4, "channels": 16, "batch": 1, "seed": 4242}
+  - {"size": 32, "kernelsize": 4, "channels": 16, "batch": 2, "seed": 5236}
+  - {"size": 64, "kernelsize": 4, "channels": 32, "batch": 1, "seed": 1001}
+  - {"size": 64, "kernelsize": 8, "channels": 32, "batch": 2, "seed": 5531}
+  - {"size": 128, "kernelsize": 8, "channels": 64, "batch": 1, "seed": 9173}
 
 benchmarks:
-  - {"size": 128, "kernel_size": 3, "channels": 64, "batch": 4, "seed": 54352}
-  - {"size": 128, "kernel_size": 5, "channels": 64, "batch": 4, "seed": 93246}
-  - {"size": 256, "kernel_size": 3, "channels": 128, "batch": 2, "seed": 6256}
-  - {"size": 256, "kernel_size": 5, "channels": 128, "batch": 2, "seed": 8841}
-  - {"size": 512, "kernel_size": 3, "channels": 256, "batch": 1, "seed": 6252} 
+  - {"size": 128, "kernelsize": 8, "channels": 64, "batch": 4, "seed": 54352}
+  - {"size": 128, "kernelsize": 16, "channels": 64, "batch": 4, "seed": 93246}
+  - {"size": 256, "kernelsize": 16, "channels": 128, "batch": 2, "seed": 6256}
+  - {"size": 256, "kernelsize": 16, "channels": 128, "batch": 2, "seed": 8841}
+  - {"size": 512, "kernelsize": 32, "channels": 256, "batch": 1, "seed": 6252} 

problems/pmpp/grayscale_py/submission.py

@@ -0,0 +1,8 @@
+from task import input_t, output_t
+import torch
+
+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)

problems/pmpp/grayscale_py/task.yml

@@ -9,7 +9,7 @@ lang: "py"
 
 description: |
   Implement an RGB to grayscale conversion kernel that matches the reference implementation.
-  The kernel should convert RGB images to grayscale using the standard coefficients:
+  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
 
   Input: RGB tensor of shape (H, W, 3) with values in [0, 1]
@@ -19,15 +19,16 @@ config:
   main: "eval.py" 
 
 tests:
-  - {"size": 127, "seed": 4242}
-  - {"size": 128, "seed": 5236}
-  - {"size": 129, "seed": 1001}
+
+  - {"size": 128, "seed": 1001}
   - {"size": 256, "seed": 5531}
   - {"size": 512, "seed": 9173}
 
 benchmarks:
-  - {"size": 1024, "seed": 54352}
-  - {"size": 2048, "seed": 93246}
-  - {"size": 4096, "seed": 6256}
-  - {"size": 8192, "seed": 8841}
-  - {"size": 16384, "seed": 6252} 
+
+  - {"size": 512, "seed": 54352}
+  - {"size": 1024, "seed": 93246}
+  - {"size": 2048, "seed": 6256}
+  - {"size": 4096, "seed": 8841}
+  - {"size": 8192, "seed": 6252}
+  - {"size": 16384, "seed": 54352}

problems/pmpp/histogram_py/reference.py

@@ -1,47 +1,62 @@
 from utils import verbose_allclose
 import torch
-from task import input_t, output_t, HistogramSpec
+from task import input_t, output_t
 
-def ref_kernel(data: input_t, spec: HistogramSpec) -> output_t:
+def ref_kernel(data: input_t) -> output_t:
     """
     Reference implementation of histogram using PyTorch.
     Args:
-        data: Input tensor to compute histogram on
-        spec: Histogram specifications (num_bins, min_val, max_val)
+        data: tensor of shape (size,)
     Returns:
         Tensor containing bin counts
     """
-    # Clip values to range
-    clipped = torch.clamp(data, spec.min_val, spec.max_val)
+    # Fixed range [0, 100]
+    min_val, max_val = 0, 100
+
+    # Number of bins is input size / 16
+    num_bins = data.shape[0] // 16
+
+    clipped = torch.clamp(data, min_val, max_val)
 
     # Scale to bin indices
-    bin_width = (spec.max_val - spec.min_val) / spec.num_bins
-    indices = ((clipped - spec.min_val) / bin_width).long()
-    indices = torch.clamp(indices, 0, spec.num_bins - 1)
+    bin_width = (max_val - min_val) / num_bins
+    indices = ((clipped - min_val) / bin_width).long()
+    indices = torch.clamp(indices, 0, num_bins - 1)
 
     # Count values in each bin
-    return torch.bincount(indices, minlength=spec.num_bins).to(torch.float32)
+    return torch.bincount(indices, minlength=num_bins).to(torch.float32)
 
 def generate_input(size: int, seed: int) -> input_t:
     """
-    Generates random input tensor with values roughly in [0, 1].
+    Generates random input tensor for histogram.
+    The number of bins is automatically set to size/16.
+
+    Args:
+        size: Size of the input tensor (must be multiple of 16)
+        seed: Random seed
     Returns:
-        Tensor to compute histogram on
+        The input tensor with values in [0, 100]
     """
     gen = torch.Generator(device='cuda')
     gen.manual_seed(seed)
-    # Generate values with normal distribution for interesting histograms
-    return torch.randn(size, device='cuda', dtype=torch.float32, generator=gen).contiguous()
+
+    # Generate integer values between 0 and 100
+    data = torch.randint(0, 101, (size,), device='cuda', dtype=torch.int32, generator=gen)
+
+    # Convert to float since the histogram implementation expects float input
+    return data.float().contiguous()
 
 def check_implementation(
     data: input_t,
-    spec: HistogramSpec,
     output: output_t,
 ) -> str:
-    expected = ref_kernel(data, spec)
+    """
+    Compare custom implementation's output to the reference output.
+    """
+    expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected)
 
     if len(reasons) > 0:
         return "mismatch found! custom implementation doesn't match reference: " + reasons[0]
 
-    return '' 
+    return ''

problems/pmpp/histogram_py/submission.py

@@ -0,0 +1,27 @@
+import torch
+from task import input_t, output_t
+
+def custom_kernel(data: input_t) -> output_t:
+    """
+    Reference implementation of histogram using PyTorch.
+    Args:
+        data: tensor of shape (size,)
+    Returns:
+        Tensor containing bin counts
+    """
+    # Fixed range [0, 100]
+    min_val, max_val = 0, 100
+
+    # Number of bins is input size / 16
+    num_bins = data.shape[0] // 16
+
+    # Clip values to range
+    clipped = torch.clamp(data, min_val, max_val)
+
+    # Scale to bin indices
+    bin_width = (max_val - min_val) / num_bins
+    indices = ((clipped - min_val) / bin_width).long()
+    indices = torch.clamp(indices, 0, num_bins - 1)
+
+    # Count values in each bin
+    return torch.bincount(indices, minlength=num_bins).to(torch.float32)

problems/pmpp/histogram_py/task.py

@@ -1,16 +1,10 @@
 from typing import TypedDict, TypeVar
 import torch
-from dataclasses import dataclass
 
 input_t = TypeVar("input_t", bound=torch.Tensor)
 output_t = TypeVar("output_t", bound=torch.Tensor)
 
-@dataclass
-class HistogramSpec:
-    num_bins: int
-    min_val: float
-    max_val: float
-
 class TestSpec(TypedDict):
     size: int
-    seed: int 
+    seed: int 
+

problems/pmpp/histogram_py/task.yml

@@ -10,22 +10,24 @@ files:
 lang: "py"
 
 description: |
-  Implement a histogram kernel using CUDA inline function that matches the reference implementation.
-  The kernel should count the number of elements falling into each bin across the specified range.
+  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:
+    - data: a tensor of shape (size,)
 
 config:
-  main: "eval.py" 
+  main: "eval.py"
 
 tests:
-  - {"size": 1023, "seed": 4242}
-  - {"size": 1024, "seed": 5236}
-  - {"size": 1025, "seed": 1001}
-  - {"size": 2048, "seed": 5531}
-  - {"size": 4096, "seed": 9173}
+  - {"size": 5120,   "seed": 9991}
+  - {"size": 7840,   "seed": 2105}
+  - {"size": 30080,  "seed": 9999}
+  - {"size": 100000, "seed": 1212}
 
 benchmarks:
-  - {"size": 8192, "seed": 54352}
-  - {"size": 16384, "seed": 93246}
-  - {"size": 32768, "seed": 6256}
-  - {"size": 65536, "seed": 8841}
-  - {"size": 131072, "seed": 6252} 
+  - {"size": 1310720, "seed": 6252}
+  - {"size": 2621440, "seed": 8841}
+  - {"size": 5242880, "seed": 6252}
+  - {"size": 10485760, "seed": 8841}

problems/pmpp/matmul_py/task.yml

@@ -12,6 +12,7 @@ lang: "py"
 description: |
   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
 
 config:
   main: "eval.py"

problems/pmpp/prefixsum_py/reference.py

@@ -22,14 +22,18 @@ def generate_input(size: int, seed: int) -> input_t:
     gen.manual_seed(seed)
     return torch.randn(size, device='cuda', dtype=torch.float32, generator=gen).contiguous()
 
-def check_implementation(
-    data: input_t,
-    output: output_t,
-) -> str:
+# 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:
     expected = ref_kernel(data)
-    reasons = verbose_allclose(output, expected, rtol=1e-5, atol=1e-5)
 
+    # Then get the size for scaling the tolerance
+    n = data.numel()
+
+    scale_factor = n ** 0.5  # Square root of input size
+    rtol = 1e-5 * scale_factor
+    atol = 1e-5 * scale_factor
+    reasons = verbose_allclose(output, expected, rtol=rtol, atol=atol)
     if len(reasons) > 0:
         return "mismatch found! custom implementation doesn't match reference: " + reasons[0]
-
-    return '' 
+    return ''