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 depthwise_conv2d_vertical(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                        int64_t stride, int64_t padding, int64_t dilation, bool has_bias);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void dw_vert3_4rows_vec4_kernel(const float* __restrict__ x,
                                           const float* __restrict__ w,
                                           const float* __restrict__ b,
                                           float* __restrict__ y,
                                           int C, int H, int W, int Ho,
                                           int vecs_per_row, bool has_bias) {
    int w4 = blockIdx.x * blockDim.x + threadIdx.x;
    if (w4 >= vecs_per_row) return;
    int g = blockIdx.y;
    int oh = g << 2;
    int nc = blockIdx.z;
    int c = nc % C;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
    long long xbase = ((long long)nc * H + oh) * vecs_per_row + w4;
    long long ybase = ((long long)nc * Ho + oh) * vecs_per_row + w4;
    float k0 = w[(long long)c * 3 + 0];
    float k1 = w[(long long)c * 3 + 1];
    float k2 = w[(long long)c * 3 + 2];
    float bv = has_bias ? b[c] : 0.0f;

    if (oh + 3 < Ho) {
        float4 a0 = x4[xbase];
        float4 a1 = x4[xbase + vecs_per_row];
        float4 a2 = x4[xbase + 2LL*vecs_per_row];
        float4 a3 = x4[xbase + 3LL*vecs_per_row];
        float4 a4 = x4[xbase + 4LL*vecs_per_row];
        float4 a5 = x4[xbase + 5LL*vecs_per_row];
        float4 o0, o1, o2, o3;
        o0.x = a0.x*k0 + a1.x*k1 + a2.x*k2 + bv;
        o0.y = a0.y*k0 + a1.y*k1 + a2.y*k2 + bv;
        o0.z = a0.z*k0 + a1.z*k1 + a2.z*k2 + bv;
        o0.w = a0.w*k0 + a1.w*k1 + a2.w*k2 + bv;
        o1.x = a1.x*k0 + a2.x*k1 + a3.x*k2 + bv;
        o1.y = a1.y*k0 + a2.y*k1 + a3.y*k2 + bv;
        o1.z = a1.z*k0 + a2.z*k1 + a3.z*k2 + bv;
        o1.w = a1.w*k0 + a2.w*k1 + a3.w*k2 + bv;
        o2.x = a2.x*k0 + a3.x*k1 + a4.x*k2 + bv;
        o2.y = a2.y*k0 + a3.y*k1 + a4.y*k2 + bv;
        o2.z = a2.z*k0 + a3.z*k1 + a4.z*k2 + bv;
        o2.w = a2.w*k0 + a3.w*k1 + a4.w*k2 + bv;
        o3.x = a3.x*k0 + a4.x*k1 + a5.x*k2 + bv;
        o3.y = a3.y*k0 + a4.y*k1 + a5.y*k2 + bv;
        o3.z = a3.z*k0 + a4.z*k1 + a5.z*k2 + bv;
        o3.w = a3.w*k0 + a4.w*k1 + a5.w*k2 + bv;
        y4[ybase] = o0;
        y4[ybase + vecs_per_row] = o1;
        y4[ybase + 2LL*vecs_per_row] = o2;
        y4[ybase + 3LL*vecs_per_row] = o3;
    } else {
        for (int r = 0; r < 4 && oh + r < Ho; ++r) {
            float4 a0 = x4[xbase + (long long)r*vecs_per_row];
            float4 a1 = x4[xbase + (long long)(r+1)*vecs_per_row];
            float4 a2 = x4[xbase + (long long)(r+2)*vecs_per_row];
            float4 o;
            o.x = a0.x*k0 + a1.x*k1 + a2.x*k2 + bv;
            o.y = a0.y*k0 + a1.y*k1 + a2.y*k2 + bv;
            o.z = a0.z*k0 + a1.z*k1 + a2.z*k2 + bv;
            o.w = a0.w*k0 + a1.w*k1 + a2.w*k2 + bv;
            y4[ybase + (long long)r*vecs_per_row] = o;
        }
    }
}

__global__ void dw_vertical_generic_kernel(const float* __restrict__ x,
                                           const float* __restrict__ w,
                                           const float* __restrict__ b,
                                           float* __restrict__ y,
                                           int N, int C, int H, int W, int Ho, int Wo,
                                           int K, int stride, int padding, int dilation,
                                           bool has_bias) {
    long long total = (long long)N * C * Ho * Wo;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < total;
         idx += (long long)blockDim.x * gridDim.x) {
        int ow = idx % Wo;
        long long t = idx / Wo;
        int oh = t % Ho;
        t /= Ho;
        int c = t % C;
        int n = t / C;
        float acc = has_bias ? b[c] : 0.0f;
        int iw = ow * stride - padding;
        if ((unsigned)iw < (unsigned)W) {
            for (int kh = 0; kh < K; ++kh) {
                int ih = oh * stride - padding + kh * dilation;
                if ((unsigned)ih < (unsigned)H) {
                    acc += x[((long long)n*C + c)*H*W + (long long)ih*W + iw] * w[(long long)c*K + kh];
                }
            }
        }
        y[idx] = acc;
    }
}

torch::Tensor depthwise_conv2d_vertical(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                        int64_t stride, int64_t padding, int64_t dilation, bool has_bias) {
    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 supported");
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int K = (int)weight.size(2);
    int Ho = (H + 2*(int)padding - (int)dilation*(K-1) - 1) / (int)stride + 1;
    int Wo = (W + 2*(int)padding - 1) / (int)stride + 1;
    auto y = torch::empty({N, C, Ho, Wo}, x.options());
    if (N == 0 || C == 0 || Ho <= 0 || Wo <= 0) return y;
    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;
    if (K == 3 && stride == 1 && padding == 0 && dilation == 1 && Wo == W && (W % 4 == 0)) {
        int vecs = W >> 2;
        dim3 block(128);
        dim3 grid((vecs + block.x - 1) / block.x, (Ho + 3) >> 2, N * C);
        dw_vert3_4rows_vec4_kernel<<<grid, block>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), C, H, W, Ho, vecs, has_bias);
    } else {
        const int threads = 256;
        long long total = (long long)N * C * Ho * Wo;
        int blocks = (int)min((long long)131072, (total + threads - 1) / threads);
        dw_vertical_generic_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), N, C, H, W, Ho, Wo, K, (int)stride, (int)padding, (int)dilation, has_bias);
    }
    return y;
}
'''

_ext = load_inline(name='dw_vertical_ext_v3', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['depthwise_conv2d_vertical'], verbose=False,
                   extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])

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 = stride
        self.padding = padding
        self.dilation = dilation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv2d.bias
        if b is None:
            b = self.conv2d.weight.new_empty((0,))
            has_b = False
        else:
            has_b = True
        return _ext.depthwise_conv2d_vertical(x, self.conv2d.weight, b, self.stride, self.padding, self.dilation, has_b)