import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise 2D convolution with a square input and an asymmetric kernel.

    Args:
        in_channels (int): Number of channels in the input tensor.
        kernel_size (int): Size of the convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(Model, self).__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:
        """
        Performs the depthwise 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, in_channels, height_out, width_out).
        """
        return self.conv2d(x)

# Test code
batch_size = 64
in_channels = 8
kernel_size = 3
width = 512
height = 512
stride = 1
padding = 0
dilation = 1

def get_inputs():
    x = torch.rand(batch_size, in_channels, height, width)
    return [x]

def get_init_inputs():
    return [in_channels, kernel_size, stride, padding, dilation]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor vertical_conv_forward(torch::Tensor x,
                                    torch::Tensor weight,
                                    torch::Tensor bias,
                                    int64_t stride,
                                    int64_t padding,
                                    int64_t dilation);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void dw_vertical_k3_s1p0_vec4_kernel(const float* __restrict__ x,
                                                const float* __restrict__ w,
                                                const float* __restrict__ bias,
                                                float* __restrict__ out,
                                                int N, int C, int H, int W,
                                                int Hout, int Wout,
                                                bool has_bias) {
    int ow = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    int oh0 = blockIdx.y * 4;
    int nc = blockIdx.z;
    if (ow >= Wout || oh0 >= Hout) return;

    int c = nc % C;
    float w0 = w[c * 3 + 0];
    float w1 = w[c * 3 + 1];
    float w2 = w[c * 3 + 2];
    float bv = has_bias ? bias[c] : 0.0f;

    const float* in = x + ((int64_t)nc * H) * W;
    float* o = out + ((int64_t)nc * Hout) * Wout;

    // Fast path for the overwhelmingly common full 4-column by 4-row tile.
    if (ow + 3 < Wout && oh0 + 3 < Hout) {
        const float4 r0 = *reinterpret_cast<const float4*>(in + (int64_t)(oh0 + 0) * W + ow);
        const float4 r1 = *reinterpret_cast<const float4*>(in + (int64_t)(oh0 + 1) * W + ow);
        const float4 r2 = *reinterpret_cast<const float4*>(in + (int64_t)(oh0 + 2) * W + ow);
        const float4 r3 = *reinterpret_cast<const float4*>(in + (int64_t)(oh0 + 3) * W + ow);
        const float4 r4 = *reinterpret_cast<const float4*>(in + (int64_t)(oh0 + 4) * W + ow);
        const float4 r5 = *reinterpret_cast<const float4*>(in + (int64_t)(oh0 + 5) * W + ow);

        float4 y0, y1, y2, y3;
        y0.x = r0.x*w0 + r1.x*w1 + r2.x*w2 + bv;
        y0.y = r0.y*w0 + r1.y*w1 + r2.y*w2 + bv;
        y0.z = r0.z*w0 + r1.z*w1 + r2.z*w2 + bv;
        y0.w = r0.w*w0 + r1.w*w1 + r2.w*w2 + bv;

        y1.x = r1.x*w0 + r2.x*w1 + r3.x*w2 + bv;
        y1.y = r1.y*w0 + r2.y*w1 + r3.y*w2 + bv;
        y1.z = r1.z*w0 + r2.z*w1 + r3.z*w2 + bv;
        y1.w = r1.w*w0 + r2.w*w1 + r3.w*w2 + bv;

        y2.x = r2.x*w0 + r3.x*w1 + r4.x*w2 + bv;
        y2.y = r2.y*w0 + r3.y*w1 + r4.y*w2 + bv;
        y2.z = r2.z*w0 + r3.z*w1 + r4.z*w2 + bv;
        y2.w = r2.w*w0 + r3.w*w1 + r4.w*w2 + bv;

        y3.x = r3.x*w0 + r4.x*w1 + r5.x*w2 + bv;
        y3.y = r3.y*w0 + r4.y*w1 + r5.y*w2 + bv;
        y3.z = r3.z*w0 + r4.z*w1 + r5.z*w2 + bv;
        y3.w = r3.w*w0 + r4.w*w1 + r5.w*w2 + bv;

        *reinterpret_cast<float4*>(o + (int64_t)(oh0 + 0) * Wout + ow) = y0;
        *reinterpret_cast<float4*>(o + (int64_t)(oh0 + 1) * Wout + ow) = y1;
        *reinterpret_cast<float4*>(o + (int64_t)(oh0 + 2) * Wout + ow) = y2;
        *reinterpret_cast<float4*>(o + (int64_t)(oh0 + 3) * Wout + ow) = y3;
    } else {
        // Boundary tiles (only the last row tile for the benchmark shape).
        #pragma unroll
        for (int rr = 0; rr < 4; ++rr) {
            int oh = oh0 + rr;
            if (oh < Hout) {
                #pragma unroll
                for (int jj = 0; jj < 4; ++jj) {
                    int col = ow + jj;
                    if (col < Wout) {
                        float v = in[(int64_t)(oh + 0) * W + col] * w0
                                + in[(int64_t)(oh + 1) * W + col] * w1
                                + in[(int64_t)(oh + 2) * W + col] * w2
                                + bv;
                        o[(int64_t)oh * Wout + col] = v;
                    }
                }
            }
        }
    }
}

