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 depthwise_conv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride, int64_t padding);
"""

_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")

static constexpr int TILE_H = 8;

__global__ void dw3x3_s1p0_tiled_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 Hout, int Wout, int has_bias)
{
    extern __shared__ float sh[];
    float* sw = sh + (TILE_H + 2) * W;

    const int tx = threadIdx.x;       // 0..63
    const int ty = threadIdx.y;       // 0..7
    const int z = blockIdx.z;
    const int c = z % C;
    const int n = z / C;
    const int h0 = blockIdx.y * TILE_H;
    const int tile_rows = min(TILE_H, Hout - h0);
    const int load_rows = tile_rows + 2;

    const float* xbase = x + ((n * C + c) * H + h0) * W;
    const float* wbase = w + c * 9;

    const int tid = ty * blockDim.x + tx;
    if (tid < 9) sw[tid] = wbase[tid];

    // Load the input tile: (tile_rows+2) x W.  This avoids the 9x rereads
    // that a direct depthwise implementation would perform.
    for (int rr = ty; rr < load_rows; rr += blockDim.y) {
        const float* src = xbase + rr * W;
        float* dst = sh + rr * W;
        for (int col = tx; col < W; col += blockDim.x) {
            dst[col] = src[col];
        }
    }
    __syncthreads();

    if (ty < tile_rows) {
        const float w0 = sw[0], w1 = sw[1], w2 = sw[2];
        const float w3 = sw[3], w4 = sw[4], w5 = sw[5];
        const float w6 = sw[6], w7 = sw[7], w8 = sw[8];
        const float bv = has_bias ? b[c] : 0.0f;
        float* yrow = y + ((n * C + c) * Hout + (h0 + ty)) * Wout;

        for (int ow = tx; ow < Wout; ow += blockDim.x) {
            const float* p = sh + ty * W + ow;
            float acc = bv;
            acc = fmaf(p[0],       w0, acc);
            acc = fmaf(p[1],       w1, acc);
            acc = fmaf(p[2],       w2, acc);
            acc = fmaf(p[W],       w3, acc);
            acc = fmaf(p[W + 1],   w4, acc);
            acc = fmaf(p[W + 2],   w5, acc);
            acc = fmaf(p[2 * W],   w6, acc);
            acc = fmaf(p[2 * W+1], w7, acc);
            acc = fmaf(p[2 * W+2], w8, acc);
            yrow[ow] = acc;
        }
    }
}

__global__ void depthwise_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    long long total, int C, int H, int W, int K, int stride, int pad, int Hout, int Wout, int has_bias)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

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

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

    int ih0 = oh * stride - pad;
    int iw0 = ow * stride - pad;
    for (int kh = 0; kh < K; ++kh) {
        int ih = ih0 + kh;
        if ((unsigned)ih >= (unsigned)H) continue;
        for (int kw = 0; kw < K; ++kw) {
            int iw = iw0 + kw;
            if ((unsigned)iw < (unsigned)W) {
                acc = fmaf(xnc[ih * W + iw], wc[kh * K + kw], acc);
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor depthwise_conv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride64, int64_t padding64) {
    CHECK_CUDA(x);
    CHECK_CUDA(weight);
    CHECK_CONTIGUOUS(x);
    CHECK_CONTIGUOUS(weight);
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "only float32 weights are supported");

    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 K = (int)weight.size(2);
    const int stride = (int)stride64;
    const int pad = (int)padding64;
    const int Hout = (H + 2 * pad - K) / stride + 1;
    const int Wout = (W + 2 * pad - K) / stride + 1;

    auto y = torch::empty({N, C, Hout, Wout}, x.options());
    const int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && pad == 0 && W > 0 && Wout > 0) {
        dim3 block(64, TILE_H, 1);
        dim3 grid(1, (Hout + TILE_H - 1) / TILE_H, N * C);
        size_t smem = ((TILE_H + 2) * (size_t)W + 9) * sizeof(float);
        dw3x3_s1p0_tiled_kernel<<<grid, block, smem, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(),
            N, C, H, W, Hout, Wout, has_bias);
    } else {
        long long total = (long long)N * C * Hout * Wout;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        depthwise_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(),
            total, C, H, W, K, stride, pad, Hout, Wout, has_bias);
    }
    return y;
}
"""

_ext = load_inline(
    name="depthwise_conv2d_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["depthwise_conv_forward"],
    extra_cflags=["-O3"],
    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)
        self.stride = stride
        self.padding = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bias = self.conv2d.bias
        if bias is None:
            bias = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.depthwise_conv_forward(x, self.conv2d.weight, bias, int(self.stride), int(self.padding))