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 os
os.environ['TMPDIR'] = '/tmp'
os.environ['TEMP'] = '/tmp'
os.environ['TMP'] = '/tmp'
os.environ.setdefault('MAX_JOBS', '1')
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = """
#include <torch/extension.h>
torch::Tensor depthwise_asym_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias, int stride, int padding, int dilation);
"""

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

__device__ __forceinline__ float4 madd3(float4 a, float aw, float4 b, float bw, float4 c, float cw, float bv) {
    float4 o;
    o.x = fmaf(c.x, cw, fmaf(b.x, bw, a.x * aw)) + bv;
    o.y = fmaf(c.y, cw, fmaf(b.y, bw, a.y * aw)) + bv;
    o.z = fmaf(c.z, cw, fmaf(b.z, bw, a.z * aw)) + bv;
    o.w = fmaf(c.w, cw, fmaf(b.w, bw, a.w * aw)) + bv;
    return o;
}

__global__ void dw_k3s1p0_row6_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ b,
                                      float* __restrict__ y,
                                      int total, bool has_bias) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int step = blockDim.x * gridDim.x;
    for (int idx = tid; idx < total; idx += step) {
        int vecw = idx & 127;
        int q = idx >> 7;
        int rg = q % 85;
        int nc = q / 85;
        int c = nc & 7;
        int oh = rg * 6;
        int offx = (nc * 512 + oh) * 512 + (vecw << 2);
        int offy = (nc * 510 + oh) * 512 + (vecw << 2);
        float w0 = __ldg(w + c * 3 + 0);
        float w1 = __ldg(w + c * 3 + 1);
        float w2 = __ldg(w + c * 3 + 2);
        float bv = has_bias ? __ldg(b + c) : 0.0f;
        float4 r0 = *reinterpret_cast<const float4*>(x + offx);
        float4 r1 = *reinterpret_cast<const float4*>(x + offx + 512);
        float4 r2 = *reinterpret_cast<const float4*>(x + offx + 1024);
        *reinterpret_cast<float4*>(y + offy) = madd3(r0, w0, r1, w1, r2, w2, bv);
        float4 r3 = *reinterpret_cast<const float4*>(x + offx + 1536);
        *reinterpret_cast<float4*>(y + offy + 512) = madd3(r1, w0, r2, w1, r3, w2, bv);
        float4 r4 = *reinterpret_cast<const float4*>(x + offx + 2048);
        *reinterpret_cast<float4*>(y + offy + 1024) = madd3(r2, w0, r3, w1, r4, w2, bv);
        float4 r5 = *reinterpret_cast<const float4*>(x + offx + 2560);
        *reinterpret_cast<float4*>(y + offy + 1536) = madd3(r3, w0, r4, w1, r5, w2, bv);
        float4 r6 = *reinterpret_cast<const float4*>(x + offx + 3072);
        *reinterpret_cast<float4*>(y + offy + 2048) = madd3(r4, w0, r5, w1, r6, w2, bv);
        float4 r7 = *reinterpret_cast<const float4*>(x + offx + 3584);
        *reinterpret_cast<float4*>(y + offy + 2560) = madd3(r5, w0, r6, w1, r7, w2, bv);
    }
}

__global__ void dw_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,
                                  const float* __restrict__ b, float* __restrict__ y,
                                  long total, int C, int H, int W, int Hout, int Wout, int K,
                                  int stride, int padding, int dilation, bool has_bias) {
    long tid = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long step = (long)blockDim.x * gridDim.x;
    for (long idx = tid; idx < total; idx += step) {
        int ow = idx % Wout;
        long q = idx / Wout;
        int oh = q % Hout; q /= Hout;
        int c = q % C;
        long n = q / C;
        float acc = has_bias ? b[c] : 0.0f;
        int ih0 = oh * stride - padding;
        int iw = ow * stride - padding;
        if ((unsigned)iw < (unsigned)W) {
            const float* xp = x + ((n*C + c)*H)*W + iw;
            const float* wp = w + c*K;
            for (int kh = 0; kh < K; ++kh) {
                int ih = ih0 + kh * dilation;
                if ((unsigned)ih < (unsigned)H) acc = fmaf(xp[ih * W], wp[kh], acc);
            }
        }
        y[idx] = acc;
    }
}

torch::Tensor depthwise_asym_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias, int stride, int padding, int dilation) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int K = (int)weight.size(2);
    int Hout = (H + 2*padding - dilation*(K-1) - 1) / stride + 1;
    int Wout = (W + 2*padding - 1) / stride + 1;
    auto y = torch::empty({N, C, Hout, Wout}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (K == 3 && C == 8 && H == 512 && W == 512 && stride == 1 && padding == 0 && dilation == 1) {
        int total = N * C * 85 * 128;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 65535) blocks = 65535;
        dw_k3s1p0_row6_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(), total, has_bias);
    } else {
        long total = (long)N*C*Hout*Wout;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        dw_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(), total, C, H, W, Hout, Wout, K, stride, padding, dilation, has_bias);
    }
    return y;
}
"""

_ext = load_inline(name="dw_asym_ext_final", cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=["depthwise_asym_forward"], 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
        self.register_buffer('_dummy_bias', torch.empty(0), persistent=False)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv2d.bias
        if b is None:
            return _ext.depthwise_asym_forward(x, self.conv2d.weight, self._dummy_bias, False, self._stride, self._padding, self._dilation)
        return _ext.depthwise_asym_forward(x, self.conv2d.weight, b, True, self._stride, self._padding, self._dilation)