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

# Direct tiled 2D convolution CUDA kernel.
# Each thread block computes a 32x32 tile of output pixels for one
# (batch, out_channel) pair. Threads cooperatively stage the required
# input halo region and the kernel weights in shared memory, and each
# thread accumulates 2x2 outputs in registers for data reuse.
conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

#define THREADS_X 16
#define THREADS_Y 16
#define OUT_X 2
#define OUT_Y 2
#define TILE_OUT_X (THREADS_X * OUT_X)
#define TILE_OUT_Y (THREADS_Y * OUT_Y)

__global__ void conv2d_direct_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int N, const int C_in, const int H_in, const int W_in,
    const int C_out, const int H_out, const int W_out,
    const int KH, const int KW,
    const int stride_h, const int stride_w,
    const int pad_h, const int pad_w,
    const int dil_h, const int dil_w,
    const int tile_h, const int tile_w)
{
    extern __shared__ float smem[];
    float* tile = smem;                       // tile_h * tile_w floats
    float* wgt  = smem + tile_h * tile_w;     // KH * KW floats

    const int oc = blockIdx.z % C_out;
    const int n  = blockIdx.z / C_out;

    const int out_x0 = blockIdx.x * TILE_OUT_X;
    const int out_y0 = blockIdx.y * TILE_OUT_Y;

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

    const int in_x0 = out_x0 * stride_w - pad_w;
    const int in_y0 = out_y0 * stride_h - pad_h;

    const float* in_nc = input + (int64_t)n * C_in * H_in * W_in;
    const float* w_oc  = weight + (int64_t)oc * C_in * KH * KW;

    float acc[OUT_Y][OUT_X];
    #pragma unroll
    for (int i = 0; i < OUT_Y; ++i) {
        #pragma unroll
        for (int j = 0; j < OUT_X; ++j) {
            acc[i][j] = 0.0f;
        }
    }

    const int tile_elems = tile_h * tile_w;
    const int w_elems = KH * KW;

    for (int ic = 0; ic < C_in; ++ic) {
        const float* in_c = in_nc + (int64_t)ic * H_in * W_in;
        const float* w_c  = w_oc + (int64_t)ic * KH * KW;

        // Cooperative load of the input tile (zero padded outside image).
        for (int idx = tid; idx < tile_elems; idx += nthreads) {
            const int yy = idx / tile_w;
            const int xx = idx - yy * tile_w;
            const int iy = in_y0 + yy;
            const int ix = in_x0 + xx;
            float v = 0.0f;
            if ((unsigned)iy < (unsigned)H_in && (unsigned)ix < (unsigned)W_in) {
                v = in_c[iy * W_in + ix];
            }
            tile[idx] = v;
        }
        // Cooperative load of this channel's weights.
        for (int idx = tid; idx < w_elems; idx += nthreads) {
            wgt[idx] = w_c[idx];
        }
        __syncthreads();

        const int base_y = ty * OUT_Y * stride_h;
        const int base_x = tx * OUT_X * stride_w;

        for (int kh = 0; kh < KH; ++kh) {
            const int row_y = base_y + kh * dil_h;
            const float* wrow = wgt + kh * KW;
            for (int kw = 0; kw < KW; ++kw) {
                const float wv = wrow[kw];
                const int col = base_x + kw * dil_w;
                #pragma unroll
                for (int i = 0; i < OUT_Y; ++i) {
                    const float* trow = tile + (row_y + i * stride_h) * tile_w + col;
                    #pragma unroll
                    for (int j = 0; j < OUT_X; ++j) {
                        acc[i][j] += trow[j * stride_w] * wv;
                    }
                }
            }
        }
        __syncthreads();
    }

    float bv = 0.0f;
    if (bias != nullptr) {
        bv = bias[oc];
    }

    const int64_t out_plane = (int64_t)H_out * W_out;
    #pragma unroll
    for (int i = 0; i < OUT_Y; ++i) {
        const int oy = out_y0 + ty * OUT_Y + i;
        if (oy < H_out) {
            #pragma unroll
            for (int j = 0; j < OUT_X; ++j) {
                const int ox = out_x0 + tx * OUT_X + j;
                if (ox < W_out) {
                    output[((int64_t)n * C_out + oc) * out_plane + (int64_t)oy * W_out + ox] = acc[i][j] + bv;
                }
            }
        }
    }
}