__global__ void dw_vertical_generic_kernel(const float* __restrict__ x,
                                           const float* __restrict__ w,
                                           const float* __restrict__ bias,
                                           float* __restrict__ out,
                                           int N, int C, int H, int W,
                                           int Hout, int Wout, int K,
                                           int stride, int padding, int dilation,
                                           bool has_bias,
                                           int64_t total) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = idx % Wout;
    int64_t t = idx / Wout;
    int oh = t % Hout;
    t /= Hout;
    int c = t % C;
    int n = t / C;

    int ih0 = oh * stride - padding;
    int iw = ow * stride - padding;
    float sum = 0.0f;
    if ((unsigned)iw < (unsigned)W) {
        for (int k = 0; k < K; ++k) {
            int ih = ih0 + k * dilation;
            if ((unsigned)ih < (unsigned)H) {
                sum += x[((int64_t)n * C + c) * H * W + (int64_t)ih * W + iw] * w[c * K + k];
            }
        }
    }
    if (has_bias) sum += bias[c];
    out[idx] = sum;
}

torch::Tensor vertical_conv_forward(torch::Tensor x,
                                    torch::Tensor weight,
                                    torch::Tensor bias,
                                    int64_t stride64,
                                    int64_t padding64,
                                    int64_t dilation64) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 is supported");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "only float32 weights are supported");
    c10::cuda::CUDAGuard device_guard(x.device());

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.numel() ? bias.contiguous() : bias;

    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int H = (int)xc.size(2);
    int W = (int)xc.size(3);
    int K = (int)wc.size(2);
    int stride = (int)stride64;
    int padding = (int)padding64;
    int dilation = (int)dilation64;

    int Hout = (H + 2 * padding - dilation * (K - 1) - 1) / stride + 1;
    int Wout = (W + 2 * padding - 1) / stride + 1;
    auto out = torch::empty({N, C, Hout, Wout}, xc.options());

    bool has_bias = (bc.numel() > 0);
    const float* bptr = has_bias ? bc.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && padding == 0 && dilation == 1 && Wout == W) {
        dim3 block(128);
        dim3 grid((Wout + (int)block.x * 4 - 1) / ((int)block.x * 4),
                  (Hout + 3) / 4,
                  N * C);
        dw_vertical_k3_s1p0_vec4_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bptr, out.data_ptr<float>(),
            N, C, H, W, Hout, Wout, has_bias);
    } else {
        int64_t total = (int64_t)N * C * Hout * Wout;
        dim3 block(256);
        dim3 grid((unsigned int)((total + block.x - 1) / block.x));
        dw_vertical_generic_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bptr, out.data_ptr<float>(),
            N, C, H, W, Hout, Wout, K, stride, padding, dilation, has_bias, total);
    }
    return out;
}
'''

_ext = load_inline(
    name="dw_vertical_conv_inline_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["vertical_conv_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(ModelNew, self).__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)
        self._stride_i = stride if isinstance(stride, int) else stride[0]
        self._padding_i = padding if isinstance(padding, int) else padding[0]
        self._dilation_i = dilation if isinstance(dilation, int) else dilation[0]
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv2d.bias
        if b is None:
            b = self._empty_bias
            if b.device != x.device or b.dtype != x.dtype:
                b = x.new_empty((0,))
        return _ext.vertical_conv_forward(x, self.conv2d.weight, b, self._stride_i, self._padding_i, self._dilation_i)