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

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

#define OC_BLOCK 4
#define TILE_W 64
#define TILE_H 16
#define REG_W 4
#define REG_H 2
#define THREADS_X (TILE_W / REG_W)
#define THREADS_Y (TILE_H / REG_H)
#define NUM_THREADS (THREADS_X * THREADS_Y)

// ---------------------------------------------------------------------------
// Tiled direct convolution kernel (stride=1, dilation=1) with compile-time
// kernel size. Each block computes OC_BLOCK output channels x a TILE_H x
// TILE_W output tile. Each thread computes REG_H x REG_W outputs for each of
// the OC_BLOCK output channels. Input tile + weights are staged in shared
// memory, accumulators live in registers.
// ---------------------------------------------------------------------------
template <int KH, int KW>
__global__ void conv2d_direct_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int C, const int H, const int W,
    const int OC, const int Ho, const int Wo,
    const int padH, const int padW,
    const int oc_blocks,
    const bool has_bias)
{
    constexpr int SMEM_W = TILE_W + KW - 1;
    constexpr int SMEM_H = TILE_H + KH - 1;
    constexpr int W_SH_SIZE = OC_BLOCK * KH * KW;

    __shared__ float tile[SMEM_H * SMEM_W];
    __shared__ float w_sh[W_SH_SIZE];

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

    const int n  = blockIdx.z / oc_blocks;
    const int o0 = (blockIdx.z % oc_blocks) * OC_BLOCK;

    const int out_x0 = blockIdx.x * TILE_W;
    const int out_y0 = blockIdx.y * TILE_H;
    const int in_x0  = out_x0 - padW;
    const int in_y0  = out_y0 - padH;

    const long long x_img_stride = (long long)C * H * W;
    const long long w_oc_stride  = (long long)C * KH * KW;

    float acc[OC_BLOCK][REG_H][REG_W];
    #pragma unroll
    for (int o = 0; o < OC_BLOCK; ++o)
        #pragma unroll
        for (int i = 0; i < REG_H; ++i)
            #pragma unroll
            for (int j = 0; j < REG_W; ++j)
                acc[o][i][j] = 0.f;

    for (int c = 0; c < C; ++c) {
        const float* xc = x + (long long)n * x_img_stride + (long long)c * H * W;

        // Stage input tile (with halo) into shared memory.
        #pragma unroll
        for (int idx = tid; idx < SMEM_H * SMEM_W; idx += NUM_THREADS) {
            const int r   = idx / SMEM_W;
            const int col = idx % SMEM_W;
            const int gy  = in_y0 + r;
            const int gx  = in_x0 + col;
            float v = 0.f;
            if (gy >= 0 && gy < H && gx >= 0 && gx < W)
                v = xc[(long long)gy * W + gx];
            tile[idx] = v;
        }

        // Stage the OC_BLOCK x KH x KW weights for this input channel.
        #pragma unroll
        for (int idx = tid; idx < W_SH_SIZE; idx += NUM_THREADS) {
            const int oc = idx / (KH * KW);
            const int r  = idx % (KH * KW);
            float v = 0.f;
            if (o0 + oc < OC)
                v = w[(long long)(o0 + oc) * w_oc_stride + (long long)c * KH * KW + r];
            w_sh[idx] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                float tv[REG_H][REG_W];
                #pragma unroll
                for (int i = 0; i < REG_H; ++i)
                    #pragma unroll
                    for (int j = 0; j < REG_W; ++j)
                        tv[i][j] = tile[(ty * REG_H + i + kh) * SMEM_W + tx * REG_W + j + kw];
                #pragma unroll
                for (int o = 0; o < OC_BLOCK; ++o) {
                    const float wv = w_sh[o * KH * KW + kh * KW + kw];
                    #pragma unroll
                    for (int i = 0; i < REG_H; ++i)
                        #pragma unroll
                        for (int j = 0; j < REG_W; ++j)
                            acc[o][i][j] = fmaf(tv[i][j], wv, acc[o][i][j]);
                }
            }
        }
        __syncthreads();
    }

    if (has_bias) {
        #pragma unroll
        for (int o = 0; o < OC_BLOCK; ++o) {
            const float bv = (o0 + o < OC) ? bias[o0 + o] : 0.f;
            #pragma unroll
            for (int i = 0; i < REG_H; ++i)
                #pragma unroll
                for (int j = 0; j < REG_W; ++j)
                    acc[o][i][j] += bv;
        }
    }

    #pragma unroll
    for (int o = 0; o < OC_BLOCK; ++o) {
        const int og = o0 + o;
        if (og >= OC) continue;
        #pragma unroll
        for (int i = 0; i < REG_H; ++i) {
            const int ho = out_y0 + ty * REG_H + i;
            if (ho >= Ho) continue;
            #pragma unroll
            for (int j = 0; j < REG_W; ++j) {
                const int wo = out_x0 + tx * REG_W + j;
                if (wo >= Wo) continue;
                y[(((long long)n * OC + og) * Ho + ho) * Wo + wo] = acc[o][i][j];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Generic tiled kernel: same strategy but runtime kernel size (stride=1,
// dilation=1). Uses dynamic shared memory.
// ---------------------------------------------------------------------------
__global__ void conv2d_direct_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int C, const int H, const int W,
    const int OC, const int KH, const int KW,
    const int Ho, const int Wo,
    const int padH, const int padW,
    const int oc_blocks,
    const bool has_bias)
{
    extern __shared__ float smem[];
    const int SMEM_W = TILE_W + KW - 1;
    const int SMEM_H = TILE_H + KH - 1;
    const int W_SH_SIZE = OC_BLOCK * KH * KW;
    float* tile = smem;
    float* w_sh = smem + SMEM_H * SMEM_W;

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

    const int n  = blockIdx.z / oc_blocks;
    const int o0 = (blockIdx.z % oc_blocks) * OC_BLOCK;

    const int out_x0 = blockIdx.x * TILE_W;
    const int out_y0 = blockIdx.y * TILE_H;
    const int in_x0  = out_x0 - padW;
    const int in_y0  = out_y0 - padH;

    const long long x_img_stride = (long long)C * H * W;
    const long long w_oc_stride  = (long long)C * KH * KW;

    float acc[OC_BLOCK][REG_H][REG_W];
    #pragma unroll
    for (int o = 0; o < OC_BLOCK; ++o)
        #pragma unroll
        for (int i = 0; i < REG_H; ++i)
            #pragma unroll
            for (int j = 0; j < REG_W; ++j)
                acc[o][i][j] = 0.f;

    for (int c = 0; c < C; ++c) {
        const float* xc = x + (long long)n * x_img_stride + (long long)c * H * W;

        for (int idx = tid; idx < SMEM_H * SMEM_W; idx += NUM_THREADS) {
            const int r   = idx / SMEM_W;
            const int col = idx % SMEM_W;
            const int gy  = in_y0 + r;
            const int gx  = in_x0 + col;
            float v = 0.f;
            if (gy >= 0 && gy < H && gx >= 0 && gx < W)
                v = xc[(long long)gy * W + gx];
            tile[idx] = v;
        }

        for (int idx = tid; idx < W_SH_SIZE; idx += NUM_THREADS) {
            const int oc = idx / (KH * KW);
            const int r  = idx % (KH * KW);
            float v = 0.f;
            if (o0 + oc < OC)
                v = w[(long long)(o0 + oc) * w_oc_stride + (long long)c * KH * KW + r];
            w_sh[idx] = v;
        }
        __syncthreads();

        for (int kh = 0; kh < KH; ++kh) {
            for (int kw = 0; kw < KW; ++kw) {
                float tv[REG_H][REG_W];
                #pragma unroll
                for (int i = 0; i < REG_H; ++i)
                    #pragma unroll
                    for (int j = 0; j < REG_W; ++j)
                        tv[i][j] = tile[(ty * REG_H + i + kh) * SMEM_W + tx * REG_W + j + kw];
                for (int o = 0; o < OC_BLOCK; ++o) {
                    const float wv = w_sh[o * KH * KW + kh * KW + kw];
                    #pragma unroll
                    for (int i = 0; i < REG_H; ++i)
                        #pragma unroll
                        for (int j = 0; j < REG_W; ++j)
                            acc[o][i][j] = fmaf(tv[i][j], wv, acc[o][i][j]);
                }
            }
        }
        __syncthreads();
    }

    if (has_bias) {
        #pragma unroll
        for (int o = 0; o < OC_BLOCK; ++o) {
            const float bv = (o0 + o < OC) ? bias[o0 + o] : 0.f;
            #pragma unroll
            for (int i = 0; i < REG_H; ++i)
                #pragma unroll
                for (int j = 0; j < REG_W; ++j)
                    acc[o][i][j] += bv;
        }
    }

    #pragma unroll
    for (int o = 0; o < OC_BLOCK; ++o) {
        const int og = o0 + o;
        if (og >= OC) continue;
        #pragma unroll
        for (int i = 0; i < REG_H; ++i) {
            const int ho = out_y0 + ty * REG_H + i;
            if (ho >= Ho) continue;
            #pragma unroll
            for (int j = 0; j < REG_W; ++j) {
                const int wo = out_x0 + tx * REG_W + j;
                if (wo >= Wo) continue;
                y[(((long long)n * OC + og) * Ho + ho) * Wo + wo] = acc[o][i][j];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Fully general fallback: one output element per thread. Supports arbitrary
// stride / padding / dilation (groups = 1).
// ---------------------------------------------------------------------------
__global__ void conv2d_naive_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int C, const int H, const int W,
    const int OC, const int KH, const int KW,
    const int Ho, const int Wo,
    const int sH, const int sW,
    const int pH, const int pW,
    const int dH, const int dW,
    const bool has_bias)
{
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    const long long total = (long long)N * OC * Ho * Wo;
    if (idx >= total) return;

    const int wo = (int)(idx % Wo);
    const int ho = (int)((idx / Wo) % Ho);
    const int o  = (int)((idx / ((long long)Wo * Ho)) % OC);
    const int n  = (int)(idx / ((long long)Wo * Ho * OC));

    float sum = 0.f;
    for (int c = 0; c < C; ++c) {
        for (int kh = 0; kh < KH; ++kh) {
            const int ih = ho * sH - pH + kh * dH;
            if (ih < 0 || ih >= H) continue;
            for (int kw = 0; kw < KW; ++kw) {
                const int iw = wo * sW - pW + kw * dW;
                if (iw < 0 || iw >= W) continue;
                sum = fmaf(x[(((long long)n * C + c) * H + ih) * W + iw],
                           w[(((long long)o * C + c) * KH + kh) * KW + kw],
                           sum);
            }
        }
    }
    if (has_bias) sum += bias[o];
    y[idx] = sum;
}

// ---------------------------------------------------------------------------
// Host entry point.
// ---------------------------------------------------------------------------
torch::Tensor conv2d_forward(
    torch::Tensor x,
    torch::Tensor w,
    torch::Tensor bias,
    int64_t strideH, int64_t strideW,
    int64_t padH, int64_t padW,
    int64_t dilH, int64_t dilW)
{
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "conv2d_forward: inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "conv2d_forward: only float32 is supported");
    TORCH_CHECK(x.dim() == 4 && w.dim() == 4, "conv2d_forward: expected 4D tensors");

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int N  = (int)xc.size(0);
    const int C  = (int)xc.size(1);
    const int H  = (int)xc.size(2);
    const int W  = (int)xc.size(3);
    const int OC = (int)wc.size(0);
    const int KH = (int)wc.size(2);
    const int KW = (int)wc.size(3);
    TORCH_CHECK(wc.size(1) == C, "conv2d_forward: channel mismatch (groups must be 1)");

    const int Ho = (int)((H + 2 * padH - dilH * (KH - 1) - 1) / strideH + 1);
    const int Wo = (int)((W + 2 * padW - dilW * (KW - 1) - 1) / strideW + 1);

    auto y = at::empty({N, OC, (long)Ho, (long)Wo}, xc.options());
    const long long total = (long long)N * OC * Ho * Wo;
    if (total <= 0) return y;

    const bool has_bias = bias.defined() && bias.numel() > 0;
    torch::Tensor bc;
    const float* bptr = nullptr;
    if (has_bias) {
        bc = bias.contiguous();
        bptr = bc.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int oc_blocks = (OC + OC_BLOCK - 1) / OC_BLOCK;

    if (strideH == 1 && strideW == 1 && dilH == 1 && dilW == 1 &&
        (long long)N * oc_blocks <= 65535 &&
        (Wo + TILE_W - 1) / TILE_W <= 2147483647LL &&
        (Ho + TILE_H - 1) / TILE_H <= 65535) {

        dim3 block(THREADS_X, THREADS_Y);
        dim3 grid((Wo + TILE_W - 1) / TILE_W,
                  (Ho + TILE_H - 1) / TILE_H,
                  N * oc_blocks);

        const float* xp = xc.data_ptr<float>();
        const float* wp = wc.data_ptr<float>();
        float* yp = y.data_ptr<float>();

        if (KH == 5 && KW == 9) {
            conv2d_direct_kernel<5, 9><<<grid, block, 0, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, Ho, Wo, (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        } else if (KH == 3 && KW == 3) {
            conv2d_direct_kernel<3, 3><<<grid, block, 0, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, Ho, Wo, (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        } else if (KH == 1 && KW == 1) {
            conv2d_direct_kernel<1, 1><<<grid, block, 0, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, Ho, Wo, (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        } else if (KH == 5 && KW == 5) {
            conv2d_direct_kernel<5, 5><<<grid, block, 0, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, Ho, Wo, (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        } else if (KH == 7 && KW == 7) {
            conv2d_direct_kernel<7, 7><<<grid, block, 0, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, Ho, Wo, (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        } else if (KH == 9 && KW == 5) {
            conv2d_direct_kernel<9, 5><<<grid, block, 0, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, Ho, Wo, (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        }

        const size_t smem = ((size_t)(TILE_H + KH - 1) * (TILE_W + KW - 1) +
                             (size_t)OC_BLOCK * KH * KW) * sizeof(float);
        if (smem <= 46 * 1024) {
            conv2d_direct_generic_kernel<<<grid, block, smem, stream>>>(
                xp, wp, bptr, yp, N, C, H, W, OC, KH, KW, Ho, Wo,
                (int)padH, (int)padW, oc_blocks, has_bias);
            return y;
        }
    }

    // General fallback.
    const int threads = 256;
    const long long blocks = (total + threads - 1) / threads;
    conv2d_naive_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bptr, y.data_ptr<float>(),
        N, C, H, W, OC, KH, KW, Ho, Wo,
        (int)strideH, (int)strideW, (int)padH, (int)padW, (int)dilH, (int)dilW,
        has_bias);
    return y;
}
"""

conv2d_cpp_source = (
    "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "int64_t strideH, int64_t strideW, int64_t padH, int64_t padW, "
    "int64_t dilH, int64_t dilW);"
)

conv2d_ext = load_inline(
    name="fast_conv2d_ext",
    cpp_sources=conv2d_cpp_source,
    cuda_sources=conv2d_cuda_source,
    functions=["conv2d_forward"],
    verbose=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class _FastConv2d(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, stride, padding, dilation):
        ctx.stride = stride
        ctx.padding = padding
        ctx.dilation = dilation
        ctx.has_bias = bias is not None
        ctx.save_for_backward(x, weight, bias if bias is not None else torch.empty(0, device=x.device, dtype=x.dtype))
        bias_arg = bias if bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
        return conv2d_ext.conv2d_forward(
            x, weight, bias_arg,
            stride[0], stride[1],
            padding[0], padding[1],
            dilation[0], dilation[1],
        )

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, _bias = ctx.saved_tensors
        needs = ctx.needs_input_grad
        grad_x = grad_w = grad_b = None
        if needs[0]:
            grad_x = torch.nn.grad.conv2d_input(
                x.shape, weight, grad_output, ctx.stride, ctx.padding, ctx.dilation, 1
            )
        if needs[1]:
            grad_w = torch.nn.grad.conv2d_weight(
                x, weight.shape, grad_output, ctx.stride, ctx.padding, ctx.dilation, 1
            )
        if ctx.has_bias and needs[2]:
            grad_b = grad_output.sum(dim=(0, 2, 3))
        return grad_x, grad_w, grad_b, None, None, None


class ModelNew(nn.Module):
    """
    2D convolution accelerated with a custom tiled CUDA kernel.

    The forward pass runs a hand-written direct-convolution CUDA kernel
    (shared-memory tiling + register blocking over output channels and
    pixels). Falls back to the standard cuDNN path for unsupported
    configurations (groups != 1, non-CUDA or non-float32 tensors).
    """

    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)
        self.stride = self.conv2d.stride
        self.padding = self.conv2d.padding
        self.dilation = self.conv2d.dilation
        self.groups = self.conv2d.groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.conv2d.weight
        if (x.is_cuda and x.dim() == 4 and self.groups == 1
                and x.dtype == torch.float32 and weight.dtype == torch.float32):
            return _FastConv2d.apply(
                x, weight, self.conv2d.bias,
                self.stride, self.padding, self.dilation,
            )
        return self.conv2d(x)
```