import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Size of the convolution kernel (height, width).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int or tuple, optional): Padding applied to the input. Defaults to 0.
        dilation (int or tuple, optional): Spacing between kernel elements. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. 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: tuple, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 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).
        """
        return self.conv2d(x)

# Test code
batch_size = 8
in_channels = 32
out_channels = 64
kernel_size = (5, 9)
width = 512
height = 512

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]  # Provide in_channels, out_channels, kernel_size for initialization

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

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

torch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  bool has_bias, int stride_h, int stride_w,
                                  int pad_h, int pad_w, int dil_h, int dil_w, int groups);
"""

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

#define TH 4
#define TW 32
#define KH 5
#define KW 9
#define IC 32
#define OCT 16
#define INH (TH + KH - 1)
#define INW (TW + KW - 1)

__global__ __launch_bounds__(128, 2)
void conv_32_64_5x9_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N, int H, int W, int OH, int OW,
                           int OC, bool has_bias) {
    const int tid = threadIdx.x;
    const int lx = tid & 31;       // 0..31
    const int ly = tid >> 5;       // 0..3
    const int ow0 = blockIdx.x * TW;
    const int oh0 = blockIdx.y * TH;
    const int oc_groups = OC / OCT;
    const int n = blockIdx.z / oc_groups;
    const int oc0 = (blockIdx.z - n * oc_groups) * OCT;

    __shared__ float xs[IC * INH * INW];
    __shared__ float ws[OCT * KH * KW];

    // Load the input patch for this (n, spatial tile).  The final W tile reads
    // a little past the image; those entries are zero and are only used by
    // threads whose outputs are out of bounds.
    const int total_x = IC * INH * INW;
    for (int idx = tid; idx < total_x; idx += 128) {
        int t = idx;
        const int ix = t % INW; t /= INW;
        const int iy = t % INH; t /= INH;
        const int c = t;
        const int ih = oh0 + iy;
        const int iw = ow0 + ix;
        float v = 0.0f;
        if (ih < H && iw < W) {
            v = x[((n * IC + c) * H + ih) * W + iw];
        }
        xs[idx] = v;
    }
    __syncthreads();

    float acc[OCT];
#pragma unroll
    for (int i = 0; i < OCT; ++i) acc[i] = 0.0f;

    const bool valid = (oh0 + ly < OH) && (ow0 + lx < OW);
    const int base_xy = ly * INW + lx;

#pragma unroll
    for (int c = 0; c < IC; ++c) {
        // Load this input-channel's 16x5x9 weights once per block into shared
        // memory.  This avoids the otherwise very expensive warp-uniform global
        // reloads of the same coefficients by all 128 threads.
        for (int idx = tid; idx < OCT * KH * KW; idx += 128) {
            const int oo = idx / (KH * KW);
            const int k = idx - oo * (KH * KW);
            ws[idx] = w[((oc0 + oo) * IC + c) * (KH * KW) + k];
        }
        __syncthreads();

        if (valid) {
            const float* __restrict__ xp = xs + c * (INH * INW) + base_xy;
#pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
#pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    const int k = kh * KW + kw;
                    const float xv = xp[kh * INW + kw];
#pragma unroll
                    for (int oo = 0; oo < OCT; ++oo) {
                        acc[oo] = fmaf(xv, ws[oo * (KH * KW) + k], acc[oo]);
                    }
                }
            }
        }
        __syncthreads();
    }

    if (valid) {
        const int oh = oh0 + ly;
        const int ow = ow0 + lx;
#pragma unroll
        for (int oo = 0; oo < OCT; ++oo) {
            float v = acc[oo];
            if (has_bias) v += b[oc0 + oo];
            y[((n * OC + oc0 + oo) * OH + oh) * OW + ow] = v;
        }
    }
}

torch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  bool has_bias, int stride_h, int stride_w,
                                  int pad_h, int pad_w, int dil_h, int dil_w, int groups) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 && weight.scalar_type() == torch::kFloat32, "float32 only");
    TORCH_CHECK(x.dim() == 4 && weight.dim() == 4, "expected NCHW input and OIHW weight");
    TORCH_CHECK(stride_h == 1 && stride_w == 1 && pad_h == 0 && pad_w == 0 && dil_h == 1 && dil_w == 1 && groups == 1,
                "this optimized kernel supports only stride=1, padding=0, dilation=1, groups=1");
    TORCH_CHECK(x.size(1) == 32 && weight.size(1) == 32 && weight.size(2) == 5 && weight.size(3) == 9,
                "this optimized kernel supports C=32 and kernel 5x9");
    TORCH_CHECK(weight.size(0) % 16 == 0, "out_channels must be divisible by 16");

    x = x.contiguous();
    weight = weight.contiguous();
    if (has_bias) bias = bias.contiguous();

    const int N = (int)x.size(0);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int OC = (int)weight.size(0);
    const int OH = H - 5 + 1;
    const int OW = W - 9 + 1;
    auto y = torch::empty({N, OC, OH, OW}, x.options());

    dim3 block(128);
    dim3 grid((OW + TW - 1) / TW, (OH + TH - 1) / TH, N * (OC / OCT));
    conv_32_64_5x9_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
        y.data_ptr<float>(), N, H, W, OH, OW, OC, has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        if isinstance(stride, tuple):
            self._stride_h, self._stride_w = int(stride[0]), int(stride[1])
        else:
            self._stride_h = self._stride_w = int(stride)
        if isinstance(padding, tuple):
            self._pad_h, self._pad_w = int(padding[0]), int(padding[1])
        else:
            self._pad_h = self._pad_w = int(padding)
        if isinstance(dilation, tuple):
            self._dil_h, self._dil_w = int(dilation[0]), int(dilation[1])
        else:
            self._dil_h = self._dil_w = int(dilation)
        self._groups = int(groups)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bias_t = self.conv2d.bias if self.conv2d.bias is not None else self.conv2d.weight
        return _ext.conv2d_forward_cuda(x, self.conv2d.weight, bias_t, self.conv2d.bias is not None,
                                        self._stride_h, self._stride_w, self._pad_h, self._pad_w,
                                        self._dil_h, self._dil_w, self._groups)