import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Tuple of two integers representing the height and width of the convolution kernel.
        stride (tuple, optional): Tuple of two integers representing the stride in the height and width dimensions. Defaults to (1, 1).
        padding (tuple, optional): Tuple of two integers representing the padding in the height and width dimensions. Defaults to (0, 0).
        dilation (tuple, optional): Tuple of two integers representing the dilation in the height and width dimensions. Defaults to (1, 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: tuple = (1, 1), padding: tuple = (0, 0), dilation: tuple = (1, 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 = 64
out_channels = 128
kernel_size = (5, 7)
height = 512
width = 256

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(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias,
                             int64_t stride_h, int64_t stride_w, int64_t pad_h, int64_t pad_w,
                             int64_t dil_h, int64_t dil_w, int64_t groups);
"""

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

// Specialized kernel for the benchmark shape:
// x: [8,64,512,256], w: [128,64,5,7], stride=1, padding=0, dilation=1, groups=1
// y: [8,128,508,250]
// Tile: 64 output channels x 32 spatial positions.  Each thread computes 4 OC x 2 spatial.
__global__ __launch_bounds__(256, 2)
void conv2d_5x7_s1p0_kernel(const float* __restrict__ x,
                            const float* __restrict__ w,
                            const float* __restrict__ b,
                            float* __restrict__ y,
                            const int N) {
    constexpr int IC = 64;
    constexpr int OC = 128;
    constexpr int H = 512;
    constexpr int W = 256;
    constexpr int OH = 508;
    constexpr int OW = 250;
    constexpr int SP = OH * OW;       // 127000
    constexpr int KRS = 35;           // 5*7
    constexpr int MTILE = 32;
    constexpr int OCTILE = 64;

    __shared__ float As[KRS][MTILE];      // [kernel-position][spatial]
    __shared__ float Ws[OCTILE][36];      // [oc][kernel-position], padded

    const int tx = threadIdx.x;  // 0..15 spatial lane
    const int ty = threadIdx.y;  // 0..15 oc lane
    const int tid = ty * 16 + tx;

    const int oc_base = blockIdx.x * OCTILE;
    const int m_base  = blockIdx.y * MTILE;

    float acc00 = 0.0f, acc01 = 0.0f;
    float acc10 = 0.0f, acc11 = 0.0f;
    float acc20 = 0.0f, acc21 = 0.0f;
    float acc30 = 0.0f, acc31 = 0.0f;

    #pragma unroll 1
    for (int ic = 0; ic < IC; ++ic) {
        // Load input tile: 35 x 32 = 1120 floats, 5 per 256-thread block (last partly inactive)
        #pragma unroll
        for (int l = 0; l < 5; ++l) {
            int idx = tid + l * 256;
            if (idx < KRS * MTILE) {
                int k = idx / MTILE;
                int r = idx - k * MTILE;
                int m = m_base + r;
                float v = 0.0f;
                if (m < N * SP) {
                    int spatial = m % SP;
                    int n = m / SP;
                    int oh = spatial / OW;
                    int ow = spatial - oh * OW;
                    int kh = k / 7;
                    int kw = k - kh * 7;
                    int in_idx = ((n * IC + ic) * H + (oh + kh)) * W + (ow + kw);
                    v = __ldg(x + in_idx);
                }
                As[k][r] = v;
            }
        }

        // Load weight tile: 64 x 35 = 2240 floats, 9 per thread (last partly inactive)
        #pragma unroll
        for (int l = 0; l < 9; ++l) {
            int idx = tid + l * 256;
            if (idx < OCTILE * KRS) {
                int oc_l = idx / KRS;
                int k = idx - oc_l * KRS;
                int oc = oc_base + oc_l;
                float v = 0.0f;
                if (oc < OC) {
                    v = __ldg(w + oc * (IC * KRS) + ic * KRS + k);
                }
                Ws[oc_l][k] = v;
            }
        }
        __syncthreads();

        const int r0 = tx;
        const int r1 = tx + 16;
        const int o0 = ty;
        const int o1 = ty + 16;
        const int o2 = ty + 32;
        const int o3 = ty + 48;

        #pragma unroll
        for (int k = 0; k < KRS; ++k) {
            float a0 = As[k][r0];
            float a1 = As[k][r1];
            float w0 = Ws[o0][k];
            float w1 = Ws[o1][k];
            float w2 = Ws[o2][k];
            float w3 = Ws[o3][k];
            acc00 = fmaf(a0, w0, acc00); acc01 = fmaf(a1, w0, acc01);
            acc10 = fmaf(a0, w1, acc10); acc11 = fmaf(a1, w1, acc11);
            acc20 = fmaf(a0, w2, acc20); acc21 = fmaf(a1, w2, acc21);
            acc30 = fmaf(a0, w3, acc30); acc31 = fmaf(a1, w3, acc31);
        }
        __syncthreads();
    }

    int m0 = m_base + tx;
    int m1 = m_base + tx + 16;

    int oc0 = oc_base + ty;
    int oc1 = oc0 + 16;
    int oc2 = oc0 + 32;
    int oc3 = oc0 + 48;

    if (m0 < N * SP) {
        int n = m0 / SP;
        int spatial = m0 - n * SP;
        int base = n * OC * SP + spatial;
        float bb0 = (b != nullptr ? __ldg(b + oc0) : 0.0f);
        float bb1 = (b != nullptr ? __ldg(b + oc1) : 0.0f);
        float bb2 = (b != nullptr ? __ldg(b + oc2) : 0.0f);
        float bb3 = (b != nullptr ? __ldg(b + oc3) : 0.0f);
        y[base + oc0 * SP] = acc00 + bb0;
        y[base + oc1 * SP] = acc10 + bb1;
        y[base + oc2 * SP] = acc20 + bb2;
        y[base + oc3 * SP] = acc30 + bb3;
    }
    if (m1 < N * SP) {
        int n = m1 / SP;
        int spatial = m1 - n * SP;
        int base = n * OC * SP + spatial;
        float bb0 = (b != nullptr ? __ldg(b + oc0) : 0.0f);
        float bb1 = (b != nullptr ? __ldg(b + oc1) : 0.0f);
        float bb2 = (b != nullptr ? __ldg(b + oc2) : 0.0f);
        float bb3 = (b != nullptr ? __ldg(b + oc3) : 0.0f);
        y[base + oc0 * SP] = acc01 + bb0;
        y[base + oc1 * SP] = acc11 + bb1;
        y[base + oc2 * SP] = acc21 + bb2;
        y[base + oc3 * SP] = acc31 + bb3;
    }
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias,
                             int64_t stride_h, int64_t stride_w, int64_t pad_h, int64_t pad_w,
                             int64_t dil_h, int64_t dil_w, int64_t groups) {
    const int64_t N = x.size(0);
    const int64_t H = x.size(2);
    const int64_t W = x.size(3);
    const int64_t OC = weight.size(0);
    const int64_t KH = weight.size(2);
    const int64_t KW = weight.size(3);
    const int64_t OH = (H + 2 * pad_h - dil_h * (KH - 1) - 1) / stride_h + 1;
    const int64_t OW = (W + 2 * pad_w - dil_w * (KW - 1) - 1) / stride_w + 1;

    auto y = torch::empty({N, OC, OH, OW}, x.options());

    dim3 block(16, 16);
    dim3 grid((unsigned)((OC + 63) / 64), (unsigned)((N * OH * OW + 31) / 32));
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    conv2d_5x7_s1p0_kernel<<<grid, block>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr,
                                             y.data_ptr<float>(), (int)N);
    return y;
}
"""

_ext = load_inline(
    name="conv2d_5x7_s1p0_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv2d_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: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), dilation: tuple = (1, 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)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.has_bias = bias
        self.register_buffer("_empty_bias", torch.empty(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv2d.bias if self.conv2d.bias is not None else self._empty_bias
        return _ext.conv2d_forward(x, self.conv2d.weight, b, self.conv2d.bias is not None,
                                   self.stride[0], self.stride[1], self.padding[0], self.padding[1],
                                   self.dilation[0], self.dilation[1], self.groups)