```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 kernel: direct tiled 2D convolution specialized for
# kernel_size=(5,9), stride=1, dilation=(2,3) (any padding / channels / batch).
#
# Design:
#  - Each thread block computes a 32x32 spatial tile of the output for
#    8 output channels of one image.
#  - The dilated receptive field of a 32x32 output tile is a 40x56 input
#    window; 4 input channels of this window are staged into shared memory
#    at a time (zero-filled outside the image to implement padding).
#  - Weights for 8 output channels are staged into shared memory as float4
#    (broadcast reads).
#  - Each thread accumulates 8 output channels x 2x2 spatial positions = 16
#    independent FMA chains, giving high arithmetic intensity and ILP.
conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TILE_Y 32
#define TILE_X 32
#define KH 5
#define KW 9
#define DIL_H 2
#define DIL_W 3
#define IN_H (TILE_Y + (KH - 1) * DIL_H)   // 40
#define IN_W (TILE_X + (KW - 1) * DIL_W)   // 56
#define IC_CHUNK 4
#define OC_PER_BLOCK 8
#define S_IN_FLOATS (IC_CHUNK * IN_H * IN_W)          // 8960
#define S_W_FLOAT4S (IC_CHUNK * KH * KW)              // 180

__global__ __launch_bounds__(256, 2)
void conv2d_5x9_dil23_kernel(const float* __restrict__ x,
                             const float* __restrict__ w,
                             const float* __restrict__ bias,
                             float* __restrict__ out,
                             int N, int C_in, int H, int W,
                             int C_out, int H_out, int W_out,
                             int pad_h, int pad_w,
                             int oc_blocks) {
    extern __shared__ float smem[];
    float*  s_in = smem;                                          // 8960 floats
    float4* s_w0 = reinterpret_cast<float4*>(smem + S_IN_FLOATS); // oc 0..3
    float4* s_w1 = reinterpret_cast<float4*>(smem + S_IN_FLOATS + S_W_FLOAT4S * 4); // oc 4..7

    const int tile_x = blockIdx.x;
    const int tile_y = blockIdx.y;
    const int n      = blockIdx.z / oc_blocks;
    const int ocb    = blockIdx.z % oc_blocks;
    const int oc0    = ocb * OC_PER_BLOCK;

    const int out_y0 = tile_y * TILE_Y;
    const int out_x0 = tile_x * TILE_X;

    const int tx  = threadIdx.x;  // 0..15
    const int ty  = threadIdx.y;  // 0..15
    const int tid = ty * 16 + tx;

    // Accumulators: [oc][dy][dx]
    // rows: out_y0 + ty*2 + dy ; cols: out_x0 + tx + dx*16
    float acc[OC_PER_BLOCK][2][2];
    #pragma unroll
    for (int j = 0; j < OC_PER_BLOCK; ++j)
        #pragma unroll
        for (int dy = 0; dy < 2; ++dy)
            #pragma unroll
            for (int dx = 0; dx < 2; ++dx)
                acc[j][dy][dx] = 0.0f;

    const float* x_n = x + (long)n * C_in * H * W;

    for (int ic0 = 0; ic0 < C_in; ic0 += IC_CHUNK) {
        // ---- stage input tile: 4 channels x 40 x 56 (zero padded) ----
        #pragma unroll 4
        for (int idx = tid; idx < S_IN_FLOATS; idx += 256) {
            int c   = idx / (IN_H * IN_W);
            int rem = idx - c * (IN_H * IN_W);
            int r   = rem / IN_W;
            int cc  = rem - r * IN_W;
            int gy  = out_y0 + r  - pad_h;
            int gx  = out_x0 + cc - pad_w;
            int ic  = ic0 + c;
            float v = 0.0f;
            if (ic < C_in && gy >= 0 && gy < H && gx >= 0 && gx < W) {
                v = x_n[((long)ic * H + gy) * W + gx];
            }
            s_in[idx] = v;
        }
        // ---- stage weights: 8 oc x 4 ic x 5 x 9 as float4 pairs ----
        if (tid < S_W_FLOAT4S) {
            int icc = tid / (KH * KW);
            int k   = tid - icc * (KH * KW);
            int kh  = k / KW;
            int kw  = k - kh * KW;
            int ic  = ic0 + icc;
            bool icok = (ic < C_in);
            long wbase = (long)ic * KH * KW + kh * KW + kw;
            long woc   = (long)C_in * KH * KW;
            float4 lo, hi;
            lo.x = (icok && (oc0 + 0) < C_out) ? w[wbase + woc * (oc0 + 0)] : 0.0f;
            lo.y = (icok && (oc0 + 1) < C_out) ? w[wbase + woc * (oc0 + 1)] : 0.0f;
            lo.z = (icok && (oc0 + 2) < C_out) ? w[wbase + woc * (oc0 + 2)] : 0.0f;
            lo.w = (icok && (oc0 + 3) < C_out) ? w[wbase + woc * (oc0 + 3)] : 0.0f;
            hi.x = (icok && (oc0 + 4) < C_out) ? w[wbase + woc * (oc0 + 4)] : 0.0f;
            hi.y = (icok && (oc0 + 5) < C_out) ? w[wbase + woc * (oc0 + 5)] : 0.0f;
            hi.z = (icok && (oc0 + 6) < C_out) ? w[wbase + woc * (oc0 + 6)] : 0.0f;
            hi.w = (icok && (oc0 + 7) < C_out) ? w[wbase + woc * (oc0 + 7)] : 0.0f;
            s_w0[tid] = lo;
            s_w1[tid] = hi;
        }
        __syncthreads();

        // ---- compute ----
        #pragma unroll
        for (int icc = 0; icc < IC_CHUNK; ++icc) {
            const int base = icc * (IN_H * IN_W) + ty * (2 * IN_W) + tx;
            const float4* w0p = s_w0 + icc * (KH * KW);
            const float4* w1p = s_w1 + icc * (KH * KW);
            #pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
                #pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    const int a = base + kh * (DIL_H * IN_W) + kw * DIL_W;
                    float iv[2][2];
                    iv[0][0] = s_in[a];
                    iv[0][1] = s_in[a + 16];
                    iv[1][0] = s_in[a + IN_W];
                    iv[1][1] = s_in[a + IN_W + 16];
                    float4 wlo = w0p[kh * KW + kw];
                    float4 whi = w1p[kh * KW + kw];
                    float wv[8] = {wlo.x, wlo.y, wlo.z, wlo.w,
                                   whi.x, whi.y, whi.z, whi.w};
                    #pragma unroll
                    for (int j = 0; j < OC_PER_BLOCK; ++j)
                        #pragma unroll
                        for (int dy = 0; dy < 2; ++dy)
                            #pragma unroll
                            for (int dx = 0; dx < 2; ++dx)
                                acc[j][dy][dx] = fmaf(iv[dy][dx], wv[j], acc[j][dy][dx]);
                }
            }
        }
        __syncthreads();
    }

    // ---- epilogue ----
    #pragma unroll
    for (int j = 0; j < OC_PER_BLOCK; ++j) {
        int oc = oc0 + j;
        if (oc < C_out) {
            float bv = (bias != nullptr) ? bias[oc] : 0.0f;
            #pragma unroll
            for (int dy = 0; dy < 2; ++dy) {
                int row = out_y0 + ty * 2 + dy;
                if (row < H_out) {
                    #pragma unroll
                    for (int dx = 0; dx < 2; ++dx) {
                        int col = out_x0 + tx + dx * 16;
                        if (col < W_out) {
                            out[(((long)n * C_out + oc) * H_out + row) * W_out + col] =
                                acc[j][dy][dx] + bv;
                        }
                    }
                }
            }
        }
    }
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                             int64_t pad_h, int64_t pad_w) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    const int N     = x.size(0);
    const int C_in  = x.size(1);
    const int H     = x.size(2);
    const int W     = x.size(3);
    const int C_out = w.size(0);
    const int H_out = H + 2 * (int)pad_h - DIL_H * (KH - 1);
    const int W_out = W + 2 * (int)pad_w - DIL_W * (KW - 1);

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

    const int tiles_x   = (W_out + TILE_X - 1) / TILE_X;
    const int tiles_y   = (H_out + TILE_Y - 1) / TILE_Y;
    const int oc_blocks = (C_out + OC_PER_BLOCK - 1) / OC_PER_BLOCK;

    dim3 grid(tiles_x, tiles_y, N * oc_blocks);
    dim3 block(16, 16);
    const size_t shmem = (size_t)(S_IN_FLOATS + 2 * S_W_FLOAT4S * 4) * sizeof(float);

    const float* bias_ptr = (b.numel() > 0) ? b.data_ptr<float>() : nullptr;
    auto stream = at::cuda::getCurrentCUDAStream();

    conv2d_5x9_dil23_kernel<<<grid, block, shmem, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(),
        N, C_in, H, W, C_out, H_out, W_out,
        (int)pad_h, (int)pad_w, oc_blocks);
    return out;
}
"""

conv2d_cpp_source = (
    "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,"
    "                             int64_t pad_h, int64_t pad_w);"
)

conv2d_ext = load_inline(
    name="conv2d_5x9_dil23_ext",
    cpp_sources=conv2d_cpp_source,
    cuda_sources=conv2d_cuda_source,
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: the 2D convolution (kernel 5x9, stride 1,
    dilation (2,3)) is performed by a custom tiled CUDA kernel using shared
    memory staging and register-blocked accumulation. Falls back to the
    standard nn.Conv2d for any unsupported configuration.
    """

    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 = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
        sh, sw = (stride, stride) if isinstance(stride, int) else stride
        ph, pw = (padding, padding) if isinstance(padding, int) else padding
        dh, dw = (dilation, dilation) if isinstance(dilation, int) else dilation

        self._pad_h = int(ph)
        self._pad_w = int(pw)
        self._use_fast = (int(kh) == 5 and int(kw) == 9 and
                          int(sh) == 1 and int(sw) == 1 and
                          int(dh) == 2 and int(dw) == 3)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._use_fast and x.is_cuda and x.dtype == torch.float32:
            weight = self.conv2d.weight
            bias = self.conv2d.bias
            if bias is None:
                bias = torch.empty(0, device=x.device, dtype=torch.float32)
            return conv2d_ext.conv2d_forward(
                x.contiguous(), weight.contiguous(), bias.contiguous(),
                self._pad_h, self._pad_w)
        return self.conv2d(x)
```