```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


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


_CPP_SRC = "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias);"

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

#define IN_C {IN_C}
#define KH {KH}
#define KW {KW}
#define SH {SH}
#define SW {SW}
#define DH {DH}
#define DW {DW}

#define KDIM (IN_C * KH * KW)
#define BK 32
#define NCHUNK (KDIM / BK)
#define BM 128
#define BN 64
#define NTHREADS 256
#define APITCH (BM + 8)
#define BPITCH (BK + 4)

#define SM_A (2 * BK * APITCH)
#define SM_B (2 * BN * BPITCH)
#define SM_FLOATS (2 * SM_A + SM_B + KDIM)
#define SM_BYTES (SM_FLOATS * 4)

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

__device__ __forceinline__ void split_tf32(float x, unsigned &hi, unsigned &lo) {
    hi = cvt_tf32(x);
    lo = cvt_tf32(x - __uint_as_float(hi));
}

__device__ __forceinline__ void mma_tf32(float &c0, float &c1, float &c2, float &c3,
                                         unsigned a0, unsigned a1, unsigned a2, unsigned a3,
                                         unsigned b0, unsigned 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"(c0), "+f"(c1), "+f"(c2), "+f"(c3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

template <bool HAS_BIAS, bool NOK>
__global__ void __launch_bounds__(NTHREADS, 1)
igemm3x_kernel(const float* __restrict__ x,
               const float* __restrict__ w,
               const float* __restrict__ bias,
               float* __restrict__ y,
               const int H, const int W,
               const int Ho, const int Wo, const long M,
               const int OC)
{
    extern __shared__ float smem[];
    float* sAhi = smem;                  // [2][BK][APITCH]
    float* sAlo = smem + SM_A;           // [2][BK][APITCH]
    float* sB   = smem + 2 * SM_A;       // [2][BN][BPITCH]
    int*   sK   = (int*)(smem + 2 * SM_A + SM_B);  // [KDIM]

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int lg   = lane >> 2;
    const int lt   = lane & 3;
    const int wm   = warp >> 1;   // 0..3, M direction (32 rows each)
    const int wn   = warp & 1;    // 0..1, N direction (32 cols each)

    const long m_tile = (long)blockIdx.x * BM;
    const int  n_blk  = blockIdx.y * BN;
    const long HoWo   = (long)Ho * (long)Wo;
    const long HW     = (long)H * (long)W;

    // ---- shared k offsets: k = ic*(KH*KW) + kh*KW + kw ----
    for (int k = tid; k < KDIM; k += NTHREADS) {
        const int kw = k % KW;
        const int t  = k / KW;
        const int kh = t % KH;
        const int ic = t / KH;
        sK[k] = (int)((long)ic * HW + (long)kh * DH * W + (long)kw * DW);
    }
    __syncthreads();   // make sK visible before any staging reads it

    // ---- A staging coordinates: 2 threads per row ----
    const int arow_s = tid & (BM - 1);       // 0..127
    const int akoff  = (tid >> 7) * 16;      // 0 or 16

    const long mg = m_tile + arow_s;
    long abase;
    if (mg < M) {
        const long nb = mg / HoWo;
        const long r  = mg - nb * HoWo;
        const int oh  = (int)(r / Wo);
        const int ow  = (int)(r - (long)oh * Wo);
        abase = nb * ((long)IN_C * HW) + (long)oh * SH * W + (long)ow * SW;
    } else {
        abase = 0;
    }

    // ---- B staging coordinates ----
    const int bn  = tid >> 2;           // 0..63
    const int bk8 = (tid & 3) * 8;      // 0,8,16,24
    const int gn  = n_blk + bn;
    const bool nvalid = NOK || (gn < OC);

    float ast[16];
    float bst[8];

    auto load_stage = [&](int kn0) {
        if (nvalid) {
            const float4 p0 = *reinterpret_cast<const float4*>(w + (long)gn * KDIM + kn0 + bk8);
            const float4 p1 = *reinterpret_cast<const float4*>(w + (long)gn * KDIM + kn0 + bk8 + 4);
            bst[0] = p0.x; bst[1] = p0.y; bst[2] = p0.z; bst[3] = p0.w;
            bst[4] = p1.x; bst[5] = p1.y; bst[6] = p1.z; bst[7] = p1.w;
        } else {
#pragma unroll
            for (int i = 0; i < 8; ++i) bst[i] = 0.f;
        }
#pragma unroll
        for (int i = 0; i < 16; ++i) {
            ast[i] = x[abase + (long)sK[kn0 + akoff + i]];
        }
    };

    auto store_stage = [&](int buf) {
        unsigned* ah = reinterpret_cast<unsigned*>(sAhi + buf * (BK * APITCH));
        unsigned* al = reinterpret_cast<unsigned*>(sAlo + buf * (BK * APITCH));
#pragma unroll
        for (int i = 0; i < 16; ++i) {
            unsigned hi, lo;
            split_tf32(ast[i], hi, lo);
            ah[(akoff + i) * APITCH + arow_s] = hi;
            al[(akoff + i) * APITCH + arow_s] = lo;
        }
        float* sb = sB + buf * (BN * BPITCH) + bn * BPITCH + bk8;
#pragma unroll
        for (int i = 0; i < 8; ++i) sb[i] = bst[i];
    };

    load_stage(0);
    store_stage(0);
    __syncthreads();

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

    const int arow = wm * 32 + lg;
    const int bcol = wn * 32 + lg;

    int buf = 0;
#pragma unroll 1
    for (int c = 0; c < NCHUNK; ++c) {
        const bool has_next = (c + 1) < NCHUNK;
        if (has_next) load_stage((c + 1) * BK);

        const float* sah = sAhi + buf * (BK * APITCH);
        const float* sal = sAlo + buf * (BK * APITCH);
        const float* sb  = sB   + buf * (BN * BPITCH);

#pragma unroll
        for (int ks = 0; ks < BK; ks += 8) {
            unsigned afh[2][4], afl[2][4];
#pragma unroll
            for (int mt = 0; mt < 2; ++mt) {
                const int r = arow + mt * 16;
                const float* p0h = sah + (ks + lt) * APITCH + r;
                const float* p1h = sah + (ks + lt + 4) * APITCH + r;
                const float* p0l = sal + (ks + lt) * APITCH + r;
                const float* p1l = sal + (ks + lt + 4) * APITCH + r;
                afh[mt][0] = __float_as_uint(p0h[0]);
                afh[mt][1] = __float_as_uint(p0h[8]);
                afh[mt][2] = __float_as_uint(p1h[0]);
                afh[mt][3] = __float_as_uint(p1h[8]);
                afl[mt][0] = __float_as_uint(p0l[0]);
                afl[mt][1] = __float_as_uint(p0l[8]);
                afl[mt][2] = __float_as_uint(p1l[0]);
                afl[mt][3] = __float_as_uint(p1l[8]);
            }
            unsigned bfh[4][2], bfl[4][2];
#pragma unroll
            for (int nt = 0; nt < 4; ++nt) {
                const float* p = sb + (bcol + nt * 8) * BPITCH + ks + lt;
                const float bv0 = p[0];
                const float bv1 = p[4];
                split_tf32(bv0, bfh[nt][0], bfl[nt][0]);
                split_tf32(bv1, bfh[nt][1], bfl[nt][1]);
            }
#pragma unroll
            for (int mt = 0; mt < 2; ++mt) {
#pragma unroll
                for (int nt = 0; nt < 4; ++nt) {
                    float* c = acc[mt][nt];
                    mma_tf32(c[0], c[1], c[2], c[3],
                             afh[mt][0], afh[mt][1], afh[mt][2], afh[mt][3],
                             bfh[nt][0], bfh[nt][1]);
                    mma_tf32(c[0], c[1], c[2], c[3],
                             afh[mt][0], afh[mt][1], afh[mt][2], afh[mt][3],
                             bfl[nt][0], bfl[nt][1]);
                    mma_tf32(c[0], c[1], c[2], c[3],
                             afl[mt][0], afl[mt][1], afl[mt][2], afl[mt][3],
                             bfh[nt][0], bfh[nt][1]);
                }
            }
        }

        if (has_next) store_stage(buf ^ 1);
        __syncthreads();
        buf ^= 1;
    }

    // ---- epilogue: write directly in NCHW layout ----
    const long ocoff = (long)(OC - 1) * HoWo;
#pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
#pragma unroll
        for (int h = 0; h < 2; ++h) {
            const long row = m_tile + wm * 32 + mt * 16 + h * 8 + lg;
            if (row < M) {
                const long nb = row / HoWo;
                float* yb = y + row + nb * ocoff;
#pragma unroll
                for (int nt = 0; nt < 4; ++nt) {
                    const int col = n_blk + wn * 32 + nt * 8 + lt * 2;
                    float v0 = acc[mt][nt][2 * h];
                    float v1 = acc[mt][nt][2 * h + 1];
                    if (HAS_BIAS) {
                        v0 += bias[col];
                        v1 += bias[col + 1];
                    }
                    float* p = yb + (long)col * HoWo;
                    if (NOK) {
                        p[0] = v0;
                        p[HoWo] = v1;
                    } else {
                        if (col < OC) p[0] = v0;
                        if (col + 1 < OC) p[HoWo] = v1;
                    }
                }
            }
        }
    }
}

template <bool HAS_BIAS, bool NOK>
static void launch_igemm(const float* x, const float* w, const float* bias, float* y,
                         int H, int W, int Ho, int Wo, long M, int OC,
                         unsigned gx, unsigned gy, cudaStream_t stream) {
    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(igemm3x_kernel<HAS_BIAS, NOK>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SM_BYTES);
        configured = true;
    }
    igemm3x_kernel<HAS_BIAS, NOK><<<dim3(gx, gy), NTHREADS, SM_BYTES, stream>>>(
        x, w, bias, y, H, W, Ho, Wo, M, OC);
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only fp32 supported");
    TORCH_CHECK(w.size(1) == IN_C && w.size(2) == KH && w.size(3) == KW,
                "weight shape mismatch");
    TORCH_CHECK(x.size(1) == IN_C, "input channel mismatch");

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const long N  = xc.size(0);
    const long H  = xc.size(2);
    const long W  = xc.size(3);
    const long OC = wc.size(0);
    const long Ho = (H - DH * (KH - 1) - 1) / SH + 1;
    const long Wo = (W - DW * (KW - 1) - 1) / SW + 1;
    const long M  = N * Ho * Wo;

    auto y = torch::empty({N, OC, Ho, Wo}, xc.options());
    if (M == 0 || OC == 0) return y;

    const bool has_bias = bias.numel() > 0;
    const float* bias_ptr = nullptr;
    torch::Tensor bc;
    if (has_bias) {
        bc = bias.contiguous();
        bias_ptr = bc.data_ptr<float>();
    }

    const unsigned gx = (unsigned)((M + BM - 1) / BM);
    const unsigned gy = (unsigned)((OC + BN - 1) / BN);
    auto stream = at::cuda::getCurrentCUDAStream();

    const bool nok = (OC % BN) == 0;
    if (has_bias) {
        if (nok) launch_igemm<true, true>(xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
                                          y.data_ptr<float>(), (int)H, (int)W, (int)Ho, (int)Wo,
                                          M, (int)OC, gx, gy, stream.stream());
        else launch_igemm<true, false>(xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
                                       y.data_ptr<float>(), (int)H, (int)W, (int)Ho, (int)Wo,
                                       M, (int)OC, gx, gy, stream.stream());
    } else {
        if (nok) launch_igemm<false, true>(xc.data_ptr<float>(), wc.data_ptr<float>(), nullptr,
                                           y.data_ptr<float>(), (int)H, (int)W, (int)Ho, (int)Wo,
                                           M, (int)OC, gx, gy, stream.stream());
        else launch_igemm<false, false>(xc.data_ptr<float>(), wc.data_ptr<float>(), nullptr,
                                        y.data_ptr<float>(), (int)H, (int)W, (int)Ho, (int)Wo,
                                        M, (int)OC, gx, gy, stream.stream());
    }

    return y;
}
"""


class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 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,
                                stride=stride, padding=padding,
                                dilation=dilation, groups=groups, bias=bias)

        kh, kw = _pair(kernel_size)
        st = _pair(stride)
        pd = _pair(padding)
        dl = _pair(dilation)

        kdim = int(in_channels) * kh * kw

        self._native = True
        self._ext = None
        self._fn = None
        self._empty_bias = None

        ok = (
            groups == 1
            and pd == (0, 0)
            and (kdim % 32) == 0
            and (int(out_channels) % 8) == 0
            and int(in_channels) >= 1
        )

        if ok and torch.cuda.is_available():
            try:
                major, minor = torch.cuda.get_device_capability()
                prop = torch.cuda.get_device_properties(torch.cuda.current_device())
                sm_need = (2 * (2 * 32 * 136) + (2 * 64 * 36) + kdim) * 4
                ok = (major >= 8) and (prop.shared_memory_per_block_optin >= sm_need)
            except Exception:
                ok = False
        else:
            ok = False

        if ok:
            macros = {
                "IN_C": int(in_channels),
                "KH": kh,
                "KW": kw,
                "SH": st[0],
                "SW": st[1],
                "DH": dl[0],
                "DW": dl[1],
            }
            src = _TEMPLATE
            for k, v in macros.items():
                src = src.replace("{%s}" % k, str(v))
            name = "igemm3xtf32_c{}_k{}x{}_s{}x{}_d{}x{}_v5".format(
                in_channels, kh, kw, st[0], st[1], dl[0], dl[1])
            self._ext = load_inline(
                name=name,
                cpp_sources=_CPP_SRC,
                cuda_sources=src,
                functions=["conv2d_forward"],
                verbose=False,
                extra_cuda_cflags=["-O3"],
            )
            self._fn = self._ext.conv2d_forward
            self._native = False
            if bias:
                self._empty_bias = None
            else:
                self._empty_bias = torch.empty(0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        fn = self._fn
        if fn is not None and x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            b = self.conv2d.bias
            if b is None:
                b = self._empty_bias
                if b.device != x.device:
                    b = x.new_empty(0)
                    self._empty_bias = b
            return fn(x, self.conv2d.weight, b)
        return self.conv2d(x)
```