import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        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: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=(kernel_size, kernel_size), stride=stride, padding=padding, 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_in, width_in).

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

# Test code
batch_size = 64
in_channels = 128
out_channels = 128
kernel_size = 3
width_in = 512
height_in = 256
stride = 1
padding = 0

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

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

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

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

// Specialized fast path for the benchmark case: depthwise 3x3, stride 1, padding 0,
// NCHW contiguous, channel multiplier 1.
template<int BW, int BH, int WPT>
__global__ __launch_bounds__(256, 2)
void dw3x3_s1p0_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 has_bias) {
    // blockDim = (32, 8).  Each thread computes WPT outputs along x.
    constexpr int TX = 32;
    constexpr int SHW = BW + 2;
    constexpr int SHH = BH + 2;
    __shared__ float smem[SHH * SHW];
    __shared__ float sw[9];
    __shared__ float sbias;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;

    const int nc = blockIdx.z;
    const int n = nc / C;
    const int c = nc - n * C;
    const int h0 = blockIdx.y * BH;
    const int w0 = blockIdx.x * BW;

    if (tid < 9) sw[tid] = w[c * 9 + tid];
    if (tid == 9) sbias = has_bias ? b[c] : 0.0f;

    // Cooperative load of the input tile needed for the BH x BW output tile.
    const int tile_elems = SHH * SHW;
    const int base = (n * C + c) * H * W;
    for (int idx = tid; idx < tile_elems; idx += 256) {
        int ih = idx / SHW;
        int iw = idx - ih * SHW;
        int gh = h0 + ih;
        int gw = w0 + iw;
        float v = 0.0f;
        // For padding=0 benchmark, right/bottom edge tiles can overrun the logical output tile.
        if (gh < H && gw < W) v = x[base + gh * W + gw];
        smem[idx] = v;
    }
    __syncthreads();

    const int oh = h0 + ty;
    if (oh < Ho) {
        #pragma unroll
        for (int i = 0; i < WPT; ++i) {
            int ox_l = tx + i * TX;
            int ow = w0 + ox_l;
            if (ow < Wo) {
                const int sidx = ty * SHW + ox_l;
                float acc = sbias;
                acc = fmaf(smem[sidx + 0],       sw[0], acc);
                acc = fmaf(smem[sidx + 1],       sw[1], acc);
                acc = fmaf(smem[sidx + 2],       sw[2], acc);
                acc = fmaf(smem[sidx + SHW + 0], sw[3], acc);
                acc = fmaf(smem[sidx + SHW + 1], sw[4], acc);
                acc = fmaf(smem[sidx + SHW + 2], sw[5], acc);
                acc = fmaf(smem[sidx + 2*SHW + 0], sw[6], acc);
                acc = fmaf(smem[sidx + 2*SHW + 1], sw[7], acc);
                acc = fmaf(smem[sidx + 2*SHW + 2], sw[8], acc);
                y[((n * C + c) * Ho + oh) * Wo + ow] = acc;
            }
        }
    }
}

// More general, still simple, fallback for non-special constructor settings (keeps semantics correct
// for common depthwise cases without calling back into PyTorch conv).
__global__ void dw_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 OC, int K, int stride, int pad,
                                  int Ho, int Wo, int has_bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * OC * Ho * Wo;
    int mult = OC / C;
    for (; idx < total; idx += blockDim.x * gridDim.x) {
        int t = idx;
        int ow = t % Wo; t /= Wo;
        int oh = t % Ho; t /= Ho;
        int oc = t % OC; t /= OC;
        int n = t;
        int c = oc / mult;
        float acc = has_bias ? b[oc] : 0.0f;
        for (int kh = 0; kh < K; ++kh) {
            int ih = oh * stride + kh - pad;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < K; ++kw) {
                int iw = ow * stride + kw - pad;
                if ((unsigned)iw < (unsigned)W) {
                    acc = fmaf(x[((n*C + c)*H + ih)*W + iw], w[(oc*K + kh)*K + kw], acc);
                }
            }
        }
        y[idx] = acc;
    }
}

torch::Tensor depthwise_conv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int64_t stride, int64_t padding) {
    CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight); CHECK_CONTIGUOUS(bias);
    CHECK_FLOAT(x); CHECK_FLOAT(weight); CHECK_FLOAT(bias);
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    TORCH_CHECK(weight.dim() == 4, "weight must be OIHW");
    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 OC = (int)weight.size(0);
    int K = (int)weight.size(2);
    int has_bias = (bias.numel() != 0);
    int Ho = (H + 2*(int)padding - K) / (int)stride + 1;
    int Wo = (W + 2*(int)padding - K) / (int)stride + 1;
    auto y = torch::empty({N, OC, Ho, Wo}, x.options());

    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && padding == 0 && OC == C) {
        constexpr int BW = 128;
        constexpr int BH = 8;
        constexpr int WPT = 4;
        dim3 block(32, 8, 1);
        dim3 grid((Wo + BW - 1) / BW, (Ho + BH - 1) / BH, N * C);
        dw3x3_s1p0_kernel<BW, BH, WPT><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),
            N, C, H, W, Ho, Wo, has_bias);
    } else {
        int total = N * OC * Ho * Wo;
        int threads = 256;
        int blocks = min((total + threads - 1) / threads, 65535);
        dw_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),
            N, C, H, W, OC, K, (int)stride, (int)padding, Ho, Wo, has_bias);
    }
    return y;
}
'''

_ext = load_inline(
    name='dwconv3x3_s1p0_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    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: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=(kernel_size, kernel_size), stride=stride, padding=padding, groups=in_channels, bias=bias)
        self.stride = stride
        self.padding = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # The benchmark supplies contiguous CUDA fp32 tensors.  Keeping this lean avoids Python overhead.
        b = self.conv2d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.depthwise_conv_forward(x, self.conv2d.weight, b, int(self.stride), int(self.padding))