import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise 2D convolution operation with square input and square 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.
        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, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, 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, width).

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

# Test code
batch_size = 16
in_channels = 64
kernel_size = 3
width = 512
height = 512
stride = 1
padding = 0

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]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor dwconv3x3_nobias(torch::Tensor x, torch::Tensor weight);
'''

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

#define TILE 32
#define TX 16
#define TY 16

__global__ void dwconv3x3_kernel(const float* __restrict__ x,
                                 const float* __restrict__ w,
                                 float* __restrict__ y,
                                 int C, int H, int W, int Hout, int Wout) {
    __shared__ float sh[TILE + 2][TILE + 2];
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int tid = ty * TX + tx;
    int plane = blockIdx.z;
    int n = plane / C;
    int c = plane - n * C;
    int ox0 = blockIdx.x * TILE;
    int oy0 = blockIdx.y * TILE;

    const float* xp = x + ((n * C + c) * H) * W;
    float* yp = y + ((n * C + c) * Hout) * Wout;

    for (int idx = tid; idx < (TILE + 2) * (TILE + 2); idx += TX * TY) {
        int ly = idx / (TILE + 2);
        int lx = idx - ly * (TILE + 2);
        int iy = oy0 + ly;
        int ix = ox0 + lx;
        sh[ly][lx] = (iy < H && ix < W) ? xp[iy * W + ix] : 0.0f;
    }
    __syncthreads();

    float w00 = w[c * 9 + 0], w01 = w[c * 9 + 1], w02 = w[c * 9 + 2];
    float w10 = w[c * 9 + 3], w11 = w[c * 9 + 4], w12 = w[c * 9 + 5];
    float w20 = w[c * 9 + 6], w21 = w[c * 9 + 7], w22 = w[c * 9 + 8];

    int oxb = ox0 + tx * 2;
    int oyb = oy0 + ty * 2;
    int lxb = tx * 2;
    int lyb = ty * 2;

    #pragma unroll
    for (int dy = 0; dy < 2; ++dy) {
        int oy = oyb + dy;
        int ly = lyb + dy;
        if (oy < Hout) {
            #pragma unroll
            for (int dx = 0; dx < 2; ++dx) {
                int ox = oxb + dx;
                int lx = lxb + dx;
                if (ox < Wout) {
                    float acc = 0.0f;
                    acc = fmaf(sh[ly + 0][lx + 0], w00, acc);
                    acc = fmaf(sh[ly + 0][lx + 1], w01, acc);
                    acc = fmaf(sh[ly + 0][lx + 2], w02, acc);
                    acc = fmaf(sh[ly + 1][lx + 0], w10, acc);
                    acc = fmaf(sh[ly + 1][lx + 1], w11, acc);
                    acc = fmaf(sh[ly + 1][lx + 2], w12, acc);
                    acc = fmaf(sh[ly + 2][lx + 0], w20, acc);
                    acc = fmaf(sh[ly + 2][lx + 1], w21, acc);
                    acc = fmaf(sh[ly + 2][lx + 2], w22, acc);
                    yp[oy * Wout + ox] = acc;
                }
            }
        }
    }
}

torch::Tensor dwconv3x3_nobias(torch::Tensor x, torch::Tensor weight) {
    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 Hout = H - 2;
    int Wout = W - 2;
    auto y = torch::empty({N, C, Hout, Wout}, x.options());
    dim3 block(TX, TY);
    dim3 grid((Wout + TILE - 1) / TILE, (Hout + TILE - 1) / TILE, N * C);
    dwconv3x3_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>(), C, H, W, Hout, Wout);
    return y;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.dwconv3x3_nobias(x, self.conv2d.weight)