submission 551260

- import torch
- import torch.nn as nn
- from torch.utils.cpp_extension import load_inline
-  
-  
- _depthwise_conv_cpp = r"""
- torch::Tensor depthwise_conv2d_kh1_cuda(
-     torch::Tensor x,
-     torch::Tensor w,
-     torch::Tensor b,
-     int64_t stride_h,
-     int64_t stride_w,
-     int64_t pad_h,
-     int64_t pad_w,
-     int64_t dil_h,
-     int64_t dil_w
- );
- """
-  
-  
- _depthwise_conv_cu = r"""
- #include <torch/extension.h>
- #include <cuda.h>
- #include <cuda_runtime.h>
- #include <vector>
-  
- #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
- #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
- #define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be float32")
-  
- __global__ void depthwise_conv2d_kh1_kernel(
-     const float* __restrict__ x,
-     const float* __restrict__ w,
-     const float* __restrict__ b,
-     float* __restrict__ out,
-     int N, int C, int H, int W,
-     int K_h,
-     int outH, int outW,
-     int stride_h, int stride_w,
-     int pad_h, int pad_w,
-     int dil_h, int dil_w,
-     bool has_bias
- ) {
-     int idx = blockIdx.x * blockDim.x + threadIdx.x;
-     int total = N * C * outH * outW;
-     if (idx >= total) return;
-  
-     int ow = idx % outW;
-     int t1 = idx / outW;
-     int oh = t1 % outH;
-     int t2 = t1 / outH;
-     int c = t2 % C;
-     int n = t2 / C;
-  
-     float acc = has_bias ? b[c] : 0.0f;
-  
-     int ih_base = oh * stride_h - pad_h;
-     int iw = ow * stride_w - pad_w;
-  
-     if ((unsigned)iw < (unsigned)W) {
-         const float* x_base = x + ((n * C + c) * H) * W;
-         const float* w_base = w + c * K_h;
-  
-         #pragma unroll
-         for (int kh = 0; kh < K_h; ++kh) {
-             int ih = ih_base + kh * dil_h;
-             if ((unsigned)ih < (unsigned)H) {
-                 acc += x_base[ih * W + iw] * w_base[kh];
-             }
-         }
-     }
-  
-     out[idx] = acc;
- }
-  
- torch::Tensor depthwise_conv2d_kh1_cuda(
-     torch::Tensor x,
-     torch::Tensor w,
-     torch::Tensor b,
-     int64_t stride_h,
-     int64_t stride_w,
-     int64_t pad_h,
-     int64_t pad_w,
-     int64_t dil_h,
-     int64_t dil_w
- ) {
-     CHECK_CUDA(x);
-     CHECK_CUDA(w);
-     CHECK_CONTIGUOUS(x);
-     CHECK_CONTIGUOUS(w);
-     CHECK_FLOAT(x);
-     CHECK_FLOAT(w);
-     TORCH_CHECK(x.dim() == 4, "x must be 4D NCHW");
-     TORCH_CHECK(w.dim() == 4, "w must be 4D [C,1,Kh,1]");
-  
-     const auto N = (int)x.size(0);
-     const auto C = (int)x.size(1);
-     const auto H = (int)x.size(2);
-     const auto W = (int)x.size(3);
-  
-     TORCH_CHECK((int)w.size(0) == C, "w.size(0) must equal input channels");
-     TORCH_CHECK((int)w.size(1) == 1, "depthwise weight second dim must be 1");
-     const auto K_h = (int)w.size(2);
-     TORCH_CHECK((int)w.size(3) == 1, "This kernel supports only kw=1");
-  
-     bool has_bias = b.defined() && b.numel() > 0;
-     if (has_bias) {
-         CHECK_CUDA(b);
-         CHECK_CONTIGUOUS(b);
-         CHECK_FLOAT(b);
-         TORCH_CHECK(b.dim() == 1 && (int)b.size(0) == C, "bias must be shape [C]");
-     }
-  
-     const int outH = (H + 2 * (int)pad_h - (int)dil_h * (K_h - 1) - 1) / (int)stride_h + 1;
-     const int outW = (W + 2 * (int)pad_w - (int)dil_w * (1 - 1) - 1) / (int)stride_w + 1;
-     TORCH_CHECK(outH >= 0 && outW >= 0, "Invalid output shape");
-  
-     auto out = torch::empty({N, C, outH, outW}, x.options());
-  
-     const int total = N * C * outH * outW;
-     const int threads = 256;
-     const int blocks = (total + threads - 1) / threads;
-  
-     const float* bptr = has_bias ? b.data_ptr<float>() : nullptr;
-  
-     depthwise_conv2d_kh1_kernel<<<blocks, threads>>>(
-         x.data_ptr<float>(),
-         w.data_ptr<float>(),
-         bptr,
-         out.data_ptr<float>(),
-         N, C, H, W,
-         K_h,
-         outH, outW,
-         (int)stride_h, (int)stride_w,
-         (int)pad_h, (int)pad_w,
-         (int)dil_h, (int)dil_w,
-         has_bias
-     );
-  
-     return out;
- }
- """
-  
-  
- _depthwise_conv_ext = load_inline(
-     name="depthwise_conv2d_kh1_ext",
-     cpp_sources=_depthwise_conv_cpp,
-     cuda_sources=_depthwise_conv_cu,
-     functions=["depthwise_conv2d_kh1_cuda"],
-     extra_cuda_cflags=["-O3"],
-     extra_cflags=["-O3"],
-     verbose=False,
- )
-  
-  
- class ModelNew(nn.Module):
-     """
-     Optimized depthwise Conv2d for kernel shape (kernel_size, 1) using a custom CUDA kernel.
-     Falls back to nn.Conv2d for unsupported cases/dtypes/devices.
-     """
-     def __init__(
-         self,
-         in_channels: int,
-         kernel_size: int,
-         stride: int = 1,
-         padding: int = 0,
-         dilation: int = 1,
-         bias: bool = False,
-     ):
-         super().__init__()
-         self.conv2d = nn.Conv2d(
-             in_channels,
-             in_channels,
-             kernel_size=(kernel_size, 1),
-             stride=stride,
-             padding=padding,
-             dilation=dilation,
-             groups=in_channels,
-             bias=bias,
-         )
-  
-     def forward(self, x: torch.Tensor) -> torch.Tensor:
-         if not x.is_cuda:
-             return self.conv2d(x)
-         if x.dtype != torch.float32:
-             return self.conv2d(x)
-         if not x.is_contiguous():
-             x = x.contiguous()
-  
-         w = self.conv2d.weight
-         b = self.conv2d.bias
-  
-         if not w.is_cuda or w.dtype != torch.float32:
-             return self.conv2d(x)
-  
-         if not w.is_contiguous():
-             w = w.contiguous()
-  
-         if b is None:
-             b_arg = torch.empty(0, device=x.device, dtype=torch.float32)
-         else:
-             if not b.is_contiguous():
-                 b = b.contiguous()
-             b_arg = b
-  
-         stride_h, stride_w = self.conv2d.stride
-         pad_h, pad_w = self.conv2d.padding
-         dil_h, dil_w = self.conv2d.dilation
-  
-         return _depthwise_conv_ext.depthwise_conv2d_kh1_cuda(
-             x,
-             w,
-             b_arg,
-             int(stride_h),
-             int(stride_w),
-             int(pad_h),
-             int(pad_w),
-             int(dil_h),
-             int(dil_w),
-         )
-  
-  
-  
  #!POPCORN leaderboard grayscale_v2
  #!POPCORN gpu A100
   
  import torch
   
