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 = r'''
#include <torch/extension.h>
torch::Tensor depthwise_conv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int stride_h, int stride_w, int padding_h, int padding_w,
                                     int dilation_h, int dilation_w);
'''

_cuda = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type()==at::kFloat, #x " must be float")

constexpr int BW = 32;
constexpr int OW_TILE = 128;
constexpr int TH = 8;
constexpr int KH = 3;
constexpr int KW = 7;
constexpr int SHW = OW_TILE + KW - 1;
constexpr int SHH = TH + KH - 1;

__global__ void dw3x7_4col(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N, int C, int H, int W, int OH, int OW, int has_bias) {
    __shared__ float smem[SHH * SHW];
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int ow0 = blockIdx.x * OW_TILE;
    int oh0 = blockIdx.y * TH;
    int nc = blockIdx.z;
    int n = nc / C;
    int c = nc - n * C;
    const float* xptr = x + ((n * C + c) * H) * W;
    int tid = ty * BW + tx;

    for (int idx = tid; idx < SHH * SHW; idx += BW * TH) {
        int rr = idx / SHW;
        int cc = idx - rr * SHW;
        int ih = oh0 + rr;
        int iw = ow0 + cc;
        smem[idx] = (ih < H && iw < W) ? xptr[ih * W + iw] : 0.0f;
    }
    __syncthreads();

    int oh = oh0 + ty;
    if (oh < OH) {
        const float* wptr = w + c * KH * KW;
        float init = has_bias ? b[c] : 0.0f;
        float a0 = init, a1 = init, a2 = init, a3 = init;
        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            int base = (ty + kh) * SHW + tx;
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                float wt = wptr[kh * KW + kw];
                a0 = fmaf(smem[base + kw], wt, a0);
                a1 = fmaf(smem[base + 32 + kw], wt, a1);
                a2 = fmaf(smem[base + 64 + kw], wt, a2);
                a3 = fmaf(smem[base + 96 + kw], wt, a3);
            }
        }
        int outbase = ((n * C + c) * OH + oh) * OW;
        int ow = ow0 + tx;
        if (ow < OW) y[outbase + ow] = a0;
        ow += 32;
        if (ow < OW) y[outbase + ow] = a1;
        ow += 32;
        if (ow < OW) y[outbase + ow] = a2;
        ow += 32;
        if (ow < OW) y[outbase + ow] = a3;
    }
}

__global__ void dw_generic(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N, int C, int H, int W, int OH, int OW,
                           int kH, int kW, int sh, int sw, int ph, int pw, int dh, int dw,
                           int has_bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * OH * OW;
    if (idx >= total) return;
    int ow = idx % OW;
    int t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int c = t % C;
    int n = t / C;
    float acc = has_bias ? b[c] : 0.0f;
    const float* xptr = x + ((n * C + c) * H) * W;
    const float* wptr = w + c * kH * kW;
    for (int kh = 0; kh < kH; ++kh) {
        int ih = oh * sh - ph + kh * dh;
        if ((unsigned)ih < (unsigned)H) {
            for (int kw = 0; kw < kW; ++kw) {
                int iw = ow * sw - pw + kw * dw;
                if ((unsigned)iw < (unsigned)W) {
                    acc = fmaf(xptr[ih * W + iw], wptr[kh * kW + kw], acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor depthwise_conv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int stride_h, int stride_w, int padding_h, int padding_w,
                                     int dilation_h, int dilation_w) {
    CHECK_CUDA(x);
    CHECK_CUDA(weight);
    CHECK_FLOAT(x);
    CHECK_FLOAT(weight);

    int64_t N64 = x.size(0), C64 = x.size(1), H64 = x.size(2), W64 = x.size(3);
    int64_t kH64 = weight.size(2), kW64 = weight.size(3);
    int64_t OH64 = (H64 + 2 * padding_h - dilation_h * (kH64 - 1) - 1) / stride_h + 1;
    int64_t OW64 = (W64 + 2 * padding_w - dilation_w * (kW64 - 1) - 1) / stride_w + 1;
    auto y = torch::empty({N64, C64, OH64, OW64}, x.options());

    int N = (int)N64, C = (int)C64, H = (int)H64, W = (int)W64;
    int OH = (int)OH64, OW = (int)OW64, kH = (int)kH64, kW = (int)kW64;
    int has_bias = bias.numel() != 0;
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    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) {
        dim3 block(BW, TH);
        dim3 grid((OW + OW_TILE - 1) / OW_TILE, (OH + TH - 1) / TH, N * C);
        dw3x7_4col<<<grid, block, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(),
                                               bptr, y.data_ptr<float>(), N, C, H, W, OH, OW,
                                               has_bias);
    } else {
        int total = N * C * OH * OW;
        dw_generic<<<(total + 255) / 256, 256, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(),
            N, C, H, W, OH, OW, kH, kW, stride_h, stride_w, padding_h, padding_w,
            dilation_h, dilation_w, has_bias);
    }
    return y;
}
'''

_ext = load_inline(
    name='dwconv2d_asym_ext_final',
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=['depthwise_conv_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    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().__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)
        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 _ext.depthwise_conv_forward(x, self.conv2d.weight, bias,
                                           self.stride_h, self.stride_w,
                                           self.padding_h, self.padding_w,
                                           self.dilation_h, self.dilation_w)