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
from torch.utils.cpp_extension import load_inline

cpp_sources = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDABlas.h>
#include <cublas_v2.h>

void depthwise_kernel_launcher(const float*, const float*, const float*, float*,
                               int, int, int, int, int, int, int, int, int, int, bool, cudaStream_t);
void add_bias_launcher(float*, const float*, int, long, long, cudaStream_t);

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

at::Tensor depthwise_sep_forward(at::Tensor x, at::Tensor dw_w, at::Tensor pw_w,
                                 at::Tensor dw_b, at::Tensor pw_b,
                                 int64_t stride, int64_t padding, int64_t dilation) {
    CHECK_CUDA(x); CHECK_CUDA(dw_w); CHECK_CUDA(pw_w); CHECK_CUDA(dw_b); CHECK_CUDA(pw_b);
    CHECK_FLOAT(x); CHECK_FLOAT(dw_w); CHECK_FLOAT(pw_w); CHECK_FLOAT(dw_b); CHECK_FLOAT(pw_b);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(dw_w); CHECK_CONTIGUOUS(pw_w); CHECK_CONTIGUOUS(dw_b); CHECK_CONTIGUOUS(pw_b);

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

    auto inter = at::empty({N, C, Ho, Wo}, x.options());
    auto out = at::empty({N, OC, Ho, Wo}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float* db = dw_b.numel() ? dw_b.data_ptr<float>() : nullptr;
    depthwise_kernel_launcher(x.data_ptr<float>(), dw_w.data_ptr<float>(), db, inter.data_ptr<float>(),
                              N, C, H, W, K, (int)stride, (int)padding, (int)dilation, Ho, Wo,
                              dw_b.numel() != 0, stream);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, stream);
    cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH);

    const float alpha = 1.0f, beta = 0.0f;
    long long S = (long long)Ho * (long long)Wo;
    cublasStatus_t stat = cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N,
        (int)S, OC, C,
        &alpha,
        inter.data_ptr<float>(), (int)S, (long long)C * S,
        pw_w.data_ptr<float>(), C, 0LL,
        &beta,
        out.data_ptr<float>(), (int)S, (long long)OC * S,
        N);
    TORCH_CHECK(stat == CUBLAS_STATUS_SUCCESS, "cublasSgemmStridedBatched failed");

    if (pw_b.numel()) {
        add_bias_launcher(out.data_ptr<float>(), pw_b.data_ptr<float>(), OC, (long)S,
                          (long)((long long)N * OC * S), stream);
    }
    return out;
}
'''

cuda_sources = r'''
#include <cuda_runtime.h>

__global__ void dw3_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N, int C, int H, int W, bool has_bias) {
    const int ITEMS = 4;
    int plane = blockIdx.y;
    int n = plane / C;
    int c = plane - n * C;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * ITEMS;
    int S = H * W;

    const float* xb = x + ((long)n * C + c) * S;
    const float* wb = w + (long)c * 9;
    float* yb = y + ((long)n * C + c) * S;

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

    #pragma unroll
    for (int it = 0; it < ITEMS; ++it) {
        int idx = base + it;
        if (idx < S) {
            int h = idx / W;
            int col = idx - h * W;
            float acc = has_bias ? b[c] : 0.0f;

            if (h > 0) {
                int r = idx - W;
                if (col > 0)   acc = fmaf(xb[r - 1], w00, acc);
                acc = fmaf(xb[r], w01, acc);
                if (col < W-1) acc = fmaf(xb[r + 1], w02, acc);
            }
            if (col > 0)       acc = fmaf(xb[idx - 1], w10, acc);
            acc = fmaf(xb[idx], w11, acc);
            if (col < W-1)     acc = fmaf(xb[idx + 1], w12, acc);
            if (h < H-1) {
                int r = idx + W;
                if (col > 0)   acc = fmaf(xb[r - 1], w20, acc);
                acc = fmaf(xb[r], w21, acc);
                if (col < W-1) acc = fmaf(xb[r + 1], w22, acc);
            }
            yb[idx] = 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 stride, int pad, int dil, int Ho, int Wo,
                                  bool has_bias) {
    const int ITEMS = 2;
    int plane = blockIdx.y;
    int n = plane / C;
    int c = plane - n * C;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * ITEMS;
    int S = Ho * Wo;
    const float* xb = x + ((long)n * C + c) * H * W;
    const float* wb = w + (long)c * K * K;
    float* yb = y + ((long)n * C + c) * S;

    #pragma unroll
    for (int it = 0; it < ITEMS; ++it) {
        int idx = base + it;
        if (idx < S) {
            int oh = idx / Wo;
            int ow = idx - oh * Wo;
            int ih0 = oh * stride - pad;
            int iw0 = ow * stride - pad;
            float acc = has_bias ? b[c] : 0.0f;
            for (int kh = 0; kh < K; ++kh) {
                int ih = ih0 + kh * dil;
                if ((unsigned)ih < (unsigned)H) {
                    for (int kw = 0; kw < K; ++kw) {
                        int iw = iw0 + kw * dil;
                        if ((unsigned)iw < (unsigned)W) {
                            acc = fmaf(xb[ih * W + iw], wb[kh * K + kw], acc);
                        }
                    }
                }
            }
            yb[idx] = acc;
        }
    }
}

__global__ void add_bias_kernel(float* __restrict__ y, const float* __restrict__ b,
                                int OC, long S, long total) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        int oc = (int)((idx / S) % OC);
        y[idx] += b[oc];
    }
}

void depthwise_kernel_launcher(const float* x, const float* w, const float* b, float* y,
                               int N, int C, int H, int W, int K, int stride, int pad, int dil,
                               int Ho, int Wo, bool has_bias, cudaStream_t stream) {
    dim3 block(256);
    int S = Ho * Wo;
    if (K == 3 && stride == 1 && pad == 1 && dil == 1 && Ho == H && Wo == W) {
        dim3 grid((S + (int)block.x * 4 - 1) / ((int)block.x * 4), N * C);
        dw3_kernel<<<grid, block, 0, stream>>>(x, w, b, y, N, C, H, W, has_bias);
    } else {
        dim3 grid((S + (int)block.x * 2 - 1) / ((int)block.x * 2), N * C);
        dw_generic_kernel<<<grid, block, 0, stream>>>(x, w, b, y, N, C, H, W, K,
                                                      stride, pad, dil, Ho, Wo, has_bias);
    }
}

void add_bias_launcher(float* y, const float* b, int OC, long S, long total, cudaStream_t stream) {
    dim3 block(256);
    dim3 grid((total + 255) / 256);
    add_bias_kernel<<<grid, block, 0, stream>>>(y, b, OC, S, total);
}
'''

_ext = load_inline(
    name='dwsep_h100_submit_v1',
    cpp_sources=cpp_sources,
    cuda_sources=cuda_sources,
    functions=['depthwise_sep_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    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 = stride
        self.padding = padding
        self.dilation = dilation
        self._empty = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._empty is None or self._empty.device != x.device or self._empty.dtype != x.dtype:
            self._empty = torch.empty(0, device=x.device, dtype=x.dtype)
        db = self.depthwise.bias if self.depthwise.bias is not None else self._empty
        pb = self.pointwise.bias if self.pointwise.bias is not None else self._empty
        return _ext.depthwise_sep_forward(x, self.depthwise.weight, self.pointwise.weight,
                                          db, pb, self.stride, self.padding, self.dilation)