import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise-separable 2D convolution operation.

    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 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, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the depthwise-separable 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).
        """
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

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

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, stride, padding, dilation]

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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor depthwise_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias, bool has_bias,
                             int64_t stride, int64_t padding, int64_t dilation);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.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")

// Optimized path for the benchmark case: k=3, stride=1, padding=1, dilation=1.
// One CTA computes a 32x16 spatial tile for one (N,C) plane.  The input tile
// including halo is staged in shared memory, reducing the 3x3 depthwise pass to
// about one global load per output element.
__global__ __launch_bounds__(512, 2)
void dw3x3s1p1_kernel(const float* __restrict__ x,
                      const float* __restrict__ w,
                      const float* __restrict__ b,
                      float* __restrict__ y,
                      int NC, int H, int W) {
    constexpr int TW = 32;
    constexpr int TH = 16;
    constexpr int SW = TW + 2;
    constexpr int SH = TH + 2;
    __shared__ float sh[SH * SW];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    const int z = blockIdx.z;       // n*C + c
    const int c = z % NC;           // actually C is passed as NC? see below no
}

__global__ __launch_bounds__(512, 2)
void dw3x3s1p1_kernel2(const float* __restrict__ x,
                       const float* __restrict__ w,
                       const float* __restrict__ b,
                       float* __restrict__ y,
                       int C, int H, int W) {
    constexpr int TW = 32;
    constexpr int TH = 16;
    constexpr int SW = TW + 2;
    constexpr int SH = TH + 2;
    __shared__ float sh[SH * SW];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * 32 + tx;
    const int z = blockIdx.z;       // n*C + c
    const int c = z % C;
    const int ox0 = blockIdx.x * TW;
    const int oy0 = blockIdx.y * TH;

    // Cooperative load of the 34x18 tile including zero-padding halo.
    #pragma unroll
    for (int i = 0; i < SH * SW; i += 512) {
        int j = i + tid;
        if (j < SH * SW) {
            int sy = j / SW;
            int sx = j - sy * SW;
            int iy = oy0 + sy - 1;
            int ix = ox0 + sx - 1;
            float v = 0.0f;
            if ((unsigned)iy < (unsigned)H && (unsigned)ix < (unsigned)W) {
                v = x[((long long)z * H + iy) * W + ix];
            }
            sh[j] = v;
        }
    }
    __syncthreads();

    const int ox = ox0 + tx;
    const int oy = oy0 + ty;
    if (ox < W && oy < H) {
        const float* __restrict__ wc = w + c * 9;
        float acc = 0.0f;
        acc = fmaf(wc[0], sh[(ty + 0) * SW + tx + 0], acc);
        acc = fmaf(wc[1], sh[(ty + 0) * SW + tx + 1], acc);
        acc = fmaf(wc[2], sh[(ty + 0) * SW + tx + 2], acc);
        acc = fmaf(wc[3], sh[(ty + 1) * SW + tx + 0], acc);
        acc = fmaf(wc[4], sh[(ty + 1) * SW + tx + 1], acc);
        acc = fmaf(wc[5], sh[(ty + 1) * SW + tx + 2], acc);
        acc = fmaf(wc[6], sh[(ty + 2) * SW + tx + 0], acc);
        acc = fmaf(wc[7], sh[(ty + 2) * SW + tx + 1], acc);
        acc = fmaf(wc[8], sh[(ty + 2) * SW + tx + 2], acc);
        if (b != nullptr) acc += b[c];
        y[((long long)z * H + oy) * W + ox] = acc;
    }
}

__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 K,
                                  int Ho, int Wo, int stride, int pad, int dil) {
    long long total = (long long)N * C * Ho * Wo;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long step = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += step) {
        int ow = idx % Wo;
        long long t = idx / Wo;
        int oh = t % Ho;
        t /= Ho;
        int c = t % C;
        int n = t / C;
        float acc = 0.0f;
        int base_y = oh * stride - pad;
        int base_x = ow * stride - pad;
        const float* __restrict__ wc = w + (long long)c * K * K;
        for (int ky = 0; ky < K; ++ky) {
            int iy = base_y + ky * dil;
            if ((unsigned)iy >= (unsigned)H) continue;
            for (int kx = 0; kx < K; ++kx) {
                int ix = base_x + kx * dil;
                if ((unsigned)ix < (unsigned)W) {
                    float xv = x[((long long)n * C + c) * H * W + (long long)iy * W + ix];
                    acc = fmaf(wc[ky * K + kx], xv, acc);
                }
            }
        }
        if (b != nullptr) acc += b[c];
        y[idx] = acc;
    }
}

torch::Tensor depthwise_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias, bool has_bias,
                             int64_t stride, int64_t padding, int64_t dilation) {
    CHECK_CUDA(x); CHECK_CUDA(w);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(w);
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));

    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)w.size(2);
    const int Ho = (H + 2 * (int)padding - (int)dilation * (K - 1) - 1) / (int)stride + 1;
    const int Wo = (W + 2 * (int)padding - (int)dilation * (K - 1) - 1) / (int)stride + 1;

    auto y = torch::empty({N, C, 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 == 1 && dilation == 1 && Ho == H && Wo == W) {
        dim3 block(32, 16, 1);
        dim3 grid((W + 31) / 32, (H + 15) / 16, N * C);
        dw3x3s1p1_kernel2<<<grid, block, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), bp,
                                                       y.data_ptr<float>(), C, H, W);
    } else {
        int threads = 256;
        int blocks = (int)std::min<long long>(65535LL, ((long long)N * C * Ho * Wo + threads - 1) / threads);
        dw_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), bp,
                                                          y.data_ptr<float>(), N, C, H, W, K, Ho, Wo,
                                                          (int)stride, (int)padding, (int)dilation);
    }
    return y;
}
"""

_ext = load_inline(
    name="dw_sep_depthwise_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["depthwise_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size, stride=stride,
                                   padding=padding, dilation=dilation, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
        self.stride = int(stride)
        self.padding = int(padding)
        self.dilation = int(dilation)
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Custom CUDA depthwise pass, followed by the highly optimized library 1x1 convolution.
        db = self.depthwise.bias
        if db is None:
            dw = _ext.depthwise_cuda(x, self.depthwise.weight, self._empty_bias, False,
                                     self.stride, self.padding, self.dilation)
        else:
            dw = _ext.depthwise_cuda(x, self.depthwise.weight, db, True,
                                     self.stride, self.padding, self.dilation)
        return F.conv2d(dw, self.pointwise.weight, self.pointwise.bias, stride=1, padding=0, dilation=1, groups=1)