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

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

#define BLOCK_W 32
#define BLOCK_H 8

// Tiled depthwise 2D convolution kernel (NCHW, depth multiplier = 1).
// Each block computes a BLOCK_W x BLOCK_H output tile for one (batch, channel).
// The required input halo tile and the (small) filter are staged in shared memory.
template <int KH, int KW>
__global__ void depthwise_conv2d_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 stride_h, const int stride_w,
    const int pad_h, const int pad_w,
    const int dil_h, const int dil_w,
    const int has_bias)
{
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int bc = blockIdx.z;           // fused (batch * channel) index
    const int c = bc % channels;

    const int out_x = blockIdx.x * BLOCK_W + tx;
    const int out_y = blockIdx.y * BLOCK_H + ty;

    const int tile_w = (BLOCK_W - 1) * stride_w + (KW - 1) * dil_w + 1;
    const int tile_h = (BLOCK_H - 1) * stride_h + (KH - 1) * dil_h + 1;

    extern __shared__ float smem[];
    float* tile = smem;                       // tile_h x tile_w
    float* sw   = smem + tile_w * tile_h;     // KH * KW filter weights

    const int tid = ty * BLOCK_W + tx;
    if (tid < KH * KW) {
        sw[tid] = weight[c * (KH * KW) + tid];
    }

    const int in_x0 = blockIdx.x * BLOCK_W * stride_w - pad_w;
    const int in_y0 = blockIdx.y * BLOCK_H * stride_h - pad_h;
    const float* in_ptr = input + (long)bc * in_h * in_w;

    // Stage the input halo tile into shared memory (zero padding outside).
    for (int i = ty; i < tile_h; i += BLOCK_H) {
        const int iy = in_y0 + i;
        const bool row_ok = (iy >= 0) && (iy < in_h);
        for (int j = tx; j < tile_w; j += BLOCK_W) {
            const int ix = in_x0 + j;
            float v = 0.0f;
            if (row_ok && ix >= 0 && ix < in_w) {
                v = in_ptr[iy * in_w + ix];
            }
            tile[i * tile_w + j] = v;
        }
    }
    __syncthreads();

    if (out_x < out_w && out_y < out_h) {
        float acc = has_bias ? bias[c] : 0.0f;
        const int base_y = ty * stride_h;
        const int base_x = tx * stride_w;
        #pragma unroll
        for (int i = 0; i < KH; ++i) {
            #pragma unroll
            for (int j = 0; j < KW; ++j) {
                acc = fmaf(tile[(base_y + i * dil_h) * tile_w + base_x + j * dil_w],
                           sw[i * KW + j], acc);
            }
        }
        output[((long)bc * out_h + out_y) * out_w + out_x] = acc;
    }
}

// Generic fallback (runtime kernel size).
__global__ void depthwise_conv2d_kernel_generic(
    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 kh, const int kw,
    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 int has_bias)
{
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int bc = blockIdx.z;
    const int c = bc % channels;

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

    const int tile_w = (blockDim.x - 1) * stride_w + (kw - 1) * dil_w + 1;
    const int tile_h = (blockDim.y - 1) * stride_h + (kh - 1) * dil_h + 1;

    extern __shared__ float smem[];
    float* tile = smem;
    float* sw   = smem + tile_w * tile_h;

    for (int i = ty * blockDim.x + tx; i < kh * kw; i += blockDim.x * blockDim.y) {
        sw[i] = weight[c * kh * kw + i];
    }

    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 float* in_ptr = input + (long)bc * in_h * in_w;

    for (int i = ty; i < tile_h; i += blockDim.y) {
        const int iy = in_y0 + i;
        const bool row_ok = (iy >= 0) && (iy < in_h);
        for (int j = tx; j < tile_w; j += blockDim.x) {
            const int ix = in_x0 + j;
            float v = 0.0f;
            if (row_ok && ix >= 0 && ix < in_w) {
                v = in_ptr[iy * in_w + ix];
            }
            tile[i * tile_w + j] = v;
        }
    }
    __syncthreads();

    if (out_x < out_w && out_y < out_h) {
        float acc = has_bias ? bias[c] : 0.0f;
        const int base_y = ty * stride_h;
        const int base_x = tx * stride_w;
        for (int i = 0; i < kh; ++i) {
            for (int j = 0; j < kw; ++j) {
                acc = fmaf(tile[(base_y + i * dil_h) * tile_w + base_x + j * dil_w],
                           sw[i * kw + j], acc);
            }
        }
        output[((long)bc * out_h + out_y) * out_w + out_x] = acc;
    }
}

