```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA implementation of the 2D convolution.
#
# Strategy: implicit-GEMM convolution accelerated with TF32 tensor-core
# (WMMA) MMA instructions.
#   * GEMM view:  C[M, N] = A[M, K] * B[K, N]
#       - M = batch * H_out * W_out  (tiled as BM=128)
#       - N = out_channels           (tiled as BN=64)
#       - K = in_channels * Kh * Kw  (padded to multiples of BK=48)
#   * The A-tile (im2col) is generated on the fly directly from the NCHW
#     input (coalesced 128-float column loads, zero-filled at the borders)
#     into shared memory in column-major layout so WMMA can consume it.
#   * The weights are first rearranged by a tiny custom kernel into a
#     [K][N] matrix (with zero padding) so B-tiles are loaded coalesced.
#   * Each thread block computes a 128x64 output tile with 8 warps
#     (4x2), each warp producing a 32x32 tile via m16n8k8 TF32 MMAs with
#     FP32 accumulation.
#   * The epilogue writes the NCHW output directly (coalesced, via a
#     shared-memory transpose), optionally adding bias.
# ---------------------------------------------------------------------------

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

using namespace nvcuda;

#define BM 128
#define BN 64
#define BK 48

// ---------------------------------------------------------------------------
// Rearrange convolution weights [Nout][C][Kh][Kw] into a GEMM-friendly
// [C*BK][Nout] matrix where k = c*BK + kh*Kw + kw and the last
// (BK - Kh*Kw) rows of every channel block are zero padded.
// ---------------------------------------------------------------------------
__global__ void weight_transform_kernel(const float* __restrict__ W,
                                        float* __restrict__ Wt,
                                        int C, int Nout, int Kh, int Kw) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int K = C * BK;
    int total = K * Nout;
    if (idx >= total) return;
    int n = idx % Nout;
    int k = idx / Nout;
    int c = k / BK;
    int r = k % BK;
    int R = Kh * Kw;
    float v = 0.0f;
    if (r < R) {
        int kh = r / Kw;
        int kw = r % Kw;
        v = W[((long long)(n * C + c) * Kh + kh) * Kw + kw];
    }
    Wt[idx] = v;
}

// ---------------------------------------------------------------------------
// Implicit GEMM convolution kernel (tensor cores, TF32, FP32 accumulate).
// X: NCHW input, Wt: [C*BK][Nout], Y: NCHW output.
// Requires: W_out % BM == 0, Nout % BN == 0, C % 4 == 0.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(256)
conv_im2col_gemm_kernel(const float* __restrict__ X,
                        const float* __restrict__ Wt,
                        const float* __restrict__ bias,
                        float* __restrict__ Y,
                        int Nimg, int H, int W, int C, int Nout,
                        int Ho, int Wo, int Kh, int Kw,
                        int sh, int sw, int ph, int pw,
                        int dh, int dw) {
    extern __shared__ float smem[];
    float* At = smem;                    // [2][BK][BM] column-major A tiles
    float* Bt = smem + 2 * BK * BM;      // [2][BK][BN] row-major B tiles
    float* stage = smem;                 // reused as [BN][BM] in epilogue

    const int m0 = blockIdx.x * BM;
    const int n0 = blockIdx.y * BN;
    const int HWo = Ho * Wo;
    const int img = m0 / HWo;
    const int s = m0 % HWo;
    const int ho = s / Wo;
    const int wo0 = s % Wo;

    const int tid = threadIdx.x;
    const int R = Kh * Kw;

    // ---- tile loaders ------------------------------------------------------
    auto load_tile = [&](int t, int buf) {
        const int c = t;  // one input channel per k-tile
        // A tile: BK columns x BM rows, column-major (col stride BM)
        #pragma unroll
        for (int i = 0; i < 6; ++i) {
            int idx = i * 256 + tid;     // 0..1535
            int k = idx >> 5;            // 0..47
            int f4m = (idx & 31) * 4;    // 0..124
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (k < R) {
                int kh = k / Kw;
                int kw = k % Kw;
                int hi = ho * sh + kh * dh - ph;
                int wi0 = wo0 * sw + kw * dw - pw;
                if (hi >= 0 && hi < H) {
                    const float* base = X + (((long long)img * C + c) * H + hi) * W;
                    if (wi0 >= 0 && wi0 + BM <= W) {
                        v = *reinterpret_cast<const float4*>(base + wi0 + f4m);
                    } else {
                        float tmp[4] = {0.f, 0.f, 0.f, 0.f};
                        #pragma unroll
                        for (int j = 0; j < 4; ++j) {
                            int wi = wi0 + f4m + j;
                            if (wi >= 0 && wi < W) tmp[j] = base[wi];
                        }
                        v = make_float4(tmp[0], tmp[1], tmp[2], tmp[3]);
                    }
                }
            }
            *reinterpret_cast<float4*>(&At[buf * BK * BM + k * BM + f4m]) = v;
        }
        // B tile: BK rows x BN cols, row-major
        #pragma unroll
        for (int i = 0; i < 3; ++i) {
            int idx = i * 256 + tid;     // 0..767
            int k = idx >> 4;            // 0..47
            int f4n = (idx & 15) * 4;    // 0..60
            const float* g = Wt + ((long long)(t * BK + k)) * Nout + n0 + f4n;
            *reinterpret_cast<float4*>(&Bt[buf * BK * BN + k * BN + f4n]) =
                *reinterpret_cast<const float4*>(g);
        }
    };

    // ---- accumulators ------------------------------------------------------
    const int warp = tid >> 5;
    const int wm = warp >> 1;   // 0..3
    const int wn = warp & 1;    // 0..1
    wmma::fragment<wmma::accumulator, 16, 8, 8, float> acc[2][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    load_tile(0, 0);
    __syncthreads();

    int buf = 0;
    for (int t = 0; t < C; ++t) {
        if (t + 1 < C) load_tile(t + 1, buf ^ 1);
        const float* As = At + buf * BK * BM;
        const float* Bs = Bt + buf * BK * BN;
        #pragma unroll
        for (int ks = 0; ks < 6; ++ks) {
            wmma::fragment<wmma::matrix_a, 16, 8, 8, wmma::precision::tf32,
                           wmma::col_major> a[2];
            wmma::fragment<wmma::matrix_b, 16, 8, 8, wmma::precision::tf32,
                           wmma::row_major> b[4];
            #pragma unroll
            for (int i = 0; i < 2; ++i)
                wmma::load_matrix_sync(a[i], As + (ks * 8) * BM + wm * 32 + i * 16, BM);
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                wmma::load_matrix_sync(b[j], Bs + (ks * 8) * BN + wn * 32 + j * 8, BN);
            #pragma unroll
            for (int i = 0; i < 2; ++i)
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    wmma::mma_sync(acc[i][j], a[i], b[j], acc[i][j]);
        }
        __syncthreads();
        buf ^= 1;
    }

    // ---- epilogue: stage transposed (col-major store) then write NCHW ------
    __syncthreads();
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::store_matrix_sync(stage + (wn * 32 + j * 8) * BM + wm * 32 + i * 16,
                                    acc[i][j], BM, wmma::mem_col_major);
    __syncthreads();

    float* Ybase = Y + (((long long)img * Nout + n0) * Ho + ho) * Wo + wo0;
    if (bias != nullptr) {
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            int idx = i * 256 + tid;     // 0..2047
            int c = idx >> 5;            // 0..63
            int f4m = (idx & 31) * 4;
            float4 v = *reinterpret_cast<const float4*>(&stage[c * BM + f4m]);
            float bv = bias[n0 + c];
            v.x += bv; v.y += bv; v.z += bv; v.w += bv;
            *reinterpret_cast<float4*>(Ybase + (long long)c * HWo + f4m) = v;
        }
    } else {
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            int idx = i * 256 + tid;
            int c = idx >> 5;
            int f4m = (idx & 31) * 4;
            *reinterpret_cast<float4*>(Ybase + (long long)c * HWo + f4m) =
                *reinterpret_cast<const float4*>(&stage[c * BM + f4m]);
        }
    }
}

// ---------------------------------------------------------------------------
// Host launcher
// ---------------------------------------------------------------------------
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias_opt,
                             int64_t sh, int64_t sw,
                             int64_t ph, int64_t pw,
                             int64_t dh, int64_t dw) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be fp32");
    TORCH_CHECK(x.is_contiguous() && w.is_contiguous(), "need contiguous tensors");

    const int Nimg = x.size(0);
    const int C = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);
    const int Nout = w.size(0);
    const int Kh = w.size(2);
    const int Kw = w.size(3);

    const int Ho = (H + 2 * (int)ph - (int)dh * (Kh - 1) - 1) / (int)sh + 1;
    const int Wo = (W + 2 * (int)pw - (int)dw * (Kw - 1) - 1) / (int)sw + 1;

    auto y = torch::empty({Nimg, Nout, Ho, Wo}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    // transform weights into [C*BK][Nout]
    auto wt = torch::empty({(long long)C * BK, Nout}, x.options());
    int total = C * BK * Nout;
    weight_transform_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        w.data_ptr<float>(), wt.data_ptr<float>(), C, Nout, Kh, Kw);

    const int M = Nimg * Ho * Wo;
    dim3 grid(M / BM, Nout / BN);
    dim3 block(256);
    const int smem_bytes = (2 * BK * BM + 2 * BK * BN) * sizeof(float);  // 72KB

    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(conv_im2col_gemm_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             smem_bytes);
        configured = true;
    }

    const float* bias_ptr = nullptr;
    torch::Tensor bias;
    if (bias_opt.has_value() && bias_opt->defined()) {
        bias = bias_opt->contiguous();
        bias_ptr = bias.data_ptr<float>();
    }

    conv_im2col_gemm_kernel<<<grid, block, smem_bytes, stream>>>(
        x.data_ptr<float>(), wt.data_ptr<float>(), bias_ptr,
        y.data_ptr<float>(),
        Nimg, H, W, C, Nout, Ho, Wo, Kh, Kw,
        (int)sh, (int)sw, (int)ph, (int)pw, (int)dh, (int)dw);

    return y;
}
"""

cpp_source = """
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias_opt,
                             int64_t sh, int64_t sw,
                             int64_t ph, int64_t pw,
                             int64_t dh, int64_t dw);
"""

_conv_ext = load_inline(
    name="conv2d_implicit_gemm_tf32",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


def _as_pair(v):
    if isinstance(v, tuple):
        return v
    return (v, v)


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.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, bias=bias)
        kh, kw = _as_pair(kernel_size)
        ph, pw = _as_pair(padding)
        dh, dw = _as_pair(dilation)
        sh, sw = _as_pair(stride) if isinstance(stride, tuple) else (stride, stride)
        self.kernel_size_ = (kh, kw)
        self.padding_ = (ph, pw)
        self.dilation_ = (dh, dw)
        self.stride_ = (sh, sw)
        self._use_fast = self._check_fast_path()

    def _check_fast_path(self):
        if not torch.cuda.is_available():
            return False
        cap = torch.cuda.get_device_capability()
        if cap[0] < 8:  # TF32 tensor cores require sm80+
            return False
        c = self.conv2d
        if c.in_channels % 4 != 0 or c.out_channels % 64 != 0:
            return False
        if self.kernel_size_[0] * self.kernel_size_[1] > 48:
            return False
        return True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv2d
        if self._use_fast and x.is_cuda and x.dtype == torch.float32:
            x = x.contiguous()
            N, C, H, W = x.shape
            sh, sw = self.stride_
            ph, pw = self.padding_
            dh, dw = self.dilation_
            kh, kw = self.kernel_size_
            ho = (H + 2 * ph - dh * (kh - 1) - 1) // sh + 1
            wo = (W + 2 * pw - dw * (kw - 1) - 1) // sw + 1
            # fast path requires the output width to tile cleanly into 128
            # and channels to match the GEMM tiling
            if (wo % 128 == 0 and wo > 0 and ho > 0 and C % 4 == 0
                    and c.out_channels % 64 == 0):
                bias = c.bias if c.bias is not None else None
                return _conv_ext.conv2d_forward(
                    x, c.weight.contiguous(), bias, sh, sw, ph, pw, dh, dw)
        return self.conv2d(x)
```