torch::Tensor conv2d_forward(
    torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
    int64_t stride_h, int64_t stride_w,
    int64_t pad_h, int64_t pad_w,
    int64_t dil_h, int64_t dil_w)
{
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "input and weight must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "expecting 4D input and weight");

    auto input_c = input.contiguous();
    auto weight_c = weight.contiguous();

    const int64_t N = input_c.size(0);
    const int64_t C_in = input_c.size(1);
    const int64_t H_in = input_c.size(2);
    const int64_t W_in = input_c.size(3);
    const int64_t C_out = weight_c.size(0);
    const int64_t KH = weight_c.size(2);
    const int64_t KW = weight_c.size(3);

    const int64_t H_out = (H_in + 2 * pad_h - dil_h * (KH - 1) - 1) / stride_h + 1;
    const int64_t W_out = (W_in + 2 * pad_w - dil_w * (KW - 1) - 1) / stride_w + 1;
    TORCH_CHECK(H_out > 0 && W_out > 0, "computed output size must be positive");
    TORCH_CHECK(N * C_out <= 65535, "batch * out_channels exceeds grid-z limit");

    auto output = at::empty({N, C_out, H_out, W_out}, input_c.options());

    torch::Tensor bias_c;
    const float* bias_ptr = nullptr;
    if (bias.defined() && bias.numel() > 0) {
        TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == at::kFloat, "bad bias tensor");
        bias_c = bias.contiguous();
        bias_ptr = bias_c.data_ptr<float>();
    }

    const int64_t tile_h = (TILE_OUT_Y - 1) * stride_h + (KH - 1) * dil_h + 1;
    const int64_t tile_w = (TILE_OUT_X - 1) * stride_w + (KW - 1) * dil_w + 1;
    const size_t smem_bytes = (size_t)(tile_h * tile_w + KH * KW) * sizeof(float);
    TORCH_CHECK(smem_bytes <= 48 * 1024, "shared memory requirement too large for this configuration");

    dim3 block(THREADS_X, THREADS_Y);
    dim3 grid(
        (unsigned)((W_out + TILE_OUT_X - 1) / TILE_OUT_X),
        (unsigned)((H_out + TILE_OUT_Y - 1) / TILE_OUT_Y),
        (unsigned)(N * C_out));

    auto stream = at::cuda::getCurrentCUDAStream();
    conv2d_direct_kernel<<<grid, block, smem_bytes, stream>>>(
        input_c.data_ptr<float>(),
        weight_c.data_ptr<float>(),
        bias_ptr,
        output.data_ptr<float>(),
        (int)N, (int)C_in, (int)H_in, (int)W_in,
        (int)C_out, (int)H_out, (int)W_out,
        (int)KH, (int)KW,
        (int)stride_h, (int)stride_w,
        (int)pad_h, (int)pad_w,
        (int)dil_h, (int)dil_w,
        (int)tile_h, (int)tile_w);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

conv2d_cpp_source = r"""
torch::Tensor conv2d_forward(
    torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
    int64_t stride_h, int64_t stride_w,
    int64_t pad_h, int64_t pad_w,
    int64_t dil_h, int64_t dil_w);
"""

# Compile the inline CUDA extension.
_conv2d_ext = load_inline(
    name="direct_conv2d_ext",
    cpp_sources=conv2d_cpp_source,
    cuda_sources=conv2d_cuda_source,
    functions=["conv2d_forward"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: the nn.Conv2d operator is replaced by a
    custom tiled direct-convolution CUDA kernel. The nn.Conv2d module is
    kept only as the parameter container; any configuration not handled by
    the custom kernel falls back to the reference 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)
        self.register_buffer("_null_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv2d
        if (c.groups == 1
                and x.is_cuda
                and x.dtype == torch.float32
                and isinstance(c.padding, tuple)):
            bias = c.bias if c.bias is not None else self._null_bias
            return _conv2d_ext.conv2d_forward(
                x, c.weight, bias,
                c.stride[0], c.stride[1],
                c.padding[0], c.padding[1],
                c.dilation[0], c.dilation[1])
        return c(x)
```