import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise 2D convolution with asymmetric input and square kernel.

    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 square 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, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=(kernel_size, 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_in, width_in).

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

# Test code
batch_size = 64
in_channels = 128
out_channels = 128
kernel_size = 3
width_in = 512
height_in = 256
stride = 1
padding = 0

def get_inputs():
    x = torch.rand(batch_size, in_channels, height_in, width_in)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding]

import os
os.makedirs('/scratch/abaner/tmp', exist_ok=True)
os.environ['TMPDIR'] = '/scratch/abaner/tmp'
os.environ['TEMP'] = '/scratch/abaner/tmp'
os.environ['TMP'] = '/scratch/abaner/tmp'

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

cpp_src = r"""
#include <torch/extension.h>
torch::Tensor depthwise3x3_forward(torch::Tensor x, torch::Tensor weight);
"""

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

#define TILE_H 16

__global__ __launch_bounds__(256, 3)
void dw3x3_smem_kernel(const float* __restrict__ x,
                       const float* __restrict__ w,
                       float* __restrict__ y,
                       int C, int H, int W, int Ho, int Wo) {
    extern __shared__ float sh[];
    const int nc = blockIdx.x;
    const int hbase = blockIdx.y * TILE_H;
    const int tid = threadIdx.x;
    const int c = nc % C;
    const int n = nc / C;

    const float* __restrict__ xp = x + ((long long)n * C + c) * (long long)H * W;
    float* __restrict__ yp = y + ((long long)n * C + c) * (long long)Ho * Wo;

    int out_rows = Ho - hbase;
    if (out_rows > TILE_H) out_rows = TILE_H;
    const int load_rows = out_rows + 2;

    for (int idx = tid; idx < load_rows * W; idx += 256) {
        sh[idx] = xp[(long long)hbase * W + idx];
    }
    __syncthreads();

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

    #pragma unroll
    for (int rr = 0; rr < TILE_H; ++rr) {
        if (rr >= out_rows) return;
        const float* row0 = sh + rr * W;
        const float* row1 = row0 + W;
        const float* row2 = row1 + W;
        float* out = yp + (long long)(hbase + rr) * Wo;
        for (int ow = tid; ow < Wo; ow += 256) {
            float acc = row0[ow] * w00;
            acc = fmaf(row0[ow + 1], w01, acc);
            acc = fmaf(row0[ow + 2], w02, acc);
            acc = fmaf(row1[ow],     w10, acc);
            acc = fmaf(row1[ow + 1], w11, acc);
            acc = fmaf(row1[ow + 2], w12, acc);
            acc = fmaf(row2[ow],     w20, acc);
            acc = fmaf(row2[ow + 1], w21, acc);
            acc = fmaf(row2[ow + 2], w22, acc);
            out[ow] = acc;
        }
    }
}

torch::Tensor depthwise3x3_forward(torch::Tensor x, torch::Tensor weight) {
    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 Ho = H - 2;
    const int Wo = W - 2;
    auto y = torch::empty({N, C, Ho, Wo}, x.options());
    const dim3 block(256);
    const dim3 grid(N * C, (Ho + TILE_H - 1) / TILE_H);
    const size_t smem = (TILE_H + 2) * (size_t)W * sizeof(float);
    dw3x3_smem_kernel<<<grid, block, smem, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>(), C, H, W, Ho, Wo);
    return y;
}
"""

_ext = load_inline(name="dw3x3_final_smem16", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["depthwise3x3_forward"],
                   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, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=(kernel_size, kernel_size), stride=stride, padding=padding, groups=in_channels, bias=bias)

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