```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# ----------------------------------------------------------------------------
# Custom near-FP32-accurate tensor-core CUDA kernel for the dilated 2D conv
# (specialized: kernel=(5,9), stride=1, dilation=(2,3), any padding/channels).
#
# Method: implicit-GEMM style direct convolution on fp16 tensor cores using
# an error-compensated hi/lo split of both operands:
#     x = x_hi + x_lo      (each fp16; combined representation error ~2^-22)
#     a.b ~= a_hi*b_hi + a_lo*b_hi + a_hi*b_lo     (3 mmas, FP32 accumulate)
# This keeps accuracy in the same class as a plain FP32 kernel (error ~1e-7,
# well below TF32's ~1e-4), while executing on the fp16 tensor cores which
# have 16x the throughput of the FP32 CUDA cores on A100.
#
# Layout:
#   * Block = 256 threads (8 warps); tile = 8 output rows x 16 output cols,
#     covering all (padded to 64) output channels of one image.
#   * Each warp owns one m16 tile (16 consecutive output columns of one
#     output row), so dilation shifts become compile-time offsets per tap.
#   * The zero-padded input tile (16x40 per channel) is staged in shared
#     memory in chunks of 16 channels, pre-split and packed as half2 channel
#     pairs (hi and lo planes), with a padded plane stride so the A-fragment
#     gathers are bank-conflict free.
#   * Weights are transposed/split once per forward by a tiny pre-kernel
#     into [tap][ic/2][oc_pad] half2 (hi/lo) arrays; B fragments are loaded
#     straight from global memory (L1/L2 resident, reused by every warp).
#   * Uses mma.m16n8k8 (available sm_75+); anything else falls back to
#     at::conv2d (cuDNN).
# ----------------------------------------------------------------------------

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

#define KH        5
#define KW        9
#define NTAP      (KH * KW)                 // 45
#define DH        2
#define DW        3
#define TILE_OH   8
#define TILE_OW   16
#define IN_H      (TILE_OH + (KH - 1) * DH) // 16
#define IN_W      (TILE_OW + (KW - 1) * DW) // 40
#define PIX       (IN_H * IN_W)             // 640
#define PLANE     648                       // padded: PLANE % 32 == 8
#define CHUNK     16                        // input channels staged at once
#define CPAIRS    (CHUNK / 2)               // 8 half2 planes (hi and lo)
#define NTHREADS  256

#define MMA1688(C, A0, A1, B)                                              \
    asm volatile(                                                          \
        "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "                \
        "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};\n"                   \
        : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3])                   \
        : "r"(A0), "r"(A1), "r"(B))

// ---------------------------------------------------------------------------
// Pre-kernel: transpose + split weights [OC][IC][45] ->
// hi/lo u32 (half2 channel pairs) arrays [45][ICP][OCP].
// ---------------------------------------------------------------------------
__global__ void wsplit_kernel(const float* __restrict__ w,
                              unsigned* __restrict__ w_hi,
                              unsigned* __restrict__ w_lo,
                              const int IC, const int OC,
                              const int ICP, const int OCP) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = NTAP * ICP * OCP;
    if (idx >= total) return;
    const int oc  = idx % OCP;
    const int t   = idx / OCP;
    const int p   = t % ICP;
    const int tap = t / ICP;
    const int ic0 = 2 * p;
    const int ic1 = ic0 + 1;

    float v0 = 0.0f, v1 = 0.0f;
    if (oc < OC) {
        if (ic0 < IC) v0 = __ldg(w + ((long)oc * IC + ic0) * NTAP + tap);
        if (ic1 < IC) v1 = __ldg(w + ((long)oc * IC + ic1) * NTAP + tap);
    }
    const __half h0 = __float2half_rn(v0);
    const __half h1 = __float2half_rn(v1);
    const __half l0 = __float2half_rn(v0 - __half2float(h0));
    const __half l1 = __float2half_rn(v1 - __half2float(h1));
    __half2 hh = __halves2half2(h0, h1);
    __half2 ll = __halves2half2(l0, l1);
    w_hi[idx] = *reinterpret_cast<unsigned*>(&hh);
    w_lo[idx] = *reinterpret_cast<unsigned*>(&ll);
}

// ---------------------------------------------------------------------------
// Main kernel.
// grid: (ceil(WO/16), ceil(HO/8), N*ceil(OC/64));  block: 256 threads.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NTHREADS)
conv2d_f16x3_kernel(const float* __restrict__ x,
                    const unsigned* __restrict__ w_hi,  // [45][ICP][OCP]
                    const unsigned* __restrict__ w_lo,
                    const float* __restrict__ bias,     // may be null
                    float* __restrict__ out,
                    const int N, const int IC, const int H, const int W,
                    const int OC, const int HO, const int WO,
                    const int ICP, const int OCP,
                    const int pad_h, const int pad_w) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)
    __shared__ unsigned s_hi[CPAIRS * PLANE];
    __shared__ unsigned s_lo[CPAIRS * PLANE];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int lm0  = lane >> 2;   // 0..7  (pixel offset within m16 tile)
    const int lk0  = lane & 3;    // 0..3  (channel-pair index within group)

    const int oc_blocks = (OC + 63) >> 6;
    const int bz  = blockIdx.z;
    const int n   = bz / oc_blocks;
    const int oc0 = (bz % oc_blocks) << 6;
    const int oh0 = blockIdx.y * TILE_OH;
    const int ow0 = blockIdx.x * TILE_OW;

    const int r  = warp;          // this warp's output row within the tile
    const int oh = oh0 + r;

    float c[8][4];
#pragma unroll
    for (int j = 0; j < 8; ++j) {
        c[j][0] = 0.0f; c[j][1] = 0.0f; c[j][2] = 0.0f; c[j][3] = 0.0f;
    }

    const float* x_base = x + (long)n * IC * H * W;

    for (int ic0 = 0; ic0 < IC; ic0 += CHUNK) {
        // ---- Stage zero-padded input tile, split into hi/lo half2 pairs ---
        {
            const int in_r0 = oh0 - pad_h;
            const int in_c0 = ow0 - pad_w;
#pragma unroll 4
            for (int idx = tid; idx < CPAIRS * PIX; idx += NTHREADS) {
                const int p   = idx / PIX;
                const int rem = idx - p * PIX;
                const int sr  = rem / IN_W;
                const int sc  = rem - sr * IN_W;
                const int ih  = in_r0 + sr;
                const int iw  = in_c0 + sc;
                float v0 = 0.0f, v1 = 0.0f;
                const int icA = ic0 + 2 * p;
                if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W
                    && icA < IC) {
                    const long off = ((long)icA * H + ih) * W + iw;
                    v0 = __ldg(x_base + off);
                    if (icA + 1 < IC) v1 = __ldg(x_base + (long)H * W + off);
                }
                const __half h0 = __float2half_rn(v0);
                const __half h1 = __float2half_rn(v1);
                const __half l0 = __float2half_rn(v0 - __half2float(h0));
                const __half l1 = __float2half_rn(v1 - __half2float(h1));
                __half2 hh = __halves2half2(h0, h1);
                __half2 ll = __halves2half2(l0, l1);
                const int so = p * PLANE + sr * IN_W + sc;
                s_hi[so] = *reinterpret_cast<unsigned*>(&hh);
                s_lo[so] = *reinterpret_cast<unsigned*>(&ll);
            }
        }
        __syncthreads();

        // ---- Accumulate over the 45 taps ----------------------------------
        const int icp0 = ic0 >> 1;
#pragma unroll 1
        for (int tap = 0; tap < NTAP; ++tap) {
            const int kg = tap / KW;
            const int kw = tap - kg * KW;
            const int srow = (r + DH * kg) * IN_W + DW * kw;

            // A fragments (hi/lo), two k-groups of 8 channels each.
            const unsigned* sh = s_hi + srow + lm0;
            const unsigned* sl = s_lo + srow + lm0;
            unsigned ah00 = sh[lk0 * PLANE];
            unsigned ah01 = sh[lk0 * PLANE + 8];
            unsigned ah10 = sh[(lk0 + 4) * PLANE];
            unsigned ah11 = sh[(lk0 + 4) * PLANE + 8];
            unsigned al00 = sl[lk0 * PLANE];
            unsigned al01 = sl[lk0 * PLANE + 8];
            unsigned al10 = sl[(lk0 + 4) * PLANE];
            unsigned al11 = sl[(lk0 + 4) * PLANE + 8];

            // B fragments (hi/lo): 8 n-tiles of 8 output channels each.
            const unsigned* wp_h =
                w_hi + ((long)tap * ICP + icp0 + lk0) * OCP + oc0 + lm0;
            const unsigned* wp_l =
                w_lo + ((long)tap * ICP + icp0 + lk0) * OCP + oc0 + lm0;
            const int step_k = 4 * OCP;   // +4 channel pairs = +8 channels

#pragma unroll
            for (int j = 0; j < 8; ++j) {
                const unsigned bh0 = __ldg(wp_h + 8 * j);
                const unsigned bh1 = __ldg(wp_h + step_k + 8 * j);
                const unsigned bl0 = __ldg(wp_l + 8 * j);
                const unsigned bl1 = __ldg(wp_l + step_k + 8 * j);
                MMA1688(c[j], ah00, ah01, bh0);
                MMA1688(c[j], al00, al01, bh0);
                MMA1688(c[j], ah00, ah01, bl0);
                MMA1688(c[j], ah10, ah11, bh1);
                MMA1688(c[j], al10, al11, bh1);
                MMA1688(c[j], ah10, ah11, bl1);
            }
        }
        __syncthreads();
    }

    // ---- Epilogue: bias + store -------------------------------------------
    if (oh < HO) {
#pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int oc = oc0 + j * 8 + 2 * lk0;
            float c0 = c[j][0], c1 = c[j][1], c2 = c[j][2], c3 = c[j][3];
            if (bias != nullptr) {
                if (oc < OC) {
                    const float b0 = __ldg(bias + oc);
                    c0 += b0; c2 += b0;
                }
                if (oc + 1 < OC) {
                    const float b1 = __ldg(bias + oc + 1);
                    c1 += b1; c3 += b1;
                }
            }
            const int ow_a = ow0 + lm0;
            const int ow_b = ow0 + lm0 + 8;
            const bool wa = (ow_a < WO);
            const bool wb = (ow_b < WO);
            float* p0 = out + ((long)(n * OC + oc) * HO + oh) * WO;
            if (oc < OC) {
                if (wa) p0[ow_a] = c0;
                if (wb) p0[ow_b] = c2;
            }
            if (oc + 1 < OC) {
                float* p1 = p0 + (long)HO * WO;
                if (wa) p1[ow_a] = c1;
                if (wb) p1[ow_b] = c3;
            }
        }
    }
#endif
}

// ---------------------------------------------------------------------------
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias_opt,
                             std::vector<int64_t> stride,
                             std::vector<int64_t> padding,
                             std::vector<int64_t> dilation) {
    const bool has_bias =
        bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0;

    bool ok = x.is_cuda() && w.is_cuda()
        && x.scalar_type() == torch::kFloat32
        && x.dim() == 4 && w.dim() == 4
        && stride[0] == 1 && stride[1] == 1
        && dilation[0] == DH && dilation[1] == DW
        && w.size(2) == KH && w.size(3) == KW;
    if (ok) {
        auto prop = at::cuda::getCurrentDeviceProperties();
        ok = (prop->major > 7) || (prop->major == 7 && prop->minor >= 5);
    }
    if (!ok) {
        return at::conv2d(x, w, bias_opt, stride, padding, dilation, 1);
    }

    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(x));

    auto x_c = x.contiguous();
    auto w_c = w.contiguous();

    const int N  = x_c.size(0);
    const int IC = x_c.size(1);
    const int H  = x_c.size(2);
    const int W  = x_c.size(3);
    const int OC = w_c.size(0);

    const int pad_h = (int)padding[0];
    const int pad_w = (int)padding[1];

    const int HO = H + 2 * pad_h - DH * (KH - 1);   // stride 1
    const int WO = W + 2 * pad_w - DW * (KW - 1);
    if (HO <= 0 || WO <= 0) {
        return at::conv2d(x, w, bias_opt, stride, padding, dilation, 1);
    }

    const int ICP = (IC + 1) / 2;
    const int OCP = ((OC + 63) / 64) * 64;

    auto out   = torch::empty({N, OC, HO, WO}, x_c.options());
    auto uopts = w_c.options().dtype(torch::kInt32);
    auto w_hi  = torch::empty({(long)NTAP * ICP * OCP}, uopts);
    auto w_lo  = torch::empty({(long)NTAP * ICP * OCP}, uopts);

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

    {
        const int total = NTAP * ICP * OCP;
        const int threads = 256;
        const int blocks = (total + threads - 1) / threads;
        wsplit_kernel<<<blocks, threads, 0, stream>>>(
            w_c.data_ptr<float>(),
            reinterpret_cast<unsigned*>(w_hi.data_ptr<int>()),
            reinterpret_cast<unsigned*>(w_lo.data_ptr<int>()),
            IC, OC, ICP, OCP);
    }

    dim3 grid((WO + TILE_OW - 1) / TILE_OW,
              (HO + TILE_OH - 1) / TILE_OH,
              N * ((OC + 63) / 64));

    const float* bias_ptr = nullptr;
    torch::Tensor bias_c;
    if (has_bias) {
        bias_c = bias_opt->contiguous();
        bias_ptr = bias_c.data_ptr<float>();
    }

    conv2d_f16x3_kernel<<<grid, NTHREADS, 0, stream>>>(
        x_c.data_ptr<float>(),
        reinterpret_cast<const unsigned*>(w_hi.data_ptr<int>()),
        reinterpret_cast<const unsigned*>(w_lo.data_ptr<int>()),
        bias_ptr, out.data_ptr<float>(),
        N, IC, H, W, OC, HO, WO, ICP, OCP, pad_h, pad_w);

    return out;
}
"""

conv2d_cpp_source = r"""
#include <torch/extension.h>
#include <vector>
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias,
                             std::vector<int64_t> stride,
                             std::vector<int64_t> padding,
                             std::vector<int64_t> dilation);
"""

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


class ModelNew(nn.Module):
    """
    Optimized version of Model.  Keeps an nn.Conv2d submodule so that
    initialization and state_dict are identical to the reference model, but
    computes the forward pass with a custom tensor-core CUDA kernel that
    uses error-compensated fp16 hi/lo splitting (3x mma.m16n8k8 with FP32
    accumulation) to reach far above FP32 CUDA-core throughput while keeping
    FP32-class accuracy.  Falls back to the standard implementation for any
    configuration or device the kernel does not support.
    """

    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)

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

        s = _pair(stride)
        p = _pair(padding)
        d = _pair(dilation)
        self._stride = [s[0], s[1]]
        self._padding = [p[0], p[1]]
        self._dilation = [d[0], d[1]]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            return conv2d_ext.conv2d_forward(
                x, self.conv2d.weight, self.conv2d.bias,
                self._stride, self._padding, self._dilation)
        return self.conv2d(x)
```