```
# 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_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 64
#define BK 32
#define LDA 36   // BK + 4  (padding against bank conflicts, keeps 16B alignment)
#define LDB 68   // BN + 4
#define LDC 68   // BN + 4
#define NTHREADS 256

// Union of the smem regions used by the main loop (As+Bs) and the epilogue (Cs).
template <int A, int B> struct SMax { static const int value = A > B ? A : B; };

#define SMEM_AB_FLOATS (BM * LDA + BK * LDB)   // 6784 floats
#define SMEM_C_FLOATS  (BM * LDC)              // 8704 floats
#define SMEM_MAIN_BYTES (SMax<SMEM_AB_FLOATS, SMEM_C_FLOATS>::value * 4)

// Implicit-GEMM convolution (stride=1, pad=0, dilation=1, groups=1) using
// TF32 tensor cores via the WMMA API. C[M,N] = A[M,K] * B[K,N] with
// M = Nb*Ho*Wo, N = Cout, K = Cin*KH*KW, where A is the im2col of the input
// materialized on the fly into shared memory.
__global__ void __launch_bounds__(NTHREADS) implicit_gemm_conv2d_tf32(
    const float* __restrict__ input,     // [Nb, Cin, H, W]
    const float* __restrict__ weight_t,  // [K, N] row-major (tf32 pre-rounded)
    const float* __restrict__ bias,      // [N] or nullptr
    float* __restrict__ output,          // [Nb, Cout, Ho, Wo]
    int Nb, int Cin, int H, int W,
    int Cout, int Ho, int Wo,
    int KH, int KW,
    int M, int N, int K)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    extern __shared__ char smem_raw[];
    int* off_k = reinterpret_cast<int*>(smem_raw);          // K ints
    const size_t offk_bytes = (((size_t)K * sizeof(int)) + 15) & ~((size_t)15);
    float* As = reinterpret_cast<float*>(smem_raw + offk_bytes); // BM x LDA
    float* Bs = As + BM * LDA;                                     // BK x LDB
    float* Cs = As;                                                // BM x LDC (reuse)

    const int tid = threadIdx.x;
    const int m_base = blockIdx.x * BM;
    const int n_base = blockIdx.y * BN;

    // Table: linear k -> offset (ic*H + ky)*W + kx inside one sample's input plane.
    for (int gk = tid; gk < K; gk += NTHREADS) {
        int ic = gk / (KH * KW);
        int r  = gk - ic * (KH * KW);
        int ky = r / KW;
        int kx = r - ky * KW;
        off_k[gk] = (ic * H + ky) * W + kx;
    }
    __syncthreads();

    // Per-thread A row bases (m -> nb, oh, ow), constant across the K loop.
    int baseA[BM * BK / NTHREADS];  // 16
    #pragma unroll
    for (int i = 0; i < BM * BK / NTHREADS; i++) {
        int idx = tid + i * NTHREADS;
        int m = idx & (BM - 1);
        int gm = m_base + m;
        if (gm < M) {
            int nb  = gm / (Ho * Wo);
            int rem = gm - nb * (Ho * Wo);
            int oh  = rem / Wo;
            int ow  = rem - oh * Wo;
            baseA[i] = (nb * Cin * H + oh) * W + ow;
        } else {
            baseA[i] = -1;
        }
    }

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2][2];
    #pragma unroll
    for (int i = 0; i < 2; i++)
        #pragma unroll
        for (int j = 0; j < 2; j++)
            wmma::fill_fragment(acc[i][j], 0.0f);

    const int warp   = tid >> 5;
    const int warp_m = (warp & 3) * 32;   // 4 warps along M
    const int warp_n = (warp >> 2) * 32;  // 2 warps along N

    float a_reg[BM * BK / NTHREADS];  // 16
    float b_reg[BK * BN / NTHREADS];  // 8

    const int KTiles = (K + BK - 1) / BK;

    auto load_tile = [&](int kt) {
        const int k_base = kt * BK;
        #pragma unroll
        for (int i = 0; i < BM * BK / NTHREADS; i++) {
            int idx = tid + i * NTHREADS;
            int k = idx >> 7;           // idx / BM
            int gk = k_base + k;
            float v = 0.f;
            if (baseA[i] >= 0 && gk < K) v = input[baseA[i] + off_k[gk]];
            a_reg[i] = v;
        }
        #pragma unroll
        for (int i = 0; i < BK * BN / NTHREADS; i++) {
            int idx = tid + i * NTHREADS;
            int n = idx & (BN - 1);
            int k = idx / BN;
            int gk = k_base + k;
            int gn = n_base + n;
            float v = 0.f;
            if (gk < K && gn < N) v = weight_t[gk * N + gn];
            b_reg[i] = v;
        }
    };

    auto store_tile = [&]() {
        #pragma unroll
        for (int i = 0; i < BM * BK / NTHREADS; i++) {
            int idx = tid + i * NTHREADS;
            int m = idx & (BM - 1);
            int k = idx >> 7;
            As[m * LDA + k] = a_reg[i];
        }
        #pragma unroll
        for (int i = 0; i < BK * BN / NTHREADS; i++) {
            int idx = tid + i * NTHREADS;
            int n = idx & (BN - 1);
            int k = idx / BN;
            Bs[k * LDB + n] = b_reg[i];
        }
    };

    load_tile(0);
    store_tile();
    __syncthreads();

    for (int kt = 0; kt < KTiles; kt++) {
        if (kt + 1 < KTiles) load_tile(kt + 1);   // prefetch into registers

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag[2];
            #pragma unroll
            for (int i = 0; i < 2; i++)
                wmma::load_matrix_sync(a_frag[i], &As[(warp_m + i * 16) * LDA + kk], LDA);
            #pragma unroll
            for (int j = 0; j < 2; j++)
                wmma::load_matrix_sync(b_frag[j], &Bs[kk * LDB + warp_n + j * 16], LDB);
            #pragma unroll
            for (int i = 0; i < 2; i++)
                #pragma unroll
                for (int j = 0; j < 2; j++)
                    wmma::mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]);
        }

        __syncthreads();
        if (kt + 1 < KTiles) {
            store_tile();
            __syncthreads();
        }
    }

    // Epilogue: accumulators -> smem -> coalesced global stores.
    #pragma unroll
    for (int i = 0; i < 2; i++)
        #pragma unroll
        for (int j = 0; j < 2; j++)
            wmma::store_matrix_sync(&Cs[(warp_m + i * 16) * LDC + warp_n + j * 16],
                                    acc[i][j], LDC, wmma::mem_row_major);
    __syncthreads();

    const int HWo = Ho * Wo;
    #pragma unroll
    for (int i = 0; i < BM * BN / NTHREADS; i++) {   // 32
        int idx = tid + i * NTHREADS;
        int m = idx & (BM - 1);
        int n = idx >> 7;
        int gm = m_base + m;
        int gn = n_base + n;
        if (gm < M && gn < N) {
            int nb  = gm / HWo;
            int rem = gm - nb * HWo;
            int oh  = rem / Wo;
            int ow  = rem - oh * Wo;
            float v = Cs[m * LDC + n];
            if (bias != nullptr) v += bias[gn];
            output[((nb * N + gn) * Ho + oh) * Wo + ow] = v;
        }
    }
#endif
}

// Round fp32 -> tf32 (round to nearest even) keeping the value in fp32 format.
__global__ void round_to_tf32_kernel(float* __restrict__ x, long long n) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) x[i] = wmma::__float_to_tf32(x[i]);
#endif
}

torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight,
                             c10::optional<torch::Tensor> bias_opt,
                             int64_t stride, int64_t padding,
                             int64_t dilation, int64_t groups) {
    const at::cuda::CUDAGuard device_guard(input.device());
    auto inp = input.contiguous();
    auto wgt = weight.contiguous();
    torch::Tensor bias;
    bool has_bias = bias_opt.has_value() && bias_opt->defined();
    if (has_bias) bias = bias_opt->contiguous();

    const int Nb   = inp.size(0);
    const int Cin  = inp.size(1);
    const int H    = inp.size(2);
    const int W    = inp.size(3);
    const int Cout = wgt.size(0);
    const int KH   = wgt.size(2);
    const int KW   = wgt.size(3);
    const int Ho   = (H + 2 * (int)padding - (int)dilation * (KH - 1) - 1) / (int)stride + 1;
    const int Wo   = (W + 2 * (int)padding - (int)dilation * (KW - 1) - 1) / (int)stride + 1;

    static int dev_major = -1, dev_minor = 0;
    if (dev_major < 0) {
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, inp.get_device());
        dev_major = prop.major;
        dev_minor = prop.minor;
    }
    const bool sm_ok = (dev_major > 8) || (dev_major == 8 && dev_minor >= 0);

    const long long K = (long long)Cin * KH * KW;
    const bool fast_path = sm_ok &&
        inp.scalar_type() == at::kFloat && wgt.scalar_type() == at::kFloat &&
        stride == 1 && padding == 0 && dilation == 1 && groups == 1 &&
        wgt.size(1) == Cin && Ho > 0 && Wo > 0;

    if (!fast_path) {
        return at::conv2d(inp, wgt, bias_opt.has_value() ? *bias_opt : torch::Tensor(),
                          {stride, stride}, {padding, padding},
                          {dilation, dilation}, groups);
    }

    const long long M = (long long)Nb * Ho * Wo;
    const long long N = Cout;

    // [Cout, Cin*KH*KW] -> [K, N], pre-rounded to tf32.
    auto weight_t = wgt.reshape({Cout, K}).t().contiguous();
    {
        long long n = weight_t.numel();
        int threads = 256;
        long long blocks = (n + threads - 1) / threads;
        round_to_tf32_kernel<<<(unsigned)blocks, threads, 0,
            at::cuda::getCurrentCUDAStream()>>>(weight_t.data_ptr<float>(), n);
    }

    auto output = at::empty({Nb, Cout, Ho, Wo}, inp.options());

    const size_t offk_bytes = (((size_t)K * sizeof(int)) + 15) & ~((size_t)15);
    const size_t smem_bytes = offk_bytes + (size_t)SMEM_MAIN_BYTES;

    if (smem_bytes > 48 * 1024) {
        static bool attr_done = false;
        if (!attr_done) {
            cudaFuncSetAttribute(implicit_gemm_conv2d_tf32,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 160 * 1024);
            attr_done = true;
        }
    }

    dim3 grid((unsigned)((M + BM - 1) / BM), (unsigned)((N + BN - 1) / BN));
    dim3 block(NTHREADS);
    const float* bias_ptr = has_bias ? bias.data_ptr<float>() : nullptr;

    implicit_gemm_conv2d_tf32<<<grid, block, smem_bytes,
        at::cuda::getCurrentCUDAStream()>>>(
        inp.data_ptr<float>(), weight_t.data_ptr<float>(), bias_ptr,
        output.data_ptr<float>(),
        Nb, Cin, H, W, Cout, Ho, Wo, KH, KW,
        (int)M, (int)N, (int)K);

    return output;
}
"""

cpp_source = """
torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight,
                             c10::optional<torch::Tensor> bias_opt,
                             int64_t stride, int64_t padding,
                             int64_t dilation, int64_t groups);
"""

custom_conv = load_inline(
    name="implicit_gemm_conv2d_tf32",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv2d_forward"],
    verbose=True,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Conv2d accelerated with a custom implicit-GEMM CUDA kernel that uses
    TF32 tensor cores (WMMA) on sm_80+ GPUs. Unsupported configurations fall
    back to the standard cuDNN implementation.
    """
    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)
        s = self.conv2d.stride
        p = self.conv2d.padding
        d = self.conv2d.dilation
        self._stride = s[0] if isinstance(s, tuple) else s
        self._padding = p[0] if isinstance(p, tuple) else p
        self._dilation = d[0] if isinstance(d, tuple) else d
        self._groups = self.conv2d.groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return custom_conv.conv2d_forward(
            x, self.conv2d.weight, self.conv2d.bias,
            self._stride, self._padding, self._dilation, self._groups)
```