```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 OC {OC}
#define KH {KH}
#define KW {KW}

#define KDIM (IN_C * KH * KW)
#define BK 32
#define NCHUNK (KDIM / BK)
#define NTHREADS 256
#define NWARPS 8
#define BM (NWARPS * 32)

#if OC == 64
#define NSEG 1
#define NTH 2
#elif (OC % 16) == 0
#define NSEG (OC / 16)
#define NTH (OC / 8)
#else
#error "unsupported OC"
#endif

#define ACCN (2 * NTH)

__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));
}

__device__ __forceinline__ float2 load_bf(const float* s) {
    float2 v;
    asm volatile("ld.shared.v2.f32 {%0,%1}, [%2];\n" : "=f"(v.x), "=f"(v.y) : "l"(s));
    return v;
}

template <bool HAS_BIAS, bool MOK>
__global__ void __launch_bounds__(NTHREADS, 2)
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)
{
    __shared__ float sKoff[KDIM];
    __shared__ float sB[2][OC][BK + 4];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int lg = lane >> 2;   // 0..7
    const int lt = lane & 3;    // 0..3

    const long HoWo = (long)Ho * (long)Wo;
    const long chanstride = (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;
        sKoff[k] = (float)0.f + 0 * (int)sizeof(float) + 0;  // placeholder (overwritten)
        sKoff[k] = ic * (int)chanstride + kh * W + kw;
    }

    // ---- cooperative B chunk 0 load (with tf32 convert) ----
    {
        const int oc = tid >> 2;
        const int k8 = (tid & 3) * 8;
        const float* wp = w + (long)oc * KDIM + k8;
        float4 p0 = *reinterpret_cast<const float4*>(wp);
        float4 p1 = *reinterpret_cast<const float4*>(wp + 4);
        float* sp = &sB[0][oc][k8];
        sp[0] = (float)0.f;  // keep compiler from over-optimizing ordering
        reinterpret_cast<unsigned*>(sp)[0] = cvt_tf32(p0.x);
        reinterpret_cast<unsigned*>(sp)[1] = cvt_tf32(p0.y);
        reinterpret_cast<unsigned*>(sp)[2] = cvt_tf32(p0.z);
        reinterpret_cast<unsigned*>(sp)[3] = cvt_tf32(p0.w);
        reinterpret_cast<unsigned*>(sp + 4)[0] = cvt_tf32(p1.x);
        reinterpret_cast<unsigned*>(sp + 4)[1] = cvt_tf32(p1.y);
        reinterpret_cast<unsigned*>(sp + 4)[2] = cvt_tf32(p1.z);
        reinterpret_cast<unsigned*>(sp + 4)[3] = cvt_tf32(p1.w);
    }
    __syncthreads();

    // ---- per-thread row bases (4 rows owned per thread) ----
    const long m_tile = (long)blockIdx.x * BM;
    long base[4];
    bool rv[4];
#pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
#pragma unroll
        for (int h = 0; h < 2; ++h) {
            const long m = m_tile + warp * 32 + mt * 16 + lg + h * 8;
            if (MOK || m < M) {
                const long n = m / HoWo;
                long r = m - n * HoWo;
                const int oh = (int)(r / Wo);
                const int ow = (int)(r - (long)oh * Wo);
                base[mt * 2 + h] = n * (long)IN_C * chanstride + (long)oh * W + ow;
                rv[mt * 2 + h] = true;
            } else {
                base[mt * 2 + h] = 0;
                rv[mt * 2 + h] = false;
            }
        }
    }

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

    // ---- software-pipelined A fragment prefetch ----
    int globk = lt;              // k index of first prefetch
    const float* xp = x;
    float4 areg[2][2];
#pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        const int off = sKoff[globk];
        const int off4 = sKoff[globk + 4];
        const long b0 = base[mt * 2] + off;
        const long b1 = base[mt * 2 + 1] + off;
        const long b2 = base[mt * 2] + off4;
        const long b3 = base[mt * 2 + 1] + off4;
        areg[mt][0].x = xp[b0];
        areg[mt][0].y = xp[b1];
        areg[mt][0].z = xp[b2];
        areg[mt][0].w = xp[b3];
        areg[mt][1].x = 0.f;
        areg[mt][1].y = 0.f;
    }

    int buf = 0;
#pragma unroll 1
    for (int c = 0; c < NCHUNK; ++c) {
        // prefetch next B chunk into registers
        float4 bp0, bp1;
        const bool has_next_b = (c + 1) < NCHUNK;
        if (has_next_b) {
            const int oc = tid >> 2;
            const int kk = (c + 1) * BK + (tid & 3) * 8;
            const float* wp = w + (long)oc * KDIM + kk;
            bp0 = *reinterpret_cast<const float4*>(wp);
            bp1 = *reinterpret_cast<const float4*>(wp + 4);
        }
        const float* sb = &sB[buf][0][0];

#pragma unroll
        for (int ks = 0; ks < BK; ks += 8) {
            // issue prefetch for next k8 step
            const int nk = globk + 8;
            float4 anxt[2];
            if (nk < KDIM) {
                const int off = sKoff[nk];
                const int off4 = sKoff[nk + 4];
#pragma unroll
                for (int mt = 0; mt < 2; ++mt) {
                    anxt[mt].x = xp[base[mt * 2] + off];
                    anxt[mt].y = xp[base[mt * 2 + 1] + off];
                    anxt[mt].z = xp[base[mt * 2] + off4];
                    anxt[mt].w = xp[base[mt * 2 + 1] + off4];
                }
            }

            unsigned af[2][4];
#pragma unroll
            for (int mt = 0; mt < 2; ++mt) {
                af[mt][0] = cvt_tf32(areg[mt][0].x);
                af[mt][1] = cvt_tf32(areg[mt][0].y);
                af[mt][2] = cvt_tf32(areg[mt][0].z);
                af[mt][3] = cvt_tf32(areg[mt][0].w);
            }

#pragma unroll
            for (int nt = 0; nt < NTH; ++nt) {
                const float2 bv = load_bf(sb + (nt * 8 + lg) * (BK + 4) + ks + lt);
                const unsigned b0 = __float_as_uint(bv.x);
                const unsigned b1 = __float_as_uint(bv.y);
#pragma unroll
                for (int mt = 0; mt < 2; ++mt) {
                    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], b0, b1);
                }
            }

            if (nk < KDIM) {
#pragma unroll
                for (int mt = 0; mt < 2; ++mt) areg[mt][0] = anxt[mt];
                globk = nk;
            }
        }

        // store next B chunk to smem (converted)
        if (has_next_b) {
            const int oc = tid >> 2;
            const int k8 = (tid & 3) * 8;
            unsigned* sp = reinterpret_cast<unsigned*>(&sB[buf ^ 1][oc][k8]);
            sp[0] = cvt_tf32(bp0.x);
            sp[1] = cvt_tf32(bp0.y);
            sp[2] = cvt_tf32(bp0.z);
            sp[3] = cvt_tf32(bp0.w);
            sp[4] = cvt_tf32(bp1.x);
            sp[5] = cvt_tf32(bp1.y);
            sp[6] = cvt_tf32(bp1.z);
            sp[7] = cvt_tf32(bp1.w);
        }
        buf ^= 1;
        __syncthreads();
    }

    // ---- epilogue ----
    float biasv[NSEG > 1 ? NSEG : 1][8];
    if (HAS_BIAS) {
#pragma unroll
        for (int s = 0; s < (NSEG > 1 ? NSEG : 1); ++s) {
            if (NSEG == 1) {
#pragma unroll
                for (int j = 0; j < 8; ++j) biasv[0][j] = bias[lg * 2 + (j & 1)];
            } else {
                float4 v0 = *reinterpret_cast<const float4*>(bias + s * 16 + lg * 2);
                float4 v1 = *reinterpret_cast<const float4*>(bias + s * 16 + lg * 2 + 8);
                biasv[s][0] = v0.x; biasv[s][1] = v0.y;
                biasv[s][2] = v0.z; biasv[s][3] = v0.w;
                biasv[s][4] = v1.x; biasv[s][5] = v1.y;
                biasv[s][6] = v1.z; biasv[s][7] = v1.w;
            }
        }
    }

#pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
#pragma unroll
        for (int h = 0; h < 2; ++h) {
            const long m = m_tile + warp * 32 + mt * 16 + lg + h * 8;
            if (MOK || rv[mt * 2 + h]) {
#pragma unroll
                for (int s = 0; s < NSEG; ++s) {
                    float v[8];
#pragma unroll
                    for (int nt = 0; nt < 4; ++nt) {
                        v[nt * 2] = acc[mt][s * 4 + nt][2 * h];
                        v[nt * 2 + 1] = acc[mt][s * 4 + nt][2 * h + 1];
                    }
                    if (HAS_BIAS) {
#pragma unroll
                        for (int j = 0; j < 8; ++j) v[j] += biasv[NSEG > 1 ? s : 0][j];
                    }
                    float* op = y + m * (long)OC + s * 16;
                    *reinterpret_cast<float4*>(op) = make_float4(v[0], v[1], v[2], v[3]);
                    *reinterpret_cast<float4*>(op + 4) = make_float4(v[4], v[5], v[6], v[7]);
                }
            }
        }
    }
}

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(0) == OC && 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 Ho = H - KH + 1;
    const long Wo = W - KW + 1;
    const long M = N * Ho * Wo;

    auto y = torch::empty({M, (long)OC}, xc.options());
    if (M == 0) {
        return y.view({N, (long)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 blocks = (unsigned)((M + BM - 1) / BM);
    dim3 grid(blocks);
    dim3 block(NTHREADS);
    auto stream = at::cuda::getCurrentCUDAStream();

    const bool mok = (M % BM) == 0;
    if (has_bias) {
        if (mok) {
            igemm_tf32_kernel<true, true><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr, y.data_ptr<float>(),
                (int)H, (int)W, (int)Ho, (int)Wo, M);
        } else {
            igemm_tf32_kernel<true, false><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr, y.data_ptr<float>(),
                (int)H, (int)W, (int)Ho, (int)Wo, M);
        }
    } else {
        if (mok) {
            igemm_tf32_kernel<false, true><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), nullptr, y.data_ptr<float>(),
                (int)H, (int)W, (int)Ho, (int)Wo, M);
        } else {
            igemm_tf32_kernel<false, false><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), nullptr, y.data_ptr<float>(),
                (int)H, (int)W, (int)Ho, (int)Wo, M);
        }
    }

    return y.view({N, (long)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 = not (
            groups == 1
            and st == (1, 1)
            and pd == (0, 0)
            and dl == (1, 1)
            and (kdim % 32) == 0
            and (int(out_channels) == 64 or (int(out_channels) % 16) == 0)
            and int(out_channels) >= 16
            and int(in_channels) >= 1
        )

        self._ext = None
        if not self._native:
            macros = {
                "IN_C": int(in_channels),
                "OC": int(out_channels),
                "KH": kh,
                "KW": kw,
            }
            src = _TEMPLATE
            for k, v in macros.items():
                src = src.replace("{%s}" % k, str(v))
            name = "igemm_tf32_c{}_o{}_k{}x{}_v3".format(
                in_channels, out_channels, kh, kw)
            self._ext = load_inline(
                name=name,
                cpp_sources=_CPP_SRC,
                cuda_sources=src,
                functions=["conv2d_forward"],
                verbose=False,
                extra_cuda_cflags=["-O3", "--use_fast_math"],
            )

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