```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 256
#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_K (KDIM)
#define SM_BYTES ((SM_A + SM_B + SM_K) * 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 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)
igemm_tf32_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* sA = smem;                       // [2][BK][APITCH]
    float* sB = smem + SM_A;                // [2][BN][BPITCH]
    int*   sK = (int*)(smem + 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;   // groupID 0..7
    const int lt   = lane & 3;    // threadID_in_group 0..3
    const int wm   = warp >> 1;   // 0..3 (M direction, 64 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);
    }

    // ---- per-thread A row (row tid of the BM tile) ----
    const long mg = m_tile + tid;
    long base;
    bool mvalid;
    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);
        base = nb * ((long)IN_C * HW) + (long)oh * SH * W + (long)ow * SW;
        mvalid = true;
    } else {
        base = 0;
        mvalid = false;
    }

    // ---- 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;         // global output channel
    const bool nvalid = NOK || (gn < OC);

    float bst[8];

    // ---- stage chunk 0 ----
    {
        const int k0 = bk8;
        if (nvalid) {
            const float4 p0 = *reinterpret_cast<const float4*>(w + (long)gn * KDIM + k0);
            const float4 p1 = *reinterpret_cast<const float4*>(w + (long)gn * KDIM + k0 + 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;
        }
        float* sb = &sB[bn * BPITCH + bk8];
#pragma unroll
        for (int i = 0; i < 8; ++i)
            reinterpret_cast<unsigned*>(sb)[i] = cvt_tf32(bst[i]);

#pragma unroll
        for (int i = 0; i < BK; ++i) {
            const float av = x[base + (long)sK[i]];
            reinterpret_cast<unsigned*>(&sA[i * APITCH + tid])[0] = cvt_tf32(av);
        }
    }
    __syncthreads();

    float acc[4][4][4];
#pragma unroll
    for (int mt = 0; mt < 4; ++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 * 64 + 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;
        const int kn0 = (c + 1) * BK;

        // issue next-chunk B global loads
        if (has_next) {
            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;
            }
        }

        const float* sa = sA + buf * (BK * APITCH);
        const float* sb = sB + buf * (BN * BPITCH);

#pragma unroll
        for (int ks = 0; ks < BK; ks += 8) {
            unsigned af[4][4];
#pragma unroll
            for (int mt = 0; mt < 4; ++mt) {
                const int r = arow + mt * 16;
                const float* p0 = sa + (ks + lt) * APITCH + r;
                const float* p1 = sa + (ks + lt + 4) * APITCH + r;
                af[mt][0] = __float_as_uint(p0[0]);
                af[mt][1] = __float_as_uint(p0[8]);
                af[mt][2] = __float_as_uint(p1[0]);
                af[mt][3] = __float_as_uint(p1[8]);
            }
            unsigned bf[4][2];
#pragma unroll
            for (int nt = 0; nt < 4; ++nt) {
                const float* p = sb + (bcol + nt * 8) * BPITCH + ks + lt;
                bf[nt][0] = __float_as_uint(p[0]);
                bf[nt][1] = __float_as_uint(p[4]);
            }
#pragma unroll
            for (int mt = 0; mt < 4; ++mt) {
#pragma unroll
                for (int nt = 0; nt < 4; ++nt) {
                    mma_tf32(acc[mt][nt][0], acc[mt][nt][1], acc[mt][nt][2], acc[mt][nt][3],
                             af[mt][0], af[mt][1], af[mt][2], af[mt][3],
                             bf[nt][0], bf[nt][1]);
                }
            }

            // interleave next-chunk A gather with the mma pipeline
            if (has_next) {
                const int i0 = ks + 4;
                const long o0 = (long)sK[kn0 + i0];
                const long o1 = (long)sK[kn0 + i0 + 1];
                const long o2 = (long)sK[kn0 + i0 + 2];
                const long o3 = (long)sK[kn0 + i0 + 3];
                float v0 = x[base + o0];
                float v1 = x[base + o1];
                float v2 = x[base + o2];
                float v3 = x[base + o3];
                const int nb = (buf ^ 1) * (BK * APITCH);
                reinterpret_cast<unsigned*>(&smem[nb + i0 * APITCH + tid])[0]       = cvt_tf32(v0);
                reinterpret_cast<unsigned*>(&smem[nb + (i0 + 1) * APITCH + tid])[0] = cvt_tf32(v1);
                reinterpret_cast<unsigned*>(&smem[nb + (i0 + 2) * APITCH + tid])[0] = cvt_tf32(v2);
                reinterpret_cast<unsigned*>(&smem[nb + (i0 + 3) * APITCH + tid])[0] = cvt_tf32(v3);
            }
        }

        if (has_next) {
            // first 4 A rows of next chunk
            const long o0 = (long)sK[kn0];
            const long o1 = (long)sK[kn0 + 1];
            const long o2 = (long)sK[kn0 + 2];
            const long o3 = (long)sK[kn0 + 3];
            float v0 = x[base + o0];
            float v1 = x[base + o1];
            float v2 = x[base + o2];
            float v3 = x[base + o3];
            const int nb = (buf ^ 1) * (BK * APITCH);
            reinterpret_cast<unsigned*>(&smem[nb + tid])[0]                 = cvt_tf32(v0);
            reinterpret_cast<unsigned*>(&smem[nb + APITCH + tid])[0]        = cvt_tf32(v1);
            reinterpret_cast<unsigned*>(&smem[nb + 2 * APITCH + tid])[0]    = cvt_tf32(v2);
            reinterpret_cast<unsigned*>(&smem[nb + 3 * APITCH + tid])[0]    = cvt_tf32(v3);

            float* sbn = &sB[(buf ^ 1) * (BN * BPITCH) + bn * BPITCH + bk8];
#pragma unroll
            for (int i = 0; i < 8; ++i)
                reinterpret_cast<unsigned*>(sbn)[i] = cvt_tf32(bst[i]);
        }
        buf ^= 1;
        __syncthreads();
    }

    // ---- epilogue ----
    const long yrow_stride = (long)OC;
#pragma unroll
    for (int mt = 0; mt < 4; ++mt) {
#pragma unroll
        for (int h = 0; h < 2; ++h) {
            const long row = m_tile + wm * 64 + mt * 16 + lg + h * 8;
            if (mvalid || (row < M)) {
                const long rvalid = (mt == 0 && h == 0) ? (row < M) : (row < M);
                if (rvalid) {
                    float* yp = y + row * yrow_stride;
#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];
                        }
                        if (NOK || (col + 1 < OC)) {
                            *reinterpret_cast<float2*>(yp + col) = make_float2(v0, v1);
                        } else if (col < OC) {
                            yp[col] = v0;
                        }
                    }
                }
            }
        }
    }
}

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(igemm_tf32_kernel<HAS_BIAS, NOK>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SM_BYTES);
        configured = true;
    }
    igemm_tf32_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({M, OC}, xc.options());
    if (M == 0 || OC == 0) {
        return y.view({N, OC, Ho, Wo});
    }

    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.view({N, OC, Ho, Wo});
}
"""


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

        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 * 32 * 264 + 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 = "igemm_tf32_c{}_k{}x{}_s{}x{}_d{}x{}_v4".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._native = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (not self._native) and x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            w = self.conv2d.weight
            b = self.conv2d.bias
            if b is None:
                b = x.new_empty(0)
            return self._ext.conv2d_forward(x, w, b)
        return self.conv2d(x)
```