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

_CPP_SRC = "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> bias);"

_CUDA_TEMPLATE = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <c10/util/Optional.h>

#define IN_C %(in_channels)d
#define OUT_C %(out_channels)d
#define KH %(kh)d
#define KW %(kw)d
#define STRIDE_H %(stride_h)d
#define STRIDE_W %(stride_w)d
#define PAD_H %(pad_h)d
#define PAD_W %(pad_w)d
#define DIL_H %(dil_h)d
#define DIL_W %(dil_w)d

// Each thread computes OUTS consecutive output columns for a single
// (batch, out_channel, out_row). Weights are broadcast across the warp and
// reused OUTS times; input loads are coalesced across lanes.
template <int OUTS>
__global__ void __launch_bounds__(256) conv2d_direct_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int IH, const int IW, const int OH, const int OW)
{
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int groups_per_row = (OW + OUTS - 1) / OUTS;
    const int rows_per_oc = OH * groups_per_row;
    const int total_groups = N * OUT_C * rows_per_oc;
    if (tid >= total_groups) return;

    const int g  = tid %% groups_per_row;
    const int oh = (tid / groups_per_row) %% OH;
    const int t1 = tid / rows_per_oc;
    const int oc = t1 %% OUT_C;
    const int n  = t1 / OUT_C;

    const float* __restrict__ xp = x + (long)n * IN_C * IH * IW;
    const float* __restrict__ wp = w + (long)oc * IN_C * KH * KW;

    const int ow0 = g * OUTS;
    const int h0 = oh * STRIDE_H - PAD_H;
    const int w0 = ow0 * STRIDE_W - PAD_W;

    float acc[OUTS];
    #pragma unroll
    for (int m = 0; m < OUTS; ++m) acc[m] = 0.0f;

    #pragma unroll 2
    for (int ic = 0; ic < IN_C; ++ic) {
        const float* __restrict__ xic = xp + (long)ic * IH * IW;
        const float* __restrict__ wic = wp + ic * KH * KW;
        #pragma unroll
        for (int r = 0; r < KH; ++r) {
            const int ih = h0 + r * DIL_H;
            if ((unsigned)ih >= (unsigned)IH) continue;
            const float* __restrict__ xrow = xic + (long)ih * IW;
            #pragma unroll
            for (int c = 0; c < KW; ++c) {
                const float wv = __ldg(wic + r * KW + c);
                const int wb = w0 + c * DIL_W;
                #pragma unroll
                for (int m = 0; m < OUTS; ++m) {
                    const int iw = wb + m * STRIDE_W;
                    if ((unsigned)iw < (unsigned)IW) {
                        acc[m] = fmaf(__ldg(xrow + iw), wv, acc[m]);
                    }
                }
            }
        }
    }

    if (bias != nullptr) {
        const float bv = __ldg(bias + oc);
        #pragma unroll
        for (int m = 0; m < OUTS; ++m) acc[m] += bv;
    }

    float* __restrict__ yp = y + ((((long)n * OUT_C + oc) * OH) + oh) * OW;
    #pragma unroll
    for (int m = 0; m < OUTS; ++m) {
        if (ow0 + m < OW) yp[ow0 + m] = acc[m];
    }
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, c10::optional<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 float32 is supported");

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

    const int N  = (int)xc.size(0);
    const int IH = (int)xc.size(2);
    const int IW = (int)xc.size(3);
    const int OH = (IH + 2 * PAD_H - DIL_H * (KH - 1) - 1) / STRIDE_H + 1;
    const int OW = (IW + 2 * PAD_W - DIL_W * (KW - 1) - 1) / STRIDE_W + 1;

    auto y = torch::empty({N, OUT_C, OH, OW}, xc.options());
    const long total = (long)N * OUT_C * OH * OW;
    if (total == 0) return y;

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

    constexpr int OUTS = 4;
    const long groups_per_row = (OW + OUTS - 1) / OUTS;
    const long threads_needed = (long)N * OUT_C * OH * groups_per_row;
    const int block = 256;
    const long grid = (threads_needed + block - 1) / block;

    auto stream = at::cuda::getCurrentCUDAStream();
    conv2d_direct_kernel<OUTS><<<(unsigned)grid, block, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr, y.data_ptr<float>(),
        N, IH, IW, OH, OW);
    return y;
}
"""


def _pair(v):
    if isinstance(v, (tuple, list)):
        assert len(v) == 2
        return int(v[0]), int(v[1])
    return int(v), int(v)


class ModelNew(nn.Module):
    """
    2D convolution implemented with a custom, shape-specialized direct
    convolution CUDA kernel. The kernel is JIT-compiled with the convolution
    hyperparameters (channels, kernel size, stride, padding, dilation) baked
    in as compile-time constants so the inner loops are fully unrolled.
    """

    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__()
        # Keep an nn.Conv2d module so parameters/state_dict match the original.
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, bias=bias)

        kh, kw = _pair(kernel_size)
        sh, sw = _pair(stride)
        ph, pw = _pair(padding)
        dh, dw = _pair(dilation)

        params = dict(in_channels=in_channels, out_channels=out_channels,
                      kh=kh, kw=kw, stride_h=sh, stride_w=sw,
                      pad_h=ph, pad_w=pw, dil_h=dh, dil_w=dw)
        cuda_src = _CUDA_TEMPLATE % params

        ext_name = ("custom_conv2d_%d_%d_%dx%d_%dx%d_%dx%d_%dx%d"
                    % (in_channels, out_channels, kh, kw, sh, sw, ph, pw, dh, dw))
        self._op = load_inline(
            name=ext_name,
            cpp_sources=_CPP_SRC,
            cuda_sources=cuda_src,
            functions=["conv2d_forward"],
            verbose=False,
            extra_cuda_cflags=["-O3"],
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32:
            return self._op.conv2d_forward(x, self.conv2d.weight, self.conv2d.bias)
        return self.conv2d(x)
```