import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 2D convolution operation with a square 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.
        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 = 16
in_channels = 16
out_channels = 128
kernel_size = 3
width = 1024
height = 1024

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 w, torch::Tensor b,
                             int64_t stride, int64_t padding, int64_t dilation,
                             int64_t groups, bool has_bias);
"""

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")

__global__ void conv3x3_c16_oc8_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ b,
                                        float* __restrict__ y,
                                        int N, int H, int W, int O,
                                        int OH, int OW, int has_bias) {
    constexpr int C = 16;
    constexpr int OCPT = 8;
    constexpr int ITEMS = 2;
    int oc_base = blockIdx.y * OCPT;
    int tid = threadIdx.x;
    int P = N * OH * OW;
    int p0 = (blockIdx.x * (blockDim.x * ITEMS)) + tid;
    int p1 = p0 + blockDim.x;
    int64_t S = (int64_t)OH * OW;
    int64_t HW = (int64_t)H * W;

    if (p1 < P) {
        int n0 = p0 / (int)S;
        int r0 = p0 - n0 * (int)S;
        int oh0 = r0 / OW;
        int ow0 = r0 - oh0 * OW;
        int n1 = p1 / (int)S;
        int r1 = p1 - n1 * (int)S;
        int oh1 = r1 / OW;
        int ow1 = r1 - oh1 * OW;

        float a0[OCPT], a1[OCPT];
        #pragma unroll
        for (int j = 0; j < OCPT; ++j) {
            float bv = (has_bias && (oc_base + j) < O) ? __ldg(b + oc_base + j) : 0.0f;
            a0[j] = bv; a1[j] = bv;
        }

        int64_t in_base0 = ((int64_t)n0 * C) * HW + (int64_t)oh0 * W + ow0;
        int64_t in_base1 = ((int64_t)n1 * C) * HW + (int64_t)oh1 * W + ow1;
        #pragma unroll
        for (int c = 0; c < C; ++c) {
            int64_t xb0 = in_base0 + (int64_t)c * HW;
            int64_t xb1 = in_base1 + (int64_t)c * HW;
            #pragma unroll
            for (int kk = 0; kk < 9; ++kk) {
                int ky = kk / 3;
                int kx = kk - ky * 3;
                float v0 = __ldg(x + xb0 + (int64_t)ky * W + kx);
                float v1 = __ldg(x + xb1 + (int64_t)ky * W + kx);
                int woff = c * 9 + kk;
                #pragma unroll
                for (int j = 0; j < OCPT; ++j) {
                    int oc = oc_base + j;
                    if (oc < O) {
                        float ww = __ldg(w + (int64_t)oc * C * 9 + woff);
                        a0[j] = fmaf(v0, ww, a0[j]);
                        a1[j] = fmaf(v1, ww, a1[j]);
                    }
                }
            }
        }
        #pragma unroll
        for (int j = 0; j < OCPT; ++j) {
            int oc = oc_base + j;
            if (oc < O) {
                y[((int64_t)n0 * O + oc) * S + r0] = a0[j];
                y[((int64_t)n1 * O + oc) * S + r1] = a1[j];
            }
        }
    } else if (p0 < P) {
        int n0 = p0 / (int)S;
        int r0 = p0 - n0 * (int)S;
        int oh0 = r0 / OW;
        int ow0 = r0 - oh0 * OW;
        float a0[OCPT];
        #pragma unroll
        for (int j = 0; j < OCPT; ++j) {
            a0[j] = (has_bias && (oc_base + j) < O) ? __ldg(b + oc_base + j) : 0.0f;
        }
        int64_t in_base0 = ((int64_t)n0 * C) * HW + (int64_t)oh0 * W + ow0;
        #pragma unroll
        for (int c = 0; c < C; ++c) {
            int64_t xb0 = in_base0 + (int64_t)c * HW;
            #pragma unroll
            for (int kk = 0; kk < 9; ++kk) {
                int ky = kk / 3;
                int kx = kk - ky * 3;
                float v0 = __ldg(x + xb0 + (int64_t)ky * W + kx);
                int woff = c * 9 + kk;
                #pragma unroll
                for (int j = 0; j < OCPT; ++j) {
                    int oc = oc_base + j;
                    if (oc < O) {
                        float ww = __ldg(w + (int64_t)oc * C * 9 + woff);
                        a0[j] = fmaf(v0, ww, a0[j]);
                    }
                }
            }
        }
        #pragma unroll
        for (int j = 0; j < OCPT; ++j) {
            int oc = oc_base + j;
            if (oc < O) y[((int64_t)n0 * O + oc) * S + r0] = a0[j];
        }
    }
}

__global__ void conv2d_general_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ b,
                                      float* __restrict__ y,
                                      int N, int C, int H, int Wd, int O, int K,
                                      int OH, int OW, int stride, int pad,
                                      int dil, int groups, int has_bias) {
    int64_t total = (int64_t)N * O * OH * OW;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    int64_t t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int oc = t % O;
    int n = t / O;
    int ocpg = O / groups;
    int cpg = C / groups;
    int g = oc / ocpg;
    float acc = has_bias ? __ldg(b + oc) : 0.0f;
    for (int ci = 0; ci < cpg; ++ci) {
        int ic = g * cpg + ci;
        for (int ky = 0; ky < K; ++ky) {
            int ih = oh * stride - pad + ky * dil;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kx = 0; kx < K; ++kx) {
                int iw = ow * stride - pad + kx * dil;
                if ((unsigned)iw >= (unsigned)Wd) continue;
                float xv = __ldg(x + ((int64_t)n * C + ic) * H * Wd + (int64_t)ih * Wd + iw);
                float wv = __ldg(w + (((int64_t)oc * cpg + ci) * K + ky) * K + kx);
                acc = fmaf(xv, wv, acc);
            }
        }
    }
    y[idx] = acc;
}

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, bool has_bias) {
    CHECK_CUDA(x); CHECK_CUDA(w);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(w);
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int Wd = (int)x.size(3);
    int O = (int)w.size(0);
    int K = (int)w.size(2);
    int OH = (H + 2 * (int)padding - (int)dilation * (K - 1) - 1) / (int)stride + 1;
    int OW = (Wd + 2 * (int)padding - (int)dilation * (K - 1) - 1) / (int)stride + 1;
    auto y = torch::empty({N, O, OH, OW}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (C == 16 && K == 3 && stride == 1 && padding == 0 && dilation == 1 && groups == 1) {
        dim3 block(256);
        dim3 grid(((N * OH * OW) + block.x * 2 - 1) / (block.x * 2), (O + 7) / 8);
        conv3x3_c16_oc8_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(),
            has_bias ? b.data_ptr<float>() : nullptr, y.data_ptr<float>(), N, H, Wd, O, OH, OW, has_bias ? 1 : 0);
    } else {
        int threads = 256;
        int64_t total = (int64_t)N * O * OH * OW;
        int blocks = (int)((total + threads - 1) / threads);
        conv2d_general_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(),
            has_bias ? b.data_ptr<float>() : nullptr, y.data_ptr<float>(), N, C, H, Wd, O, K, OH, OW,
            (int)stride, (int)padding, (int)dilation, (int)groups, has_bias ? 1 : 0);
    }
    return y;
}
"""

_ext = load_inline(
    name="conv2d_specialized_inline_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv2d_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    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)
            has_bias = False
        else:
            has_bias = True
        return _ext.conv2d_forward(x, self.conv2d.weight, b, self._stride, self._padding, self._dilation, self._groups, has_bias)