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

# Custom CUDA implementation of 2D convolution.
# Fast path: shared-memory tiled direct convolution with register blocking
# (each thread produces OC_PT x OH_PT x OW_PT outputs) specialized for
# stride=1, dilation=1, groups=1. A fully general fallback kernel handles
# any other configuration.
conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TX 16
#define TY 16
#define NTHREADS (TX * TY)
#define OW_PT 4   // outputs per thread along width
#define OH_PT 2   // outputs per thread along height
#define OC_PT 4   // output channels per block
#define TILE_W (TX * OW_PT)
#define TILE_H (TY * OH_PT)

template <int KH, int KW>
__global__ void __launch_bounds__(NTHREADS)
conv2d_tiled_kernel(const float* __restrict__ in,
                    const float* __restrict__ wt,
                    const float* __restrict__ bias,
                    float* __restrict__ out,
                    int N, int Cin, int H, int W,
                    int Cout, int Hout, int Wout,
                    int padH, int padW) {
    constexpr int SH_H = TILE_H + KH - 1;
    constexpr int SH_W_RAW = TILE_W + KW - 1;
    constexpr int SH_W = (SH_W_RAW + 3) & ~3;   // padded to a multiple of 4 floats

    extern __shared__ float smem[];
    const int ckw = Cin * KH * KW;
    const int wsize = (OC_PT * ckw + 3) & ~3;
    float* sh_w = smem;          // staged weights for OC_PT output channels
    float* sh_in = smem + wsize; // staged input tile

    const int nocb = (Cout + OC_PT - 1) / OC_PT;
    const int n = blockIdx.z / nocb;
    const int oc_base = (blockIdx.z - n * nocb) * OC_PT;
    const int th0 = blockIdx.y * TILE_H;
    const int tw0 = blockIdx.x * TILE_W;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * TX + tx;

    // Stage weights into shared memory.
    #pragma unroll
    for (int o = 0; o < OC_PT; ++o) {
        const int oc = oc_base + o;
        float* dst = sh_w + o * ckw;
        if (oc < Cout) {
            const float* src = wt + (long long)oc * ckw;
            for (int i = tid; i < ckw; i += NTHREADS) dst[i] = src[i];
        } else {
            for (int i = tid; i < ckw; i += NTHREADS) dst[i] = 0.0f;
        }
    }

    const int oh0 = th0 + ty * OH_PT;
    const int ow0 = tw0 + tx * OW_PT;

    float acc[OC_PT][OH_PT][OW_PT];
    #pragma unroll
    for (int o = 0; o < OC_PT; ++o)
        #pragma unroll
        for (int i = 0; i < OH_PT; ++i)
            #pragma unroll
            for (int j = 0; j < OW_PT; ++j)
                acc[o][i][j] = 0.0f;

    const float* in_n = in + (long long)n * Cin * H * W;

    for (int ic = 0; ic < Cin; ++ic) {
        const float* in_ch = in_n + (long long)ic * H * W;

        // Cooperatively stage the input tile (zero padding outside the image).
        for (int i = tid; i < SH_H * SH_W; i += NTHREADS) {
            const int r = i / SH_W;
            const int c = i - r * SH_W;
            float v = 0.0f;
            if (c < SH_W_RAW) {
                const int ih = th0 + r - padH;
                const int iw = tw0 + c - padW;
                if (ih >= 0 && ih < H && iw >= 0 && iw < W)
                    v = in_ch[(long long)ih * W + iw];
            }
            sh_in[i] = v;
        }
        __syncthreads();

        if (oh0 < Hout) {
            const float* shw_ic = sh_w + ic * KH * KW;
            #pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
                float r[OH_PT][KW + OW_PT - 1];
                #pragma unroll
                for (int dr = 0; dr < OH_PT; ++dr) {
                    const float* shrow =
                        &sh_in[(ty * OH_PT + dr + kh) * SH_W + tx * OW_PT];
                    if ((KW + OW_PT - 1) % 4 == 0) {
                        #pragma unroll
                        for (int i2 = 0; i2 < (KW + OW_PT - 1) / 4; ++i2) {
                            const float4 v =
                                *reinterpret_cast<const float4*>(shrow + 4 * i2);
                            r[dr][4 * i2 + 0] = v.x;
                            r[dr][4 * i2 + 1] = v.y;
                            r[dr][4 * i2 + 2] = v.z;
                            r[dr][4 * i2 + 3] = v.w;
                        }
                    } else {
                        #pragma unroll
                        for (int i2 = 0; i2 < KW + OW_PT - 1; ++i2)
                            r[dr][i2] = shrow[i2];
                    }
                }
                #pragma unroll
                for (int o = 0; o < OC_PT; ++o) {
                    const float* shw = shw_ic + o * ckw + kh * KW;
                    float wreg[KW];
                    #pragma unroll
                    for (int kw = 0; kw < KW; ++kw) wreg[kw] = shw[kw];
                    #pragma unroll
                    for (int kw = 0; kw < KW; ++kw) {
                        #pragma unroll
                        for (int dr = 0; dr < OH_PT; ++dr) {
                            #pragma unroll
                            for (int dc = 0; dc < OW_PT; ++dc) {
                                acc[o][dr][dc] =
                                    fmaf(wreg[kw], r[dr][kw + dc], acc[o][dr][dc]);
                            }
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    // Epilogue.
    #pragma unroll
    for (int o = 0; o < OC_PT; ++o) {
        const int oc = oc_base + o;
        if (oc >= Cout) continue;
        const float bv = (bias != nullptr) ? bias[oc] : 0.0f;
        #pragma unroll
        for (int dr = 0; dr < OH_PT; ++dr) {
            const int oh = oh0 + dr;
            if (oh >= Hout) continue;
            #pragma unroll
            for (int dc = 0; dc < OW_PT; ++dc) {
                const int ow = ow0 + dc;
                if (ow < Wout) {
                    out[(((long long)n * Cout + oc) * Hout + oh) * Wout + ow] =
                        acc[o][dr][dc] + bv;
                }
            }
        }
    }
}

// Fully general direct convolution (any stride/padding/dilation/groups).
__global__ void conv2d_generic_kernel(
    const float* __restrict__ in,
    const float* __restrict__ wt,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int Cin, int H, int W,
    int Cout, int Hout, int Wout,
    int KH, int KW,
    int strideH, int strideW,
    int padH, int padW,
    int dilH, int dilW,
    int groups, long long total) {
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int ow = (int)(idx % Wout);
    const int oh = (int)((idx / Wout) % Hout);
    const int oc = (int)((idx / ((long long)Wout * Hout)) % Cout);
    const int n  = (int)(idx / ((long long)Wout * Hout * Cout));

    const int cin_g = Cin / groups;
    const int cout_g = Cout / groups;
    const int g = oc / cout_g;

    float sum = 0.0f;
    for (int ic = 0; ic < cin_g; ++ic) {
        const int icg = g * cin_g + ic;
        for (int kh = 0; kh < KH; ++kh) {
            const int ih = oh * strideH + kh * dilH - padH;
            if (ih < 0 || ih >= H) continue;
            for (int kw = 0; kw < KW; ++kw) {
                const int iw = ow * strideW + kw * dilW - padW;
                if (iw < 0 || iw >= W) continue;
                sum += in[((long long)(n * Cin + icg) * H + ih) * W + iw] *
                       wt[((long long)(oc * cin_g + ic) * KH + kh) * KW + kw];
            }
        }
    }
    if (bias != nullptr) sum += bias[oc];
    out[idx] = sum;
}

torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight,
                             torch::Tensor bias,
                             int64_t strideH, int64_t strideW,
                             int64_t padH, int64_t padW,
                             int64_t dilH, int64_t dilW,
                             int64_t groups) {
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "bad tensor ranks");

    auto x = input.contiguous();
    auto wgt = weight.contiguous();

    const int N = (int)x.size(0), Cin = (int)x.size(1);
    const int H = (int)x.size(2), W = (int)x.size(3);
    const int Cout = (int)wgt.size(0);
    const int KH = (int)wgt.size(2), KW = (int)wgt.size(3);
    const int Hout = (int)((H + 2 * padH - dilH * (KH - 1) - 1) / strideH + 1);
    const int Wout = (int)((W + 2 * padW - dilW * (KW - 1) - 1) / strideW + 1);

    auto out = at::empty({N, Cout, Hout, Wout}, x.options());
    const long long total = (long long)N * Cout * Hout * Wout;
    if (total == 0) return out;

    const float* bias_ptr =
        (bias.defined() && bias.numel() == (long long)Cout)
            ? bias.data_ptr<float>() : nullptr;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const bool tiled_ok = (groups == 1) && strideH == 1 && strideW == 1 &&
                          dilH == 1 && dilW == 1 &&
                          (long long)N * ((Cout + OC_PT - 1) / OC_PT) <= 65535;

#define LAUNCH_TILED(KHV, KWV)                                                  \
    do {                                                                        \
        constexpr int KHc = KHV;                                                \
        constexpr int KWc = KWV;                                                \
        constexpr int SH_H = TILE_H + KHc - 1;                                  \
        constexpr int SH_W = ((TILE_W + KWc - 1) + 3) & ~3;                     \
        const size_t wsize =                                                    \
            (((size_t)OC_PT * Cin * KHc * KWc) + 3) & ~(size_t)3;               \
        const size_t smem = (wsize + (size_t)SH_H * SH_W) * sizeof(float);      \
        if (smem <= 48 * 1024) {                                                \
            const int nocb = (Cout + OC_PT - 1) / OC_PT;                        \
            dim3 block(TX, TY);                                                 \
            dim3 grid((Wout + TILE_W - 1) / TILE_W,                             \
                      (Hout + TILE_H - 1) / TILE_H, N * nocb);                  \
            conv2d_tiled_kernel<KHc, KWc><<<grid, block, smem, stream>>>(       \
                x.data_ptr<float>(), wgt.data_ptr<float>(), bias_ptr,           \
                out.data_ptr<float>(), N, Cin, H, W, Cout, Hout, Wout,          \
                (int)padH, (int)padW);                                          \
            return out;                                                         \
        }                                                                       \
    } while (0)

    if (tiled_ok) {
        if (KH == 5 && KW == 9) LAUNCH_TILED(5, 9);
        if (KH == 3 && KW == 3) LAUNCH_TILED(3, 3);
        if (KH == 1 && KW == 1) LAUNCH_TILED(1, 1);
        if (KH == 5 && KW == 5) LAUNCH_TILED(5, 5);
        if (KH == 7 && KW == 7) LAUNCH_TILED(7, 7);
    }
#undef LAUNCH_TILED

    // Generic fallback.
    const int threads = 256;
    const long long blocks = (total + threads - 1) / threads;
    conv2d_generic_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), wgt.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(), N, Cin, H, W, Cout, Hout, Wout, KH, KW,
        (int)strideH, (int)strideW, (int)padH, (int)padW,
        (int)dilH, (int)dilW, (int)groups, total);
    return out;
}
"""

conv2d_cpp_source = (
    "torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t strideH, int64_t strideW, int64_t padH, "
    "int64_t padW, int64_t dilH, int64_t dilW, int64_t groups);"
)

conv2d_ext = load_inline(
    name="conv2d_tiled_ext",
    cpp_sources=conv2d_cpp_source,
    cuda_sources=conv2d_cuda_source,
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
)


class ModelNew(nn.Module):
    """
    Conv2d accelerated with a custom shared-memory tiled CUDA kernel.
    """

    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__()
        # Keep an nn.Conv2d purely as the parameter holder so initialization
        # matches the reference module exactly.
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv2d
        bias = conv.bias if conv.bias is not None else x.new_empty(0)
        return conv2d_ext.conv2d_forward(
            x, conv.weight, bias,
            conv.stride[0], conv.stride[1],
            conv.padding[0], conv.padding[1],
            conv.dilation[0], conv.dilation[1],
            conv.groups)
```