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 os
os.environ['TORCH_EXTENSIONS_DIR']='/dev/shm/torch_ext_final'
os.environ['TMPDIR']='/dev/shm'
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor dw_sep_forward(torch::Tensor x, torch::Tensor dw_w, torch::Tensor pw_w,
                             torch::Tensor dw_b, torch::Tensor pw_b,
                             int stride, int padding, int dilation);
'''

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

__global__ void depthwise_kernel(const float* __restrict__ x,
                                 const float* __restrict__ w,
                                 const float* __restrict__ b,
                                 float* __restrict__ d,
                                 int N, int C, int H, int W,
                                 int KH, int KW, int OH, int OW,
                                 int stride, int padding, int dilation,
                                 int has_b) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * OH * OW;
    if (idx >= total) return;
    int t = idx;
    int ow = t % OW; t /= OW;
    int oh = t % OH; t /= OH;
    int c = t % C;
    int n = t / C;
    float acc = 0.0f;
    #pragma unroll
    for (int r = 0; r < 3; ++r) {
        int ih = oh * stride + r * dilation - padding;
        if (r < KH && (unsigned)ih < (unsigned)H) {
            #pragma unroll
            for (int s = 0; s < 3; ++s) {
                int iw = ow * stride + s * dilation - padding;
                if (s < KW && (unsigned)iw < (unsigned)W) {
                    acc = fmaf(x[((n * C + c) * H + ih) * W + iw],
                               w[(c * KH + r) * KW + s], acc);
                }
            }
        }
    }
    if (has_b) acc += b[c];
    d[idx] = acc;
}

__global__ void add_pw_bias(float* out, const float* b, int N, int OC, int HW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * OC * HW;
    if (idx < total) {
        int oc = (idx / HW) % OC;
        out[idx] += b[oc];
    }
}

torch::Tensor dw_sep_forward(torch::Tensor x, torch::Tensor dw_w, torch::Tensor pw_w,
                             torch::Tensor dw_b, torch::Tensor pw_b,
                             int stride, int padding, int dilation) {
    x = x.contiguous();
    dw_w = dw_w.contiguous();
    pw_w = pw_w.contiguous();
    int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    int KH = dw_w.size(2), KW = dw_w.size(3), OC = pw_w.size(0);
    int OH = (H + 2 * padding - dilation * (KH - 1) - 1) / stride + 1;
    int OW = (W + 2 * padding - dilation * (KW - 1) - 1) / stride + 1;
    int HW = OH * OW;
    auto d = torch::empty({N, C, OH, OW}, x.options());
    auto out = torch::empty({N, OC, OH, OW}, x.options());

    int total = N * C * HW;
    int hdb = dw_b.numel() > 0;
    depthwise_kernel<<<(total + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), dw_w.data_ptr<float>(),
        hdb ? dw_b.contiguous().data_ptr<float>() : nullptr,
        d.data_ptr<float>(), N, C, H, W, KH, KW, OH, OW, stride, padding, dilation, hdb);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, at::cuda::getCurrentCUDAStream());
    cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);
    float alpha = 1.0f, beta = 0.0f;
    long long strideA = (long long)C * HW;
    long long strideB = 0;
    long long strideC = (long long)OC * HW;
    cublasStatus_t st = cublasSgemmStridedBatched(
        handle, CUBLAS_OP_N, CUBLAS_OP_N,
        HW, OC, C,
        &alpha,
        d.data_ptr<float>(), HW, strideA,
        pw_w.data_ptr<float>(), C, strideB,
        &beta,
        out.data_ptr<float>(), HW, strideC,
        N);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublasSgemmStridedBatched failed");

    if (pw_b.numel() > 0) {
        int tot = N * OC * HW;
        add_pw_bias<<<(tot + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
            out.data_ptr<float>(), pw_b.contiguous().data_ptr<float>(), N, OC, HW);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(name='dwsep_cublas_final', cpp_sources=[cpp_src], cuda_sources=[cuda_src],
                   functions=['dw_sep_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'],
                   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
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        e = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.dw_sep_forward(x, self.depthwise.weight, self.pointwise.weight,
                                   self.depthwise.bias if self.depthwise.bias is not None else e,
                                   self.pointwise.bias if self.pointwise.bias is not None else e,
                                   self.stride, self.padding, self.dilation)