```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

#define TM 128
#define TN 64
#define TK 32
#define LDA 136
#define LDB 68

__device__ __forceinline__ uint32_t f2tf32(float x) {
    uint32_t r;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
}

__device__ __forceinline__ uint4 to_tf32x4(uint4 v) {
    v.x = f2tf32(__uint_as_float(v.x));
    v.y = f2tf32(__uint_as_float(v.y));
    v.z = f2tf32(__uint_as_float(v.z));
    v.w = f2tf32(__uint_as_float(v.w));
    return v;
}

__device__ __forceinline__ void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                                         uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                                         uint32_t b0, uint32_t b1) {
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// Implicit-GEMM convolution with TF32 tensor cores.
// C[m, n] = sum_k A[m, k] * Wt[k, n]
//   m -> (b, h_out, w_out), n -> out_channel, k -> (c_in, kh, kw)
__global__ void __launch_bounds__(256)
conv2d_tf32_igemm_kernel(const float* __restrict__ x,
                         const uint32_t* __restrict__ wt,
                         const float* __restrict__ bias,
                         float* __restrict__ out,
                         int M, int N, int K,
                         int Bc, int C, int Hin, int Win,
                         int Ho, int Wo,
                         int kh, int kw, int dh, int dw, int ph, int pw) {
    extern __shared__ float smem[];
    uint32_t* As = reinterpret_cast<uint32_t*>(smem);                 // 2*TK*LDA
    uint32_t* Bs = As + 2 * TK * LDA;                                 // 2*TK*LDB
    int*  koff = reinterpret_cast<int*>(Bs + 2 * TK * LDB);           // K
    short* kdi = reinterpret_cast<short*>(koff + K);                  // K
    short* kdj = kdi + K;                                             // K
    float* Cs = smem;                                                 // reuse: TN*LDA

    const int tid = threadIdx.x;
    const int m_tile = blockIdx.x * TM;
    const int n_tile = blockIdx.y * TN;
    const int HWo = Ho * Wo;

    // ---- init kernel-offset tables ----
    for (int k = tid; k < K; k += 256) {
        int c = k / (kh * kw);
        int r = k - c * (kh * kw);
        int i = r / kw;
        int j = r - i * kw;
        int di = i * dh - ph;
        int dj = j * dw - pw;
        koff[k] = c * Hin * Win + di * Win + dj;
        kdi[k] = (short)di;
        kdj[k] = (short)dj;
    }
    __syncthreads();  // tables must be visible before load_chunk reads them

    // ---- per-thread staging setup ----
    const int kloc = tid >> 3;          // 0..31
    const int mloc = tid & 7;           // 0..7
    int m_thread = m_tile + mloc;
    int h0 = 0, w0 = 0, off0 = 0, e_lim = 0;
    if (m_thread < M) {
        int b0 = m_thread / HWo;
        int rem = m_thread - b0 * HWo;
        h0 = rem / Wo;
        w0 = rem - h0 * Wo;
        off0 = (b0 * C) * (Hin * Win) + h0 * Win + w0;
        int rl = M - m_thread;
        e_lim = (rl + 7) >> 3;
        if (e_lim > 16) e_lim = 16;
    }

    // ---- double-buffered staging through registers ----
    uint32_t ra[16];
    uint4 rb[2];
    const float4* wt4 = reinterpret_cast<const float4*>(wt);
    const int n_ok = (n_tile + (mloc << 3)) < N;

    auto load_chunk = [&](int chunk) {
        int k = chunk * TK + kloc;
        uint4 b0v = make_uint4(0u, 0u, 0u, 0u), b1v = b0v;
        if (k < K && n_ok) {
            const float4* p = wt4 + (long)k * (N >> 2) + (n_tile >> 2) + (mloc << 1);
            b0v = to_tf32x4(*reinterpret_cast<const uint4*>(p));
            b1v = to_tf32x4(*reinterpret_cast<const uint4*>(p + 1));
        }
        rb[0] = b0v; rb[1] = b1v;

        int ok = 0, di = 0, dj = 0;
        if (k < K) { ok = koff[k]; di = (int)kdi[k]; dj = (int)kdj[k]; }
        int ch = h0, cw = w0, coff = off0;
        const float* xp = x + coff + ok;
        #pragma unroll
        for (int e = 0; e < 16; ++e) {
            bool v = (e < e_lim) &&
                     ((unsigned)(ch + di) < (unsigned)Hin) &&
                     ((unsigned)(cw + dj) < (unsigned)Win);
            ra[e] = v ? f2tf32(__ldg(xp)) : 0u;
            cw += 8; coff += 8; xp += 8;
            while (cw >= Wo) {
                cw -= Wo; ch += 1; coff += (Win - Wo); xp += (Win - Wo);
                if (ch >= Ho) { ch -= Ho; coff += C * Hin * Win - Ho * Win; xp += C * Hin * Win - Ho * Win; }
            }
        }
    };

    auto store_chunk = [&](int buf) {
        uint32_t* ab = As + buf * TK * LDA + kloc * LDA + mloc;
        #pragma unroll
        for (int e = 0; e < 16; ++e) ab[e * 8] = ra[e];
        uint4* bb = reinterpret_cast<uint4*>(Bs + buf * TK * LDB + kloc * LDB + mloc * 8);
        bb[0] = rb[0];
        bb[1] = rb[1];
    };

    const int nchunks = (K + TK - 1) / TK;
    load_chunk(0);
    store_chunk(0);
    __syncthreads();

    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int wm = (warp & 3) * 32;
    const int wn = (warp >> 2) * 32;
    const int lr = lane >> 2;
    const int lc = lane & 3;

    float acc[2][4][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            #pragma unroll
            for (int v = 0; v < 4; ++v) acc[i][j][v] = 0.f;

    for (int chunk = 0; chunk < nchunks; ++chunk) {
        if (chunk + 1 < nchunks) load_chunk(chunk + 1);

        const uint32_t* ab = As + (chunk & 1) * TK * LDA;
        const uint32_t* bb = Bs + (chunk & 1) * TK * LDB;

        #pragma unroll
        for (int ks = 0; ks < TK / 8; ++ks) {
            int kk = ks * 8;
            uint32_t a0[2], a1[2], a2[2], a3[2];
            #pragma unroll
            for (int mi = 0; mi < 2; ++mi) {
                int mm = wm + mi * 16 + lr;
                const uint32_t* p0 = ab + (kk + lc) * LDA + mm;
                const uint32_t* p1 = ab + (kk + lc + 4) * LDA + mm;
                a0[mi] = p0[0]; a1[mi] = p0[8]; a2[mi] = p1[0]; a3[mi] = p1[8];
            }
            #pragma unroll
            for (int ni = 0; ni < 4; ++ni) {
                int nn = wn + ni * 8 + lr;
                uint32_t bv0 = bb[(kk + lc) * LDB + nn];
                uint32_t bv1 = bb[(kk + lc + 4) * LDB + nn];
                mma_tf32(acc[0][ni][0], acc[0][ni][1], acc[0][ni][2], acc[0][ni][3],
                         a0[0], a1[0], a2[0], a3[0], bv0, bv1);
                mma_tf32(acc[1][ni][0], acc[1][ni][1], acc[1][ni][2], acc[1][ni][3],
                         a0[1], a1[1], a2[1], a3[1], bv0, bv1);
            }
        }

        if (chunk + 1 < nchunks) store_chunk((chunk + 1) & 1);
        __syncthreads();
    }

    // ---- epilogue: transpose through shared, then coalesced store ----
    #pragma unroll
    for (int ni = 0; ni < 4; ++ni) {
        #pragma unroll
        for (int mi = 0; mi < 2; ++mi) {
            int mm = wm + mi * 16 + lr;
            int nn = wn + ni * 8 + lc * 2;
            Cs[nn * LDA + mm]           = acc[mi][ni][0];
            Cs[(nn + 1) * LDA + mm]     = acc[mi][ni][1];
            Cs[nn * LDA + (mm + 8)]     = acc[mi][ni][2];
            Cs[(nn + 1) * LDA + (mm + 8)] = acc[mi][ni][3];
        }
    }
    __syncthreads();

    const int m_c = tid & 127;
    const int n_c0 = tid >> 7;
    int m_g = m_tile + m_c;
    if (m_g < M) {
        int b = m_g / HWo;
        long base = (long)m_g + (long)b * (N - 1) * HWo + (long)n_tile * HWo;
        #pragma unroll 4
        for (int it = 0; it < TN / 2; ++it) {
            int nn_l = n_c0 + it * 2;
            int n = n_tile + nn_l;
            if (n < N) {
                float v = Cs[nn_l * LDA + m_c];
                if (bias) v += bias[n];
                out[base + (long)nn_l * HWo] = v;
            }
        }
    }
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor wt, torch::Tensor bias,
                             int64_t stride, int64_t ph, int64_t pw,
                             int64_t dh, int64_t dw, int64_t kh, int64_t kw) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat, "x must be CUDA fp32");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    TORCH_CHECK(wt.dim() == 2, "wt must be [K, N]");
    auto xc = x.contiguous();
    auto wc = wt.contiguous();

    const int Bc = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int Hin = (int)xc.size(2);
    const int Win = (int)xc.size(3);
    const int K = (int)wc.size(0);
    const int N = (int)wc.size(1);
    const int Ho = (Hin + 2 * (int)ph - (int)dh * ((int)kh - 1) - 1) / (int)stride + 1;
    const int Wo = (Win + 2 * (int)pw - (int)dw * ((int)kw - 1) - 1) / (int)stride + 1;

    auto out = torch::empty({Bc, N, Ho, Wo}, xc.options());

    const int M = Bc * Ho * Wo;
    const int smem = 2 * TK * LDA * 4 + 2 * TK * LDB * 4 + K * 8;
    static int configured_smem = 0;
    if (smem > configured_smem) {
        cudaFuncSetAttribute(conv2d_tf32_igemm_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
        configured_smem = smem;
    }

    dim3 grid((M + TM - 1) / TM, (N + TN - 1) / TN);
    const float* bias_ptr = (bias.defined() && bias.numel() > 0)
                                ? bias.data_ptr<float>() : nullptr;
    auto stream = at::cuda::getCurrentCUDAStream();
    conv2d_tf32_igemm_kernel<<<grid, 256, smem, stream>>>(
        xc.data_ptr<float>(),
        reinterpret_cast<const uint32_t*>(wc.data_ptr<float>()),
        bias_ptr, out.data_ptr<float>(),
        M, N, K, Bc, C, Hin, Win, Ho, Wo,
        (int)kh, (int)kw, (int)dh, (int)dw, (int)ph, (int)pw);
    return out;
}
"""

_cpp_src = "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor wt, torch::Tensor bias, int64_t stride, int64_t ph, int64_t pw, int64_t dh, int64_t dw, int64_t kh, int64_t kw);"

_conv_ext = load_inline(
    name="conv2d_tf32_igemm_v2",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


def _pair(v):
    if isinstance(v, int):
        return (v, v)
    return (int(v[0]), int(v[1]))


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.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, bias=bias)
        self._stride = _pair(stride)
        self._padding = _pair(padding)
        self._dilation = _pair(dilation)
        self._cached_wt = None
        self._cached_version = -1
        self._fast = None
        self._empty_bias = None

    def _fast_ok(self, x):
        if self._fast is None:
            ok = False
            try:
                cap = torch.cuda.get_device_capability(x.device)
                w = self.conv2d.weight
                N = w.shape[0]
                K = w.shape[1] * w.shape[2] * w.shape[3]
                smem = 2 * 32 * 136 * 4 + 2 * 32 * 68 * 4 + K * 8
                ok = (cap[0] >= 8 and (N % 8) == 0 and (N % 4) == 0
                      and smem <= 160 * 1024 and self._stride == (1, 1))
            except Exception:
                ok = False
            self._fast = ok
        return self._fast

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv2d.weight
        if (x.is_cuda and x.dtype == torch.float32 and w.dtype == torch.float32
                and self._fast_ok(x)):
            if (self._cached_wt is None or self._cached_version != w._version
                    or self._cached_wt.device != w.device):
                # transpose weights once to [C*kh*kw, N] layout for coalesced access
                self._cached_wt = w.detach().reshape(w.shape[0], -1).t().contiguous()
                self._cached_version = w._version
            b = self.conv2d.bias
            if b is None:
                if self._empty_bias is None:
                    self._empty_bias = w.new_empty(0)
                b = self._empty_bias
            return _conv_ext.conv2d_forward(
                x, self._cached_wt, b,
                1, self._padding[0], self._padding[1],
                self._dilation[0], self._dilation[1],
                w.shape[2], w.shape[3])
        return self.conv2d(x)
```