- # Manual adaptation of this kernelbench candidate to Popcorn grayscale_v2 contract.
- # Input: (rgb, output) where rgb is HxWx3 float32 CUDA and output is HxW float32 CUDA.
- @torch.inference_mode()
- def custom_kernel(data):
-     if not isinstance(data, (tuple, list)) or len(data) != 2:
-         raise RuntimeError('Expected tuple(input_rgb, output)')
+ _WEIGHTS = None
   
-     rgb, out = data
-     if not isinstance(rgb, torch.Tensor) or not isinstance(out, torch.Tensor):
-         raise RuntimeError('Expected tensor inputs')
-     if rgb.ndim != 3 or rgb.shape[-1] != 3:
-         raise RuntimeError(f'Expected rgb shape [H, W, 3], got {tuple(rgb.shape)}')
   
-     # Ensure compute uses float32 CUDA tensor contiguous layout.
-     x = rgb
-     if x.dtype != torch.float32:
-         x = x.float()
-     if not x.is_cuda:
-         x = x.cuda()
-     if not x.is_contiguous():
-         x = x.contiguous()
+ def _weights(device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+     global _WEIGHTS
+     if _WEIGHTS is None or _WEIGHTS.device != device or _WEIGHTS.dtype != dtype:
+         _WEIGHTS = torch.tensor([0.2989, 0.5870, 0.1140], device=device, dtype=dtype)
+     return _WEIGHTS
   
-     # Convert HWC RGB -> NCHW for the existing depthwise kernel path.
-     x_nchw = x.permute(2, 0, 1).contiguous().unsqueeze(0)  # [1, 3, H, W]
   
-     # Depthwise 1x1 channel weights corresponding to grayscale coefficients.
-     w = torch.tensor(
-         [0.2989, 0.5870, 0.1140],
-         device=x_nchw.device,
-         dtype=torch.float32,
-     ).view(3, 1, 1, 1).contiguous()
-     b = torch.empty(0, device=x_nchw.device, dtype=torch.float32)
-  
-     y = _depthwise_conv_ext.depthwise_conv2d_kh1_cuda(
-         x_nchw,
-         w,
-         b,
-         1, 1, 0, 0, 1, 1,
-     )  # [1, 3, H, W]
-  
-     gray = y.sum(dim=1).squeeze(0)  # [H, W]
-     if out.shape == gray.shape and out.is_cuda:
-         out[...] = gray.to(dtype=out.dtype)
-         return out
-     return gray
+ @torch.inference_mode()
+ def custom_kernel(data):
+     rgb, out = data
+     # Reference-equivalent weighted sum over channel dimension.
+     w = _weights(rgb.device, rgb.dtype)
+     gray = torch.sum(rgb * w, dim=-1)
+     out[...] = gray
+     return out