import torch
import torch.nn as nn

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

    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 (tuple, optional): Padding applied to the input (top/bottom, left/right). Defaults to (0, 0).
        dilation (tuple, optional): Spacing between kernel elements (height, width). Defaults to (1, 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: tuple = (0, 0), dilation: tuple = (1, 1), bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, 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
stride = 1
padding = (2, 4)
dilation = (2, 3)

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 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,
                                  int64_t stride, int64_t pad_h, int64_t pad_w,
                                  int64_t dil_h, int64_t dil_w);
'''

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

// Specialized kernel for the benchmark case:
// N=8, IC=32, OC=64, KH=5, KW=9, stride=1, padding=(2,4), dilation=(2,3).
// One block computes 32 spatial positions and all 64 output channels.  The input
// patch values are loaded once into shared memory and reused by all channel warps.
template <bool HAS_BIAS>
__global__ __launch_bounds__(256, 2)
void conv5x9_dilated_spec_kernel(const float* __restrict__ x,
                                  const float* __restrict__ w,
                                  const float* __restrict__ bias,
                                  float* __restrict__ out,
                                  int H, int W, int Hout, int Wout, int M) {
    constexpr int BM = 32;
    constexpr int BN = 64;
    constexpr int BK = 32;
    constexpr int KTOTAL = 32 * 5 * 9; // 1440, divisible by BK
    __shared__ float As[BK * BM];      // [k][m]
    __shared__ float Bs[BK * BN];      // [k][oc]

    const int tx = threadIdx.x; // 0..31: spatial lane
    const int ty = threadIdx.y; // 0..7: output-channel warp group
    const int tid = ty * 32 + tx;
    const int m = blockIdx.x * BM + tx;

    float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f;
    float acc4 = 0.f, acc5 = 0.f, acc6 = 0.f, acc7 = 0.f;

    int n = 0, oh = 0, ow = 0;
    if (m < M) {
        const int plane = Hout * Wout;
        n = m / plane;
        const int r = m - n * plane;
        oh = r / Wout;
        ow = r - oh * Wout;
    }

    #pragma unroll
    for (int kb = 0; kb < KTOTAL; kb += BK) {
        // load A tile: 32 spatial positions x 32 reduction entries
        #pragma unroll
        for (int idx = tid; idx < BM * BK; idx += 256) {
            const int kk_in = idx >> 5;       // / BM
            const int mm = idx & 31;          // % BM
            const int gm = blockIdx.x * BM + mm;
            float v = 0.f;
            if (gm < M) {
                const int plane = Hout * Wout;
                const int nn = gm / plane;
                const int rr = gm - nn * plane;
                const int ooh = rr / Wout;
                const int oow = rr - ooh * Wout;
                const int kglob = kb + kk_in;
                const int ic = kglob / 45;
                const int rem = kglob - ic * 45;
                const int kh = rem / 9;
                const int kw = rem - kh * 9;
                const int ih = ooh - 2 + kh * 2;
                const int iw = oow - 4 + kw * 3;
                if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) {
                    v = x[((nn * 32 + ic) * H + ih) * W + iw];
                }
            }
            As[kk_in * BM + mm] = v;
        }

        // load B tile: 32 reduction entries x 64 output channels
        #pragma unroll
        for (int idx = tid; idx < BK * BN; idx += 256) {
            const int kk_in = idx >> 6;       // / 64
            const int oc = idx & 63;          // % 64
            Bs[kk_in * BN + oc] = w[oc * KTOTAL + (kb + kk_in)];
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            const float a = As[kk * BM + tx];
            acc0 = fmaf(a, Bs[kk * BN + (ty +  0)], acc0);
            acc1 = fmaf(a, Bs[kk * BN + (ty +  8)], acc1);
            acc2 = fmaf(a, Bs[kk * BN + (ty + 16)], acc2);
            acc3 = fmaf(a, Bs[kk * BN + (ty + 24)], acc3);
            acc4 = fmaf(a, Bs[kk * BN + (ty + 32)], acc4);
            acc5 = fmaf(a, Bs[kk * BN + (ty + 40)], acc5);
            acc6 = fmaf(a, Bs[kk * BN + (ty + 48)], acc6);
            acc7 = fmaf(a, Bs[kk * BN + (ty + 56)], acc7);
        }
        __syncthreads();
    }

    if (m < M) {
        const int plane = Hout * Wout;
        const int base0 = (n * 64 * plane) + oh * Wout + ow;
        const int oc0 = ty + 0;
        if constexpr (HAS_BIAS) {
            acc0 += bias[oc0 +  0]; acc1 += bias[oc0 +  8];
            acc2 += bias[oc0 + 16]; acc3 += bias[oc0 + 24];
            acc4 += bias[oc0 + 32]; acc5 += bias[oc0 + 40];
            acc6 += bias[oc0 + 48]; acc7 += bias[oc0 + 56];
        }
        out[base0 + (oc0 +  0) * plane] = acc0;
        out[base0 + (oc0 +  8) * plane] = acc1;
        out[base0 + (oc0 + 16) * plane] = acc2;
        out[base0 + (oc0 + 24) * plane] = acc3;
        out[base0 + (oc0 + 32) * plane] = acc4;
        out[base0 + (oc0 + 40) * plane] = acc5;
        out[base0 + (oc0 + 48) * plane] = acc6;
        out[base0 + (oc0 + 56) * plane] = acc7;
    }
}

// A simple generic correctness kernel for non-benchmark constructor values.
__global__ void conv2d_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ bias,
                                      float* __restrict__ out,
                                      int N, int C, int H, int W, int OC,
                                      int KH, int KW, int Hout, int Wout,
                                      int stride, int pad_h, int pad_w, int dil_h, int dil_w,
                                      int has_bias, long total) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % Wout;
    long t = idx / Wout;
    int oh = t % Hout;
    t /= Hout;
    int oc = t % OC;
    int n = t / OC;
    float sum = has_bias ? bias[oc] : 0.f;
    for (int c = 0; c < C; ++c) {
        for (int kh = 0; kh < KH; ++kh) {
            int ih = oh * stride - pad_h + kh * dil_h;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < KW; ++kw) {
                int iw = ow * stride - pad_w + kw * dil_w;
                if ((unsigned)iw < (unsigned)W) {
                    sum = fmaf(x[((n*C + c)*H + ih)*W + iw],
                               w[((oc*C + c)*KH + kh)*KW + kw], sum);
                }
            }
        }
    }
    out[idx] = sum;
}

torch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  int64_t stride, int64_t pad_h, int64_t pad_w,
                                  int64_t dil_h, int64_t dil_w) {
    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 OC = (int)weight.size(0);
    const int KH = (int)weight.size(2);
    const int KW = (int)weight.size(3);
    const int Hout = (H + 2*(int)pad_h - (int)dil_h * (KH - 1) - 1) / (int)stride + 1;
    const int Wout = (W + 2*(int)pad_w - (int)dil_w * (KW - 1) - 1) / (int)stride + 1;
    auto out = torch::empty({N, OC, Hout, Wout}, x.options());
    const bool has_bias = bias.numel() != 0;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kFloat && C == 32 && OC == 64 && KH == 5 && KW == 9 &&
        stride == 1 && pad_h == 2 && pad_w == 4 && dil_h == 2 && dil_w == 3) {
        const int M = N * Hout * Wout;
        dim3 block(32, 8, 1);
        dim3 grid((M + 31) / 32, 1, 1);
        if (has_bias) {
            conv5x9_dilated_spec_kernel<true><<<grid, block, 0, stream>>>(
                x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
                out.data_ptr<float>(), H, W, Hout, Wout, M);
        } else {
            conv5x9_dilated_spec_kernel<false><<<grid, block, 0, stream>>>(
                x.data_ptr<float>(), weight.data_ptr<float>(), nullptr,
                out.data_ptr<float>(), H, W, Hout, Wout, M);
        }
    } else {
        long total = (long)N * OC * Hout * Wout;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        conv2d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            out.data_ptr<float>(), N, C, H, W, OC, KH, KW, Hout, Wout,
            (int)stride, (int)pad_h, (int)pad_w, (int)dil_h, (int)dil_w, has_bias ? 1 : 0, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: int = 1, padding: tuple = (0, 0), dilation: tuple = (1, 1),
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
                                padding=padding, dilation=dilation, bias=bias)
        self.stride = int(stride)
        self.padding = (int(padding[0]), int(padding[1]))
        self.dilation = (int(dilation[0]), int(dilation[1]))
        self.register_buffer('_empty_bias', torch.empty(0), persistent=False)

    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_cuda(x, self.conv2d.weight, b,
                                        self.stride, self.padding[0], self.padding[1],
                                        self.dilation[0], self.dilation[1])