```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 kernels for depthwise 2D convolution.
# Fast path: shared-memory tiled kernel specialized (via templates) for
# stride=1, dilation=1 and common kernel sizes (fully unrolled tap loops).
# Fallback: general kernel handling any stride/padding/dilation/kernel size.
depthwise_conv2d_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TILE_W 32
#define TILE_H 16
#define THREADS_X 32
#define THREADS_Y 16

// Tiled shared-memory depthwise conv for stride=1, dilation=1.
// Each block computes a TILE_H x TILE_W output patch for one (batch, channel).
template <int KH, int KW>
__global__ __launch_bounds__(THREADS_X * THREADS_Y)
void depthwise_conv2d_s1d1_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int pad_h, const int pad_w,
    const bool has_bias)
{
    constexpr int SH_W = TILE_W + KW - 1;
    constexpr int SH_H = TILE_H + KH - 1;
    constexpr int NT = THREADS_X * THREADS_Y;

    __shared__ float tile[SH_H * SH_W];
    __shared__ float wgt[KH * KW];

    const int c = blockIdx.z % channels;
    const int b = blockIdx.z / channels;

    const int ow = blockIdx.x * TILE_W + threadIdx.x;
    const int oh = blockIdx.y * TILE_H + threadIdx.y;

    const int iw0 = blockIdx.x * TILE_W - pad_w;
    const int ih0 = blockIdx.y * TILE_H - pad_h;

    const float* __restrict__ in_base = input + ((long)b * channels + c) * (long)in_h * in_w;

    const int tid = threadIdx.y * THREADS_X + threadIdx.x;

    // Stage the per-channel filter in shared memory (broadcast to all threads).
    for (int i = tid; i < KH * KW; i += NT)
        wgt[i] = weight[(long)c * KH * KW + i];

    // Stage the input tile (with halo / zero padding) in shared memory.
    for (int i = tid; i < SH_H * SH_W; i += NT) {
        const int ty = i / SH_W;
        const int tx = i - ty * SH_W;
        const int ih = ih0 + ty;
        const int iw = iw0 + tx;
        float v = 0.0f;
        if ((unsigned)ih < (unsigned)in_h && (unsigned)iw < (unsigned)in_w)
            v = in_base[(long)ih * in_w + iw];
        tile[i] = v;
    }
    __syncthreads();

    if (ow < out_w && oh < out_h) {
        const int tx = threadIdx.x;
        const int ty = threadIdx.y;
        float sum = 0.0f;
        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                sum = fmaf(tile[(ty + kh) * SH_W + (tx + kw)], wgt[kh * KW + kw], sum);
            }
        }
        if (has_bias) sum += bias[c];
        output[(((long)b * channels + c) * out_h + oh) * out_w + ow] = sum;
    }
}

// Fully general depthwise conv (any kernel/stride/padding/dilation).
__global__ void depthwise_conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int k_h, const int k_w,
    const int stride_h, const int stride_w,
    const int pad_h, const int pad_w,
    const int dil_h, const int dil_w,
    const bool has_bias)
{
    const int ow = blockIdx.x * blockDim.x + threadIdx.x;
    const int oh = blockIdx.y * blockDim.y + threadIdx.y;
    if (ow >= out_w || oh >= out_h) return;

    const int c = blockIdx.z % channels;
    const int b = blockIdx.z / channels;

    const float* __restrict__ in_base = input + ((long)b * channels + c) * (long)in_h * in_w;
    const float* __restrict__ w_base = weight + (long)c * k_h * k_w;

    float sum = 0.0f;
    for (int kh = 0; kh < k_h; ++kh) {
        const int ih = oh * stride_h - pad_h + kh * dil_h;
        if ((unsigned)ih >= (unsigned)in_h) continue;
        const float* in_row = in_base + (long)ih * in_w;
        const float* w_row = w_base + kh * k_w;
        for (int kw = 0; kw < k_w; ++kw) {
            const int iw = ow * stride_w - pad_w + kw * dil_w;
            if ((unsigned)iw < (unsigned)in_w)
                sum = fmaf(in_row[iw], w_row[kw], sum);
        }
    }
    if (has_bias) sum += bias[c];
    output[(((long)b * channels + c) * out_h + oh) * out_w + ow] = sum;
}

torch::Tensor depthwise_conv2d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t stride_h, int64_t stride_w,
    int64_t pad_h, int64_t pad_w,
    int64_t dil_h, int64_t dil_w)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");

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

    const int b  = x.size(0);
    const int c  = x.size(1);
    const int ih = x.size(2);
    const int iw = x.size(3);
    const int kh = w.size(2);
    const int kw = w.size(3);

    const int oh = (int)((ih + 2 * pad_h - dil_h * (kh - 1) - 1) / stride_h + 1);
    const int ow = (int)((iw + 2 * pad_w - dil_w * (kw - 1) - 1) / stride_w + 1);
    TORCH_CHECK(oh > 0 && ow > 0, "output size must be positive");

    auto out = at::empty({b, c, oh, ow}, x.options());

    const bool has_bias = bias.defined() && bias.numel() == c;
    const float* bias_ptr = has_bias ? bias.data_ptr<float>() : nullptr;

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

    const float* x_ptr = x.data_ptr<float>();
    const float* w_ptr = w.data_ptr<float>();
    float* o_ptr = out.data_ptr<float>();

    if (stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1 &&
        ((kh == 3 && kw == 7) || (kh == 3 && kw == 3) ||
         (kh == 5 && kw == 5) || (kh == 7 && kw == 7) ||
         (kh == 7 && kw == 3) || (kh == 1 && kw == 1)))
    {
        dim3 block(THREADS_X, THREADS_Y, 1);
        dim3 grid((ow + TILE_W - 1) / TILE_W, (oh + TILE_H - 1) / TILE_H, b * c);
        if (kh == 3 && kw == 7)
            depthwise_conv2d_s1d1_kernel<3, 7><<<grid, block, 0, stream>>>(
                x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow, (int)pad_h, (int)pad_w, has_bias);
        else if (kh == 3 && kw == 3)
            depthwise_conv2d_s1d1_kernel<3, 3><<<grid, block, 0, stream>>>(
                x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow, (int)pad_h, (int)pad_w, has_bias);
        else if (kh == 5 && kw == 5)
            depthwise_conv2d_s1d1_kernel<5, 5><<<grid, block, 0, stream>>>(
                x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow, (int)pad_h, (int)pad_w, has_bias);
        else if (kh == 7 && kw == 7)
            depthwise_conv2d_s1d1_kernel<7, 7><<<grid, block, 0, stream>>>(
                x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow, (int)pad_h, (int)pad_w, has_bias);
        else if (kh == 7 && kw == 3)
            depthwise_conv2d_s1d1_kernel<7, 3><<<grid, block, 0, stream>>>(
                x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow, (int)pad_h, (int)pad_w, has_bias);
        else
            depthwise_conv2d_s1d1_kernel<1, 1><<<grid, block, 0, stream>>>(
                x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow, (int)pad_h, (int)pad_w, has_bias);
    }
    else
    {
        dim3 block(32, 8, 1);
        dim3 grid((ow + 31) / 32, (oh + 7) / 8, b * c);
        depthwise_conv2d_generic_kernel<<<grid, block, 0, stream>>>(
            x_ptr, w_ptr, bias_ptr, o_ptr, c, ih, iw, oh, ow,
            kh, kw, (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w, (int)dil_h, (int)dil_w, has_bias);
    }

    return out;
}
"""

depthwise_conv2d_cpp_src = (
    "torch::Tensor depthwise_conv2d_cuda(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride_h, int64_t stride_w, int64_t pad_h, "
    "int64_t pad_w, int64_t dil_h, int64_t dil_w);"
)

# Compile the inline CUDA extension.
depthwise_conv2d_ext = load_inline(
    name="depthwise_conv2d_ext",
    cpp_sources=depthwise_conv2d_cpp_src,
    cuda_sources=depthwise_conv2d_cuda_src,
    functions=["depthwise_conv2d_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom shared-memory tiled
    CUDA kernel. Falls back to the standard cuDNN conv2d for unsupported
    configurations (non-CUDA / non-float32 inputs).
    """
    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 identical parameter layout to the reference model so that
        # state dicts are interchangeable.
        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
        self._dwconv = depthwise_conv2d_ext.depthwise_conv2d_cuda

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32:
            weight = self.conv2d.weight
            bias = self.conv2d.bias
            if bias is None:
                bias = x.new_empty(0)
            return self._dwconv(
                x, weight, bias,
                self.stride_h, self.stride_w,
                self.padding_h, self.padding_w,
                self.dilation_h, self.dilation_w)
        return self.conv2d(x)
```