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

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

#define KH @KH@
#define KW @KW@
#define STRIDE_H @SH@
#define STRIDE_W @SW@
#define PAD_H @PH@
#define PAD_W @PW@
#define DIL_H @DH@
#define DIL_W @DW@

#define TILE_W 32
#define TILE_H 16

#define IN_TILE_W ((TILE_W - 1) * STRIDE_W + (KW - 1) * DIL_W + 1)
#define IN_TILE_H ((TILE_H - 1) * STRIDE_H + (KH - 1) * DIL_H + 1)

// Depthwise 2D convolution with shared-memory tiled input reuse.
// Each block computes a TILE_H x TILE_W output tile for one (batch, channel).
__global__ void depthwise_conv2d_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int C, const int H, const int W,
    const int outH, const int outW)
{
    const int c = blockIdx.z % C;
    const int n = blockIdx.z / C;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int ow0 = blockIdx.x * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int iw0 = ow0 * STRIDE_W - PAD_W;
    const int ih0 = oh0 * STRIDE_H - PAD_H;

    __shared__ float tile[IN_TILE_H][IN_TILE_W];
    __shared__ float ws[KH * KW];

    const float* __restrict__ xnc = x + ((long)n * C + c) * (long)H * W;

    const int tid = ty * TILE_W + tx;
    const int nthreads = TILE_W * TILE_H;

    // Load this channel's filter into shared memory.
    for (int i = tid; i < KH * KW; i += nthreads) {
        ws[i] = weight[(long)c * KH * KW + i];
    }

    // Load the input tile (with halo), zero-filling out-of-bounds (padding).
    for (int i = tid; i < IN_TILE_H * IN_TILE_W; i += nthreads) {
        const int r = i / IN_TILE_W;
        const int col = i - r * IN_TILE_W;
        const int ih = ih0 + r;
        const int iw = iw0 + col;
        float v = 0.0f;
        if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
            v = xnc[(long)ih * W + iw];
        }
        tile[r][col] = v;
    }
    __syncthreads();

    const int oh = oh0 + ty;
    const int ow = ow0 + tx;
    if (oh < outH && ow < outW) {
        float acc = (bias != nullptr) ? bias[c] : 0.0f;
        const int sr = ty * STRIDE_H;
        const int sc = tx * STRIDE_W;
        #pragma unroll
        for (int i = 0; i < KH; ++i) {
            #pragma unroll
            for (int j = 0; j < KW; ++j) {
                acc = fmaf(tile[sr + i * DIL_H][sc + j * DIL_W], ws[i * KW + j], acc);
            }
        }
        y[(((long)n * C + c) * outH + oh) * outW + ow] = acc;
    }
}

torch::Tensor depthwise_conv2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias)
{
    TORCH_CHECK(x.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Input must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "Weight must be float32");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    const int N = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int H = (int)xc.size(2);
    const int W = (int)xc.size(3);

    const int outH = (H + 2 * PAD_H - DIL_H * (KH - 1) - 1) / STRIDE_H + 1;
    const int outW = (W + 2 * PAD_W - DIL_W * (KW - 1) - 1) / STRIDE_W + 1;

    auto y = at::empty({N, C, outH, outW}, xc.options());
    if (y.numel() == 0) {
        return y;
    }

    const float* bias_ptr = nullptr;
    torch::Tensor bc;
    if (bias.defined() && bias.numel() > 0) {
        bc = bias.contiguous();
        bias_ptr = bc.data_ptr<float>();
    }

    dim3 block(TILE_W, TILE_H);
    dim3 grid((outW + TILE_W - 1) / TILE_W, (outH + TILE_H - 1) / TILE_H, N * C);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    depthwise_conv2d_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(),
        wc.data_ptr<float>(),
        bias_ptr,
        y.data_ptr<float>(),
        N, C, H, W, outH, outW);

    return y;
}
"""

_CPP_SRC = "torch::Tensor depthwise_conv2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);"

_module_cache = {}


def _get_dwconv_module(kh, kw, sh, sw, ph, pw, dh, dw):
    key = (kh, kw, sh, sw, ph, pw, dh, dw)
    mod = _module_cache.get(key)
    if mod is not None:
        return mod
    src = (_CUDA_SRC_TEMPLATE
           .replace("@KH@", str(kh)).replace("@KW@", str(kw))
           .replace("@SH@", str(sh)).replace("@SW@", str(sw))
           .replace("@PH@", str(ph)).replace("@PW@", str(pw))
           .replace("@DH@", str(dh)).replace("@DW@", str(dw)))
    name = "depthwise_conv2d_" + "_".join(str(v) for v in key)
    mod = load_inline(
        name=name,
        cpp_sources=_CPP_SRC,
        cuda_sources=src,
        functions=["depthwise_conv2d_forward"],
        verbose=False,
        extra_cuda_cflags=["-O3", "--use_fast_math"],
    )
    _module_cache[key] = mod
    return mod


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

    The kernel is specialized at compile time for the given kernel size,
    stride, padding and dilation so the inner filter loop is fully unrolled.
    It uses a shared-memory input tile (with halo) so each input value is
    fetched from global memory only ~once per block.
    """
    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__()
        # Keep an nn.Conv2d to own the parameters (state_dict compatible).
        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._dwconv = _get_dwconv_module(
            kernel_size_h, kernel_size_w, stride_h, stride_w,
            padding_h, padding_w, dilation_h, dilation_w)
        self._no_bias = torch.empty(0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        n, c = x.shape[0], x.shape[1]
        if x.is_cuda and x.dtype == torch.float32 and n * c <= 65535:
            bias = self.conv2d.bias
            if bias is None:
                bias = self._no_bias
            return self._dwconv.depthwise_conv2d_forward(x, self.conv2d.weight, bias)
        return self.conv2d(x)
```