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_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                             int stride_h, int stride_w, int pad_h, int pad_w,
                             int dil_h, int dil_w, int groups, bool has_bias);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <mma.h>
using namespace nvcuda;

__device__ __forceinline__ float to_tf32(float x) {
    unsigned y;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(y) : "f"(x));
    return __uint_as_float(y);
}

#define WARPS 8

__global__ void conv3x3_wmma_kernel(const float* __restrict__ x,
                                    const float* __restrict__ w,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    int N, int C, int H, int W,
                                    int K, int Ho, int Wo, int has_bias) {
    const int warp_id = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    int p0 = blockIdx.x * 16;
    int oc0 = warp_id * 16;
    int P = N * Ho * Wo;

    extern __shared__ float smem[];
    float* Bs = smem;                                      // 8 x 16, common to all warps
    float* As = Bs + 8*16 + warp_id * (16*8);              // 16 x 8 per warp
    float* Cs = Bs + 8*16 + WARPS*(16*8) + warp_id*(16*16);

    wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> c_frag;
    wmma::fill_fragment(c_frag, 0.0f);

    #pragma unroll 1
    for (int kk = 0; kk < 576; kk += 8) {
        // B tile: im2col columns for 16 adjacent output pixels. Ordering ic,kh,kw
        for (int idx = threadIdx.x; idx < 128; idx += blockDim.x) {
            int r = idx >> 4;
            int col = idx & 15;
            int kidx = kk + r;
            int p = p0 + col;
            float val = 0.0f;
            if (p < P) {
                int ow = p % Wo;
                int tmp = p / Wo;
                int oh = tmp % Ho;
                int n = tmp / Ho;
                int ic = kidx / 9;
                int rem = kidx - ic * 9;
                int kh = rem / 3;
                int kw = rem - kh * 3;
                val = x[((n*C + ic)*H + oh + kh)*W + ow + kw];
            }
            Bs[idx] = to_tf32(val);
        }
        // A tile: 16 output channels per warp
        for (int idx = lane; idx < 128; idx += 32) {
            int r = idx >> 3;
            int s = idx & 7;
            int kidx = kk + s;
            int ic = kidx / 9;
            int rem = kidx - ic * 9;
            As[idx] = to_tf32(w[((oc0 + r)*C + ic)*9 + rem]);
        }
        __syncthreads();
        wmma::load_matrix_sync(a_frag, As, 8);
        wmma::load_matrix_sync(b_frag, Bs, 16);
        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
        __syncthreads();
    }

    wmma::store_matrix_sync(Cs, c_frag, 16, wmma::mem_row_major);
    __syncwarp();
    for (int idx = lane; idx < 256; idx += 32) {
        int r = idx >> 4;
        int col = idx & 15;
        int p = p0 + col;
        int oc = oc0 + r;
        if (p < P) {
            int local = p % (Ho * Wo);
            int n = p / (Ho * Wo);
            y[(n*K + oc)*Ho*Wo + local] = Cs[idx] + (has_bias ? b[oc] : 0.0f);
        }
    }
}

__global__ void conv3x3_scalar_tf32(const float* __restrict__ x,
                                    const float* __restrict__ w,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    int N, int C, int H, int W,
                                    int K, int Ho, int Wo, int has_bias) {
    int p = blockIdx.x * blockDim.x + threadIdx.x;
    int oc = blockIdx.y * blockDim.y + threadIdx.y;
    int P = N * Ho * Wo;
    if (p >= P || oc >= K) return;
    int ow = p % Wo;
    int t = p / Wo;
    int oh = t % Ho;
    int n = t / Ho;
    float acc = has_bias ? b[oc] : 0.0f;
    for (int ic = 0; ic < C; ++ic)
        for (int kh = 0; kh < 3; ++kh)
            for (int kw = 0; kw < 3; ++kw)
                acc = fmaf(to_tf32(x[((n*C + ic)*H + oh + kh)*W + ow + kw]),
                           to_tf32(w[((oc*C + ic)*3 + kh)*3 + kw]), acc);
    y[(n*K + oc)*Ho*Wo + oh*Wo + ow] = acc;
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                             int sh, int sw, int ph, int pw, int dh, int dw,
                             int groups, bool has_bias) {
    x = x.contiguous();
    w = w.contiguous();
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int K = (int)w.size(0), R = (int)w.size(2);
    int Ho = (H + 2*ph - dh*(R-1) - 1) / sh + 1;
    int Wo = (W + 2*pw - dw*(R-1) - 1) / sw + 1;
    auto y = torch::empty({N, K, Ho, Wo}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (C == 64 && K == 128 && R == 3 && sh == 1 && sw == 1 && ph == 0 && pw == 0 && dh == 1 && dw == 1 && groups == 1) {
        dim3 block(256);
        dim3 grid((N * Ho * Wo + 15) / 16);
        conv3x3_wmma_kernel<<<grid, block, (8*16 + WARPS*16*8 + WARPS*16*16) * sizeof(float), stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? b.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), N, C, H, W, K, Ho, Wo, has_bias);
    } else {
        dim3 block(32, 8);
        dim3 grid((N * Ho * Wo + 31) / 32, (K + 7) / 8);
        conv3x3_scalar_tf32<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? b.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), N, C, H, W, K, Ho, Wo, has_bias);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='conv2d_asym_3x3_wmma_final',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['conv2d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bias = self.conv2d.bias if self.conv2d.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
        sh, sw = self.conv2d.stride
        ph, pw = self.conv2d.padding
        dh, dw = self.conv2d.dilation
        return _ext.conv2d_forward(x, self.conv2d.weight, bias, sh, sw, ph, pw, dh, dw,
                                   self.conv2d.groups, self.conv2d.bias is not None)