import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 2D convolution operation with an asymmetric input and a 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.
        dilation (int, 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: int, 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, 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
# smaller spatial dims
height = 512
width = 1024
in_channels = 64  # increased channels
out_channels = 128
kernel_size = 3
# asymmetric input: make width considerably larger than height

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_3x3_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                              int64_t stride, int64_t padding, int64_t dilation, int64_t groups);

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                             int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {
    return conv2d_3x3_cuda(x, w, b, stride, padding, dilation, groups);
}
'''

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

#define TILE_M 64
#define TILE_N 32
#define BK 8

__global__ __launch_bounds__(128, 2)
void conv3x3_s1p0_kernel(const float* __restrict__ x,
                         const float* __restrict__ w,
                         const float* __restrict__ bias,
                         float* __restrict__ y,
                         int N, int C, int H, int W,
                         int Kout, int OH, int OW,
                         int tiles_h, int has_bias) {
    __shared__ float As[TILE_M * BK];
    __shared__ float Bs[BK * TILE_N];

    const int tx = threadIdx.x; // 0..127
    const int tile_x = blockIdx.x;                  // 8-wide spatial tile in W
    const int by = blockIdx.y;
    const int n = by / tiles_h;
    const int tile_y = by - n * tiles_h;            // 8-high spatial tile in H
    const int n_tile = blockIdx.z;                  // 32 output-channel tile

    const int m_group = tx & 15;       // 16 groups in M
    const int n_group = tx >> 4;       // 8 groups in N, each computes 4 channels

    float acc[4][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) acc[i][j] = 0.0f;
    }

    const int ow_base = tile_x << 3;
    const int oh_base = tile_y << 3;
    const int oc_base = n_tile * TILE_N;
    const int Ktot = C * 9;

    for (int k0 = 0; k0 < Ktot; k0 += BK) {
        // Load im2col A tile: 64 output pixels x 8 reduction entries
        #pragma unroll
        for (int l = 0; l < 4; ++l) {
            int idx = tx + l * 128; // 0..511
            int mr = idx >> 3;      // / BK
            int kk = idx & 7;
            int kval = k0 + kk;
            int ic = kval / 9;
            int rs = kval - ic * 9;
            int rr = rs / 3;
            int ss = rs - rr * 3;
            int dy = mr >> 3;
            int dx = mr & 7;
            int oh = oh_base + dy;
            int ow = ow_base + dx;
            float v = 0.0f;
            if (oh < OH && ow < OW) {
                v = __ldg(x + (((n * C + ic) * H + (oh + rr)) * W + (ow + ss)));
            }
            As[mr * BK + kk] = v;
        }

        // Load B tile: 8 reduction entries x 32 output channels.
        #pragma unroll
        for (int l = 0; l < 2; ++l) {
            int idx = tx + l * 128; // 0..255
            int col = idx >> 3;     // output-channel within tile
            int kk = idx & 7;
            int oc = oc_base + col;
            float v = 0.0f;
            if (oc < Kout) {
                v = __ldg(w + oc * Ktot + (k0 + kk));
            }
            Bs[kk * TILE_N + col] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float a0 = As[(m_group +  0) * BK + kk];
            float a1 = As[(m_group + 16) * BK + kk];
            float a2 = As[(m_group + 32) * BK + kk];
            float a3 = As[(m_group + 48) * BK + kk];
            float b0 = Bs[kk * TILE_N + n_group * 4 + 0];
            float b1 = Bs[kk * TILE_N + n_group * 4 + 1];
            float b2 = Bs[kk * TILE_N + n_group * 4 + 2];
            float b3 = Bs[kk * TILE_N + n_group * 4 + 3];
            acc[0][0] = fmaf(a0, b0, acc[0][0]);
            acc[0][1] = fmaf(a0, b1, acc[0][1]);
            acc[0][2] = fmaf(a0, b2, acc[0][2]);
            acc[0][3] = fmaf(a0, b3, acc[0][3]);
            acc[1][0] = fmaf(a1, b0, acc[1][0]);
            acc[1][1] = fmaf(a1, b1, acc[1][1]);
            acc[1][2] = fmaf(a1, b2, acc[1][2]);
            acc[1][3] = fmaf(a1, b3, acc[1][3]);
            acc[2][0] = fmaf(a2, b0, acc[2][0]);
            acc[2][1] = fmaf(a2, b1, acc[2][1]);
            acc[2][2] = fmaf(a2, b2, acc[2][2]);
            acc[2][3] = fmaf(a2, b3, acc[2][3]);
            acc[3][0] = fmaf(a3, b0, acc[3][0]);
            acc[3][1] = fmaf(a3, b1, acc[3][1]);
            acc[3][2] = fmaf(a3, b2, acc[3][2]);
            acc[3][3] = fmaf(a3, b3, acc[3][3]);
        }
        __syncthreads();
    }

    #pragma unroll
    for (int ii = 0; ii < 4; ++ii) {
        int mr = m_group + ii * 16;
        int dy = mr >> 3;
        int dx = mr & 7;
        int oh = oh_base + dy;
        int ow = ow_base + dx;
        if (oh < OH && ow < OW) {
            #pragma unroll
            for (int jj = 0; jj < 4; ++jj) {
                int oc = oc_base + n_group * 4 + jj;
                if (oc < Kout) {
                    float v = acc[ii][jj] + (has_bias ? __ldg(bias + oc) : 0.0f);
                    y[((n * Kout + oc) * OH + oh) * OW + ow] = v;
                }
            }
        }
    }
}

// Generic but slower reference-style CUDA kernel for uncommon constructor arguments.
__global__ void conv2d_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,
                                      const float* __restrict__ bias, float* __restrict__ y,
                                      int N, int C, int H, int Wd, int Kout, int KH, int KW,
                                      int OH, int OW, int stride, int pad, int dil, int groups, int has_bias,
                                      long total) {
    long idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    long t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int oc = t % Kout;
    int n = t / Kout;
    int cpg = C / groups;
    int kpg = Kout / groups;
    int g = oc / kpg;
    int c0 = g * cpg;
    float sum = has_bias ? __ldg(bias + oc) : 0.0f;
    for (int icg = 0; icg < cpg; ++icg) {
        int ic = c0 + icg;
        for (int r = 0; r < KH; ++r) {
            int ih = oh * stride - pad + r * dil;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int s = 0; s < KW; ++s) {
                int iw = ow * stride - pad + s * dil;
                if ((unsigned)iw < (unsigned)Wd) {
                    sum = fmaf(__ldg(x + ((n*C + ic)*H + ih)*Wd + iw),
                               __ldg(w + (((oc*cpg + icg)*KH + r)*KW + s)), sum);
                }
            }
        }
    }
    y[idx] = sum;
}

torch::Tensor conv2d_3x3_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                              int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {
    x = x.contiguous();
    w = w.contiguous();
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int Wd = (int)x.size(3);
    const int Kout = (int)w.size(0);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);
    const int OH = (H + 2*(int)padding - (int)dilation*(KH-1) - 1) / (int)stride + 1;
    const int OW = (Wd + 2*(int)padding - (int)dilation*(KW-1) - 1) / (int)stride + 1;
    auto y = torch::empty({N, Kout, OH, OW}, x.options());
    const int has_bias = (b.numel() != 0);
    const float* bp = has_bias ? b.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kFloat && C == 64 && Kout == 128 && KH == 3 && KW == 3 &&
        stride == 1 && padding == 0 && dilation == 1 && groups == 1) {
        const int tiles_w = (OW + 7) >> 3;
        const int tiles_h = (OH + 7) >> 3;
        dim3 grid(tiles_w, N * tiles_h, (Kout + TILE_N - 1) / TILE_N);
        conv3x3_s1p0_kernel<<<grid, 128, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), bp, y.data_ptr<float>(),
                                                       N, C, H, Wd, Kout, OH, OW, tiles_h, has_bias);
    } else {
        long total = (long)N * Kout * OH * OW;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        conv2d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), bp, y.data_ptr<float>(),
                                                              N, C, H, Wd, Kout, KH, KW, OH, OW,
                                                              (int)stride, (int)padding, (int)dilation, (int)groups, has_bias, total);
    }
    return y;
}
'''

_ext = load_inline(
    name='conv2d_3x3_asym_ext_v1',
    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: int, 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, kernel_size), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self._stride = stride
        self._padding = padding
        self._dilation = dilation
        self._groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv2d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv2d_forward(x, self.conv2d.weight, b, self._stride, self._padding, self._dilation, self._groups)