```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernels for depthwise 2D convolution.
# Fast path: shared-memory tiled kernel (compile-time unrolled) for
# stride=1, padding=0, dilation=1 with common kernel sizes.
# Fallback: fully general depthwise kernel for any stride/pad/dilation.
dwconv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define DW_TILE_W 32
#define DW_TILE_H 8

// ---------------------------------------------------------------------------
// Tiled kernel: stride=1, pad=0, dilation=1.
// Each block computes a DW_TILE_H x DW_TILE_W tile of outputs for one (n, c).
// The (DW_TILE_H+KH-1) x (DW_TILE_W+KW-1) input halo tile is staged in shared
// memory so each input element is fetched from global memory only ~once.
// ---------------------------------------------------------------------------
template <int KH, int KW>
__global__ void dwconv2d_tile_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int C, int H, int W, int outH, int outW,
    int tilesX, int tilesPerSlice)
{
    constexpr int SH_H = DW_TILE_H + KH - 1;
    constexpr int SH_W = DW_TILE_W + KW - 1;

    __shared__ float sh[SH_H * SH_W];
    __shared__ float shw[KH * KW];

    const int tile = blockIdx.x % tilesPerSlice;
    const int nc   = blockIdx.x / tilesPerSlice;
    const int ty   = tile / tilesX;
    const int tx   = tile % tilesX;
    const int n    = nc / C;
    const int c    = nc % C;

    const int ow0 = tx * DW_TILE_W;
    const int oh0 = ty * DW_TILE_H;

    const int tid = threadIdx.y * DW_TILE_W + threadIdx.x;
    if (tid < KH * KW) {
        shw[tid] = weight[(long)c * (KH * KW) + tid];
    }

    const float* in_base = input + ((long)n * C + c) * H * W;

    #pragma unroll
    for (int sy = threadIdx.y; sy < SH_H; sy += DW_TILE_H) {
        const int gy = oh0 + sy;
        const bool rok = (gy < H);
        #pragma unroll
        for (int sx = threadIdx.x; sx < SH_W; sx += DW_TILE_W) {
            const int gx = ow0 + sx;
            sh[sy * SH_W + sx] = (rok && gx < W) ? in_base[(long)gy * W + gx] : 0.0f;
        }
    }
    __syncthreads();

    const int ow = ow0 + threadIdx.x;
    const int oh = oh0 + threadIdx.y;

    if (ow < outW && oh < outH) {
        float sum = (bias != nullptr) ? bias[c] : 0.0f;
        const int sx = threadIdx.x;
        const int sy = threadIdx.y;
        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                sum = fmaf(sh[(sy + kh) * SH_W + (sx + kw)], shw[kh * KW + kw], sum);
            }
        }
        output[((long)n * C + c) * outH * outW + (long)oh * outW + ow] = sum;
    }
}

// ---------------------------------------------------------------------------
// General depthwise kernel: any kernel size / stride / padding / dilation.
// One thread per output element; threads are laid out so that consecutive
// threads map to consecutive output widths for coalesced global access.
// ---------------------------------------------------------------------------
__global__ void dwconv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int C, int H, int W, int outH, int outW,
    int KH, int KW, int strideH, int strideW, int padH, int padW, int dilH, int dilW)
{
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long total = (long)N * C * outH * outW;
    if (idx >= total) return;

    const int ow = (int)(idx % outW);
    const int oh = (int)((idx / outW) % outH);
    const int c  = (int)((idx / ((long)outW * outH)) % C);
    const int n  = (int)(idx / ((long)outW * outH * C));

    const float* in_base = input + ((long)n * C + c) * H * W;
    const float* w_base  = weight + (long)c * KH * KW;

    float sum = (bias != nullptr) ? bias[c] : 0.0f;
    for (int kh = 0; kh < KH; ++kh) {
        const int ih = oh * strideH - padH + kh * dilH;
        if (ih < 0 || ih >= H) continue;
        const float* in_row = in_base + (long)ih * W;
        const float* w_row  = w_base + kh * KW;
        for (int kw = 0; kw < KW; ++kw) {
            const int iw = ow * strideW - padW + kw * dilW;
            if (iw < 0 || iw >= W) continue;
            sum = fmaf(in_row[iw], w_row[kw], sum);
        }
    }
    output[idx] = sum;
}

template <int KH, int KW>
static void launch_tile(
    const float* input, const float* weight, const float* bias, float* output,
    int N, int C, int H, int W, int outH, int outW, cudaStream_t stream)
{
    const int tilesX = (outW + DW_TILE_W - 1) / DW_TILE_W;
    const int tilesY = (outH + DW_TILE_H - 1) / DW_TILE_H;
    const int tilesPerSlice = tilesX * tilesY;
    const long blocks = (long)tilesPerSlice * N * C;
    dim3 grid((unsigned int)blocks);
    dim3 block(DW_TILE_W, DW_TILE_H);
    dwconv2d_tile_kernel<KH, KW><<<grid, block, 0, stream>>>(
        input, weight, bias, output, N, C, H, W, outH, outW, tilesX, tilesPerSlice);
}

torch::Tensor dwconv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t strideH, int64_t strideW,
    int64_t padH, int64_t padW,
    int64_t dilH, int64_t dilW)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");

    auto x = input.contiguous();
    auto w = weight.contiguous();

    const int N  = (int)x.size(0);
    const int C  = (int)x.size(1);
    const int H  = (int)x.size(2);
    const int W  = (int)x.size(3);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);

    const int outH = (int)((H + 2 * padH - dilH * (KH - 1) - 1) / strideH + 1);
    const int outW = (int)((W + 2 * padW - dilW * (KW - 1) - 1) / strideW + 1);

    auto out = torch::empty({N, C, outH, outW}, x.options());
    if (out.numel() == 0) return out;

    const float* in_ptr = x.data_ptr<float>();
    const float* w_ptr  = w.data_ptr<float>();
    float* out_ptr      = out.data_ptr<float>();

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

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

    const bool fast = (strideH == 1 && strideW == 1 && padH == 0 && padW == 0 &&
                       dilH == 1 && dilW == 1 && H >= KH && W >= KW);

    if (fast) {
        if (KH == 3 && KW == 7) { launch_tile<3, 7>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
        if (KH == 3 && KW == 3) { launch_tile<3, 3>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
        if (KH == 5 && KW == 5) { launch_tile<5, 5>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
        if (KH == 7 && KW == 7) { launch_tile<7, 7>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
        if (KH == 1 && KW == 1) { launch_tile<1, 1>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
        if (KH == 3 && KW == 1) { launch_tile<3, 1>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
        if (KH == 1 && KW == 3) { launch_tile<1, 3>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, stream); return out; }
    }

    const long total = (long)N * C * outH * outW;
    const int threads = 256;
    const long blocks = (total + threads - 1) / threads;
    dwconv2d_generic_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        in_ptr, w_ptr, b_ptr, out_ptr,
        N, C, H, W, outH, outW,
        KH, KW, (int)strideH, (int)strideW, (int)padH, (int)padW, (int)dilH, (int)dilW);
    return out;
}
"""

dwconv2d_cpp_source = (
    "torch::Tensor dwconv2d_forward(torch::Tensor input, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias, int64_t strideH, int64_t strideW, "
    "int64_t padH, int64_t padW, int64_t dilH, int64_t dilW);"
)

# Compile the inline CUDA extension
dwconv2d_ext = load_inline(
    name="dwconv2d_ext",
    cpp_sources=dwconv2d_cpp_source,
    cuda_sources=dwconv2d_cuda_source,
    functions=["dwconv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution backed by a custom CUDA kernel.

    The nn.Conv2d submodule is kept so that parameter initialization and
    state_dict layout exactly match the reference Model.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size_h: int, kernel_size_w: int, stride_h: int = 1, stride_w: int = 1, padding_h: int = 0, padding_w: int = 0, dilation_h: int = 1, dilation_w: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, (kernel_size_h, kernel_size_w), stride=(stride_h, stride_w), padding=(padding_h, padding_w), dilation=(dilation_h, dilation_w), groups=in_channels, bias=bias)
        self.stride_h = stride_h
        self.stride_w = stride_w
        self.padding_h = padding_h
        self.padding_w = padding_w
        self.dilation_h = dilation_h
        self.dilation_w = dilation_w

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.conv2d.weight
        bias = self.conv2d.bias
        if x.is_cuda and weight.is_cuda and x.dtype == torch.float32 and weight.dtype == torch.float32:
            return dwconv2d_ext.dwconv2d_forward(
                x, weight, bias,
                self.stride_h, self.stride_w,
                self.padding_h, self.padding_w,
                self.dilation_h, self.dilation_w)
        return F.conv2d(
            x, weight, bias,
            stride=(self.stride_h, self.stride_w),
            padding=(self.padding_h, self.padding_w),
            dilation=(self.dilation_h, self.dilation_w),
            groups=self.conv2d.groups)
```