torch::Tensor depthwise_conv2d_forward(
    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() && weight.is_cuda(), "input/weight must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "input/weight must be 4D");

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

    const int64_t batch    = x.size(0);
    const int64_t channels = x.size(1);
    const int64_t in_h     = x.size(2);
    const int64_t in_w     = x.size(3);
    const int64_t kh       = w.size(2);
    const int64_t kw       = w.size(3);
    TORCH_CHECK(w.size(0) == channels && w.size(1) == 1,
                "only depthwise convolution with depth multiplier 1 is supported");

    const int64_t out_h = (in_h + 2 * pad_h - dil_h * (kh - 1) - 1) / stride_h + 1;
    const int64_t out_w = (in_w + 2 * pad_w - dil_w * (kw - 1) - 1) / stride_w + 1;
    TORCH_CHECK(out_h > 0 && out_w > 0, "output spatial size must be positive");
    TORCH_CHECK(batch * channels <= 65535, "batch * channels too large");

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

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

    const dim3 block(BLOCK_W, BLOCK_H);
    const dim3 grid((unsigned)((out_w + BLOCK_W - 1) / BLOCK_W),
                    (unsigned)((out_h + BLOCK_H - 1) / BLOCK_H),
                    (unsigned)(batch * channels));

    const int tile_w = (BLOCK_W - 1) * (int)stride_w + ((int)kw - 1) * (int)dil_w + 1;
    const int tile_h = (BLOCK_H - 1) * (int)stride_h + ((int)kh - 1) * (int)dil_h + 1;
    const size_t shmem = (size_t)(tile_w * tile_h + kh * kw) * sizeof(float);

    auto 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>();

    #define LAUNCH_TILED(KH_, KW_)                                                        \
        depthwise_conv2d_kernel<KH_, KW_><<<grid, block, shmem, stream>>>(                \
            x_ptr, w_ptr, b_ptr, o_ptr, (int)channels, (int)in_h, (int)in_w,              \
            (int)out_h, (int)out_w, (int)stride_h, (int)stride_w,                         \
            (int)pad_h, (int)pad_w, (int)dil_h, (int)dil_w, has_bias)

    if (kh == 3 && kw == 7)      { LAUNCH_TILED(3, 7); }
    else if (kh == 3 && kw == 3) { LAUNCH_TILED(3, 3); }
    else if (kh == 5 && kw == 5) { LAUNCH_TILED(5, 5); }
    else if (kh == 7 && kw == 7) { LAUNCH_TILED(7, 7); }
    else if (kh == 1 && kw == 1) { LAUNCH_TILED(1, 1); }
    else if (kh == 3 && kw == 1) { LAUNCH_TILED(3, 1); }
    else if (kh == 1 && kw == 3) { LAUNCH_TILED(1, 3); }
    else if (kh == 5 && kw == 1) { LAUNCH_TILED(5, 1); }
    else if (kh == 1 && kw == 5) { LAUNCH_TILED(1, 5); }
    else if (kh == 7 && kw == 1) { LAUNCH_TILED(7, 1); }
    else if (kh == 1 && kw == 7) { LAUNCH_TILED(1, 7); }
    else {
        depthwise_conv2d_kernel_generic<<<grid, block, shmem, stream>>>(
            x_ptr, w_ptr, b_ptr, o_ptr, (int)channels, (int)in_h, (int)in_w,
            (int)out_h, (int)out_w, (int)kh, (int)kw, (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w, (int)dil_h, (int)dil_w, has_bias);
    }
    #undef LAUNCH_TILED

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_depthwise_conv2d_cpp_src = (
    "torch::Tensor depthwise_conv2d_forward(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);"
)

_depthwise_conv2d_ext = load_inline(
    name="depthwise_conv2d_ext_v1",
    cpp_sources=_depthwise_conv2d_cpp_src,
    cuda_sources=_depthwise_conv2d_cuda_src,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom tiled CUDA kernel
    (shared-memory input halo staging + register-cached filter + fused bias).
    """

    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
        self._op = _depthwise_conv2d_ext

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (x.is_cuda and x.dtype == torch.float32 and x.dim() == 4
                and self.conv2d.weight.size(0) == x.size(1)
                and self.conv2d.weight.size(1) == 1):
            b = self.conv2d.bias
            if b is None:
                b = torch.empty(0, device=x.device, dtype=torch.float32)
            return self._op.depthwise_conv2d_forward(
                x, self.conv2d.weight, b,
                self.stride_h, self.stride_w,
                self.padding_h, self.padding_w,
                self.dilation_h, self.dilation_w)
        return self.conv2d(x)
```