submission 548205

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)')

    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()

    # 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