```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 kernel for depthwise 2D convolution using a shared-memory
# tiled approach: each thread block cooperatively loads the input tile
# (plus halo) and the per-channel filter into shared memory, then each
# thread computes one output pixel. This minimizes global memory traffic
# for this memory-bound operator.
depthwise_conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

__global__ void depthwise_conv2d_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch, int channels,
    int in_h, int in_w,
    int out_h, int out_w,
    int k_h, int k_w,
    int stride_h, int stride_w,
    int pad_h, int pad_w,
    int dil_h, int dil_w,
    int has_bias)
{
    extern __shared__ float smem[];

    // Dimensions of the input tile needed by this block (including halo).
    const int tile_w = (blockDim.x - 1) * stride_w + (k_w - 1) * dil_w + 1;
    const int tile_h = (blockDim.y - 1) * stride_h + (k_h - 1) * dil_h + 1;

    float* s_input  = smem;                  // tile_h * tile_w floats
    float* s_weight = smem + tile_h * tile_w; // k_h * k_w floats

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

    const int out_x = blockIdx.x * blockDim.x + threadIdx.x;
    const int out_y = blockIdx.y * blockDim.y + threadIdx.y;

    // Top-left corner (in input coordinates) of the tile this block needs.
    const int in_x0 = blockIdx.x * blockDim.x * stride_w - pad_w;
    const int in_y0 = blockIdx.y * blockDim.y * stride_h - pad_h;

    const long in_base = ((long)n * channels + c) * (long)in_h * in_w;

    const int tid      = threadIdx.y * blockDim.x + threadIdx.x;
    const int nthreads = blockDim.x * blockDim.y;

    // Cooperatively load the input tile (zero-fill out-of-bounds / padding).
    const int tile_size = tile_h * tile_w;
    for (int i = tid; i < tile_size; i += nthreads) {
        int ty = i / tile_w;
        int tx = i - ty * tile_w;
        int ix = in_x0 + tx;
        int iy = in_y0 + ty;
        float v = 0.0f;
        if (ix >= 0 && ix < in_w && iy >= 0 && iy < in_h) {
            v = input[in_base + (long)iy * in_w + ix];
        }
        s_input[i] = v;
    }

    // Cooperatively load this channel's filter.
    const int ksize = k_h * k_w;
    for (int i = tid; i < ksize; i += nthreads) {
        s_weight[i] = weight[(long)c * ksize + i];
    }

    __syncthreads();

    if (out_x >= out_w || out_y >= out_h) return;

    float sum = has_bias ? bias[c] : 0.0f;
    const int sx0 = threadIdx.x * stride_w;
    const int sy0 = threadIdx.y * stride_h;

    for (int kh = 0; kh < k_h; ++kh) {
        const float* srow = s_input + (sy0 + kh * dil_h) * tile_w;
        const float* wrow = s_weight + kh * k_w;
        for (int kw = 0; kw < k_w; ++kw) {
            sum += srow[sx0 + kw * dil_w] * wrow[kw];
        }
    }

    output[(((long)n * channels + c) * out_h + out_y) * out_w + out_x] = sum;
}

torch::Tensor depthwise_conv2d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    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 batch    = (int)x.size(0);
    const int channels = (int)x.size(1);
    const int in_h     = (int)x.size(2);
    const int in_w     = (int)x.size(3);
    const int k_h      = (int)w.size(2);
    const int k_w      = (int)w.size(3);

    const int out_h = (int)((in_h + 2 * pad_h - dil_h * (k_h - 1) - 1) / stride_h + 1);
    const int out_w = (int)((in_w + 2 * pad_w - dil_w * (k_w - 1) - 1) / stride_w + 1);
    TORCH_CHECK(out_h > 0 && out_w > 0, "output size must be positive");
    TORCH_CHECK((long)batch * channels <= 65535, "batch*channels too large for grid z-dim");

    auto out = torch::empty({batch, channels, out_h, out_w}, x.options());

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

    dim3 block(32, 16);
    dim3 grid((out_w + block.x - 1) / block.x,
              (out_h + block.y - 1) / block.y,
              batch * channels);

    const int tile_w = (int)((block.x - 1) * stride_w + (k_w - 1) * dil_w + 1);
    const int tile_h = (int)((block.y - 1) * stride_h + (k_h - 1) * dil_h + 1);
    const size_t smem_bytes = (size_t)(tile_w * tile_h + k_h * k_w) * sizeof(float);

    auto stream = at::cuda::getCurrentCUDAStream();
    depthwise_conv2d_kernel<<<grid, block, smem_bytes, stream>>>(
        x.data_ptr<float>(),
        w.data_ptr<float>(),
        bias_ptr,
        out.data_ptr<float>(),
        batch, channels,
        in_h, in_w,
        out_h, out_w,
        k_h, k_w,
        (int)stride_h, (int)stride_w,
        (int)pad_h, (int)pad_w,
        (int)dil_h, (int)dil_w,
        has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

depthwise_conv2d_cpp_source = (
    "torch::Tensor depthwise_conv2d_cuda("
    "torch::Tensor input, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias_opt, "
    "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_source,
    cuda_sources=depthwise_conv2d_cuda_source,
    functions=["depthwise_conv2d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom shared-memory tiled
    CUDA kernel. Parameters are kept in an nn.Conv2d module so that the
    state_dict remains fully compatible with the original 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:
        if x.is_cuda and x.dtype == torch.float32:
            return depthwise_conv2d_ext.depthwise_conv2d_cuda(
                x,
                self.conv2d.weight,
                self.conv2d.bias,
                self.stride_h, self.stride_w,
                self.padding_h, self.padding_w,
                self.dilation_h, self.dilation_w)
        # Fallback for CPU tensors or other dtypes.
        return self.conv2d(x)
```