import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size_h (int): Height of the convolution kernel.
        kernel_size_w (int): Width of the convolution kernel.
        stride_h (int, optional): Stride of the convolution in height dimension. Defaults to 1.
        stride_w (int, optional): Stride of the convolution in width dimension. Defaults to 1.
        padding_h (int, optional): Padding applied to the input in height dimension. Defaults to 0.
        padding_w (int, optional): Padding applied to the input in width dimension. Defaults to 0.
        dilation_h (int, optional): Spacing between kernel elements in height dimension. Defaults to 1.
        dilation_w (int, optional): Spacing between kernel elements in width dimension. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size_h: int, kernel_size_w: int, stride_h: int = 1, stride_w: int = 1, padding_h: int = 0, padding_w: int = 0, dilation_h: int = 1, dilation_w: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, (kernel_size_h, kernel_size_w), stride=(stride_h, stride_w), padding=(padding_h, padding_w), dilation=(dilation_h, dilation_w), 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, out_channels, height_out, width_out).
        """
        return self.conv2d(x)

# Test code
batch_size = 32
in_channels = 128
out_channels = 128
kernel_size_h = 3
kernel_size_w = 7
width = 256
height = 128
stride_h = 1
stride_w = 1
padding_h = 0
padding_w = 0
dilation_h = 1
dilation_w = 1
groups = in_channels

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size_h, kernel_size_w, stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w, groups]

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

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

torch::Tensor dwconv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride_h, int64_t stride_w,
                             int64_t padding_h, int64_t padding_w,
                             int64_t dilation_h, int64_t dilation_w);
"""

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

__global__ __launch_bounds__(256, 2)
void dwconv3x7_w256_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N, int C, int H, int Hout, int Wout,
                           bool has_bias) {
    // One block computes up to four adjacent output rows for one (N,C) plane,
    // and the full output width.  For the benchmark W=256, Wout=250.
    __shared__ float tile[6][256];

    const int c = blockIdx.x;
    const int n = blockIdx.y;
    const int oh0 = blockIdx.z * 4;
    const int tid = threadIdx.x;

    const float* xnc = x + ((n * C + c) * H) * 256;

    #pragma unroll
    for (int r = 0; r < 6; ++r) {
        int ih = oh0 + r;
        tile[r][tid] = (ih < H) ? xnc[ih * 256 + tid] : 0.0f;
    }
    __syncthreads();

    if (tid >= Wout) return;

    const float* wc = w + c * 21;
    float acc0 = has_bias ? b[c] : 0.0f;
    float acc1 = acc0;
    float acc2 = acc0;
    float acc3 = acc0;

    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 7; ++kw) {
            float ww = __ldg(wc + kh * 7 + kw);
            int col = tid + kw;
            acc0 = fmaf(tile[kh + 0][col], ww, acc0);
            acc1 = fmaf(tile[kh + 1][col], ww, acc1);
            acc2 = fmaf(tile[kh + 2][col], ww, acc2);
            acc3 = fmaf(tile[kh + 3][col], ww, acc3);
        }
    }

    float* ync = y + ((n * C + c) * Hout + oh0) * Wout;
    ync[tid] = acc0;
    if (oh0 + 1 < Hout) ync[Wout + tid] = acc1;
    if (oh0 + 2 < Hout) ync[2 * Wout + tid] = acc2;
    if (oh0 + 3 < Hout) ync[3 * Wout + tid] = acc3;
}

__global__ __launch_bounds__(256, 2)
void dwconv_generic_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           long total,
                           int N, int C, int H, int W,
                           int Kh, int Kw, int Hout, int Wout,
                           int sh, int sw, int ph, int pw, int dh, int dw,
                           bool has_bias) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = (int)(idx % Wout);
    long t = idx / Wout;
    int oh = (int)(t % Hout);
    t /= Hout;
    int c = (int)(t % C);
    int n = (int)(t / C);

    float acc = has_bias ? b[c] : 0.0f;
    const float* xnc = x + ((n * C + c) * H) * W;
    const float* wc = w + c * Kh * Kw;

    for (int kh = 0; kh < Kh; ++kh) {
        int ih = oh * sh - ph + kh * dh;
        if ((unsigned)ih >= (unsigned)H) continue;
        for (int kw = 0; kw < Kw; ++kw) {
            int iw = ow * sw - pw + kw * dw;
            if ((unsigned)iw < (unsigned)W) {
                acc = fmaf(xnc[ih * W + iw], wc[kh * Kw + kw], acc);
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor dwconv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride_h, int64_t stride_w,
                             int64_t padding_h, int64_t padding_w,
                             int64_t dilation_h, int64_t dilation_w) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int Kh = (int)weight.size(2);
    const int Kw = (int)weight.size(3);
    const int Hout = (H + 2 * (int)padding_h - (int)dilation_h * (Kh - 1) - 1) / (int)stride_h + 1;
    const int Wout = (W + 2 * (int)padding_w - (int)dilation_w * (Kw - 1) - 1) / (int)stride_w + 1;

    auto y = torch::empty({N, C, Hout, Wout}, x.options());
    const bool has_bias = (bias.numel() != 0);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (Kh == 3 && Kw == 7 && stride_h == 1 && stride_w == 1 &&
        padding_h == 0 && padding_w == 0 && dilation_h == 1 && dilation_w == 1 &&
        W == 256 && Wout <= 250 && x.scalar_type() == at::kFloat) {
        dim3 block(256);
        dim3 grid(C, N, (Hout + 3) / 4);
        dwconv3x7_w256_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(),
            N, C, H, Hout, Wout, has_bias);
    } else {
        long total = (long)N * C * Hout * Wout;
        dim3 block(256);
        dim3 grid((unsigned int)((total + block.x - 1) / block.x));
        dwconv_generic_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(), total,
            N, C, H, W, Kh, Kw, Hout, Wout,
            (int)stride_h, (int)stride_w, (int)padding_h, (int)padding_w,
            (int)dilation_h, (int)dilation_w, has_bias);
    }
    return y;
}
"""

_dwconv_ext = load_inline(
    name="dwconv_depthwise_3x7_shared_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["dwconv_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size_h: int, kernel_size_w: int,
                 stride_h: int = 1, stride_w: int = 1, padding_h: int = 0, padding_w: int = 0,
                 dilation_h: int = 1, dilation_w: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        # Match the reference exactly: out_channels/groups arguments are ignored there.
        self.conv2d = nn.Conv2d(in_channels, in_channels, (kernel_size_h, kernel_size_w),
                                stride=(stride_h, stride_w), padding=(padding_h, padding_w),
                                dilation=(dilation_h, dilation_w), groups=in_channels, bias=bias)
        self.stride_h = stride_h
        self.stride_w = stride_w
        self.padding_h = padding_h
        self.padding_w = padding_w
        self.dilation_h = dilation_h
        self.dilation_w = dilation_w
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bias = self.conv2d.bias if self.conv2d.bias is not None else self._empty_bias
        return _dwconv_ext.dwconv_forward(x.contiguous(), self.conv2d.weight.contiguous(), bias,
                                          self.stride_h, self.stride_w,
                                          self.padding_h, self.padding_w,
                                          self.dilation_h, self.dilation_w)