```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 implementation of the dilated 2D convolution used by Model.
#
# Kernel strategy (specialized for kernel=(5,9), dilation=(2,3), stride=1):
#   * Each block computes a 48 x 72 output tile for one (batch, out_channel).
#   * Each of the 384 threads computes a 3x3 block of outputs whose elements
#     are spaced by the dilation (dh=2 in height, dw=3 in width).  Because of
#     this spacing the 9 outputs of a thread share almost all of their input
#     taps: the union of the 9 shifted 5x9 stencils is only a 7 x 11 "dilated"
#     input patch (77 values), which is held entirely in registers.
#   * The (zero-padded) input tile of 56 x 96 values per input channel and the
#     kernel weights are staged in shared memory, two input channels at a time.
#   * Accumulation is FP32 (fmaf), the 5x9 tap loop is fully unrolled.
# ----------------------------------------------------------------------------

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

#define THREADS   384
#define TILE_TX   24
#define TILE_TY   16
#define KH        5
#define KW        9
#define DIL_H     2
#define DIL_W     3
#define TILE_H    (TILE_TY / 2 * 3 * DIL_H)      // 48 output rows / block
#define TILE_W    (TILE_TX / 3 * 3 * DIL_W)      // 72 output cols / block
#define IN_H      (TILE_H + (KH - 1) * DIL_H)    // 56
#define IN_W      (TILE_W + (KW - 1) * DIL_W)    // 96
#define IN_SIZE   (IN_H * IN_W)                  // 5376
#define IC_CHUNK  2
#define NKELEM    (KH * KW)                      // 45

__global__ __launch_bounds__(THREADS)
void conv2d_dilated_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    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 pad_h, const int pad_w,
    const bool has_bias) {

    extern __shared__ float smem[];
    float* s_in = smem;                          // IC_CHUNK * IN_SIZE floats
    float* s_w  = smem + IC_CHUNK * IN_SIZE;     // IC_CHUNK * NKELEM floats

    const int t  = threadIdx.x;
    const int ty = t / TILE_TX;                  // 0..15
    const int tx = t % TILE_TX;                  // 0..23

    // Dilation-aware 3x3 output assignment (covers 48x72 tile exactly once).
    const int r0 = 6 * (ty >> 1) + (ty & 1);     // 0..43, covers both parities
    const int c0 = 9 * (tx / 3) + (tx % 3);      // 0..65, covers mod-3 classes

    const int bz = blockIdx.z;
    const int n  = bz / OC;
    const int oc = bz % OC;
    const int oh0 = blockIdx.y * TILE_H;
    const int ow0 = blockIdx.x * TILE_W;

    const int oh_base = oh0 + r0;                // output row of (i = 0)
    const int ow_base = ow0 + c0;                // output col of (j = 0)

    const int in_origin_h = oh0 - pad_h;         // tile origin in input coords
    const int in_origin_w = ow0 - pad_w;

    float acc[3][3];
    const float bias_v = has_bias ? bias[oc] : 0.0f;
    #pragma unroll
    for (int i = 0; i < 3; ++i)
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            acc[i][j] = bias_v;

    const float* x_base = x + (long)n * IC * H * W;
    const float* w_base = w + (long)oc * IC * NKELEM;

    for (int ic0 = 0; ic0 < IC; ic0 += IC_CHUNK) {
        const int chunk = min(IC_CHUNK, IC - ic0);

        // ---- Stage input tile(s), zero-padded, into shared memory ---------
        for (int idx = t; idx < chunk * IN_SIZE; idx += THREADS) {
            const int ci  = idx / IN_SIZE;
            const int rem = idx - ci * IN_SIZE;
            const int sr  = rem / IN_W;
            const int sc  = rem - sr * IN_W;
            const int ih  = in_origin_h + sr;
            const int iw  = in_origin_w + sc;
            float v = 0.0f;
            if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
                v = __ldg(x_base + ((long)(ic0 + ci) * H + ih) * W + iw);
            }
            s_in[idx] = v;
        }

        // ---- Stage weights into shared memory -----------------------------
        const float* w_src = w_base + (long)ic0 * NKELEM;
        for (int idx = t; idx < chunk * NKELEM; idx += THREADS) {
            s_w[idx] = __ldg(w_src + idx);
        }
        __syncthreads();

        // ---- Compute ------------------------------------------------------
        for (int ci = 0; ci < chunk; ++ci) {
            const float* s_in_ci = s_in + ci * IN_SIZE;
            const float* s_w_ci  = s_w + ci * NKELEM;

            // Load the 7x11 dilated register patch.  All patch reads are
            // inside the (zero-padded) shared tile: max row 43+12 = 55 < 56,
            // max col 65+30 = 95 < 96.
            float seg[7][11];
            #pragma unroll
            for (int m = 0; m < 7; ++m) {
                const float* srow = s_in_ci + (r0 + m * DIL_H) * IN_W + c0;
                #pragma unroll
                for (int q = 0; q < 11; ++q) {
                    seg[m][q] = srow[q * DIL_W];
                }
            }

            #pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
                #pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    const float wv = s_w_ci[kh * KW + kw];
                    #pragma unroll
                    for (int i = 0; i < 3; ++i) {
                        #pragma unroll
                        for (int j = 0; j < 3; ++j) {
                            acc[i][j] = fmaf(seg[i + kh][j + kw], wv, acc[i][j]);
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    // ---- Write outputs ----------------------------------------------------
    #pragma unroll
    for (int i = 0; i < 3; ++i) {
        const int oh = oh_base + i * DIL_H;
        if (oh < HO) {
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                const int ow = ow_base + j * DIL_W;
                if (ow < WO) {
                    out[((long)(n * OC + oc) * HO + oh) * WO + ow] = acc[i][j];
                }
            }
        }
    }
}

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) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x.dim() == 4 && w.dim() == 4, "x must be NCHW, w must be OIHW");
    TORCH_CHECK(stride[0] == 1 && stride[1] == 1, "kernel requires stride 1");
    TORCH_CHECK(dilation[0] == DIL_H && dilation[1] == DIL_W,
                "kernel specialized for dilation (2,3)");
    TORCH_CHECK(w.size(2) == KH && w.size(3) == KW,
                "kernel specialized for kernel size (5,9)");

    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 - DIL_H * (KH - 1) - 1) + 1;  // stride = 1
    const int WO = (W + 2 * pad_w - DIL_W * (KW - 1) - 1) + 1;
    TORCH_CHECK(HO > 0 && WO > 0, "invalid output size");

    auto out = torch::empty({N, OC, HO, WO}, x_c.options());

    dim3 grid((WO + TILE_W - 1) / TILE_W,
              (HO + TILE_H - 1) / TILE_H,
              N * OC);

    const size_t smem =
        (size_t)(IC_CHUNK * IN_SIZE + IC_CHUNK * NKELEM) * sizeof(float);

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv2d_dilated_kernel<<<grid, THREADS, smem, stream>>>(
        x_c.data_ptr<float>(), w_c.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(),
        N, IC, H, W, OC, HO, WO, pad_h, pad_w, has_bias);

    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_tiled_v2",
    cpp_sources=[conv2d_cpp_source],
    cuda_sources=[conv2d_cuda_source],
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
)


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


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 tiled CUDA kernel that exploits
    the dilation structure (register-blocked 3x3 outputs spaced by the
    dilation factors, zero-padded shared-memory input tiles, fully unrolled
    5x9 tap loop with FP32 FMA accumulation).  Falls back to the standard
    cuDNN implementation for any configuration the kernel is not
    specialized for.
    """

    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)

        k = _pair(kernel_size)
        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]]
        self._use_custom = (k == (5, 9) and s == (1, 1) and d == (2, 3))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (self._use_custom and 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)
```