```
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Custom CUDA implementation of depthwise 2D convolution (NCHW, fp32).
#
# Two kernels are provided:
#   1) A fast templated kernel for the common case stride=1, padding=0,
#      dilation=1. Each thread computes TILE_W consecutive outputs along the
#      width dimension so input values are reused from registers, weights are
#      cached in registers, and (when alignment allows) float4 loads are used.
#   2) A fully generic kernel that supports any kernel size / stride /
#      padding / dilation and optional bias.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <optional>

namespace {

constexpr int TILE_W = 4;

// Fast path: stride == 1, padding == 0, dilation == 1.
// Grid: (ceil(out_w / (blockDim.x * TILE_W)), out_h, batch * channels)
// Each thread computes TILE_W consecutive outputs along width.
template <int KH, int KW>
__global__ void dwconv2d_fast_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 bool has_bias,
    const bool vec_ok)
{
    const int ow = (blockIdx.x * blockDim.x + threadIdx.x) * TILE_W;
    const int oh = blockIdx.y;
    const int bc = blockIdx.z;

    if (ow >= out_w) return;

    const int c = bc % channels;

    // Cache the (per-channel) filter in registers.
    const float* wptr = weight + (size_t)c * (KH * KW);
    float wr[KH][KW];
#pragma unroll
    for (int i = 0; i < KH; ++i) {
#pragma unroll
        for (int j = 0; j < KW; ++j) {
            wr[i][j] = wptr[i * KW + j];
        }
    }

    const float* in_base = input + (size_t)bc * in_h * in_w;
    float* out_ptr = output + ((size_t)bc * out_h + oh) * out_w + ow;

    float acc[TILE_W];
    const float bval = has_bias ? bias[c] : 0.0f;
#pragma unroll
    for (int t = 0; t < TILE_W; ++t) acc[t] = bval;

    if (ow + TILE_W <= out_w) {
        // Interior thread: every access is guaranteed in-bounds.
#pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            const float* in_row = in_base + (size_t)(oh + kh) * in_w + ow;
            float v[TILE_W + KW - 1];
            if (vec_ok && (TILE_W + KW - 1) >= 8) {
                const float4 a = *reinterpret_cast<const float4*>(in_row);
                const float4 b = *reinterpret_cast<const float4*>(in_row + 4);
                v[0] = a.x; v[1] = a.y; v[2] = a.z; v[3] = a.w;
                v[4] = b.x; v[5] = b.y; v[6] = b.z; v[7] = b.w;
#pragma unroll
                for (int j = 8; j < TILE_W + KW - 1; ++j) v[j] = __ldg(in_row + j);
            } else {
#pragma unroll
                for (int j = 0; j < TILE_W + KW - 1; ++j) v[j] = __ldg(in_row + j);
            }
#pragma unroll
            for (int t = 0; t < TILE_W; ++t) {
#pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    acc[t] = fmaf(v[t + kw], wr[kh][kw], acc[t]);
                }
            }
        }
#pragma unroll
        for (int t = 0; t < TILE_W; ++t) out_ptr[t] = acc[t];
    } else {
        // Right-edge thread: guard loads and stores.
#pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            const float* in_row = in_base + (size_t)(oh + kh) * in_w + ow;
            float v[TILE_W + KW - 1];
#pragma unroll
            for (int j = 0; j < TILE_W + KW - 1; ++j) {
                v[j] = (ow + j < in_w) ? __ldg(in_row + j) : 0.0f;
            }
#pragma unroll
            for (int t = 0; t < TILE_W; ++t) {
#pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    acc[t] = fmaf(v[t + kw], wr[kh][kw], acc[t]);
                }
            }
        }
#pragma unroll
        for (int t = 0; t < TILE_W; ++t) {
            if (ow + t < out_w) out_ptr[t] = acc[t];
        }
    }
}

// Generic fallback: any kernel size / stride / padding / dilation, one
// thread per output element.
__global__ void dwconv2d_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 s_h, const int s_w,
    const int p_h, const int p_w,
    const int d_h, const int d_w,
    const bool has_bias,
    const long long total)
{
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int ow = (int)(idx % out_w);
    const long long tmp = idx / out_w;
    const int oh = (int)(tmp % out_h);
    const long long bc = tmp / out_h;
    const int c = (int)(bc % channels);

    const float* in_base = input + (size_t)bc * in_h * in_w;
    const float* wptr = weight + (size_t)c * k_h * k_w;

    float acc = has_bias ? bias[c] : 0.0f;
    for (int kh = 0; kh < k_h; ++kh) {
        const int ih = oh * s_h - p_h + kh * d_h;
        if (ih < 0 || ih >= in_h) continue;
        const float* in_row = in_base + (size_t)ih * in_w;
        const float* w_row = wptr + kh * k_w;
        for (int kw = 0; kw < k_w; ++kw) {
            const int iw = ow * s_w - p_w + kw * d_w;
            if (iw >= 0 && iw < in_w) {
                acc = fmaf(in_row[iw], w_row[kw], acc);
            }
        }
    }
    output[idx] = acc;
}

}  // namespace

torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    std::optional<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.dim() == 4, "input must be 4D (NCHW)");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(weight.is_cuda() && weight.dim() == 4, "weight must be a 4D CUDA tensor");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");

    const c10::cuda::CUDAGuard device_guard(input.device());

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

    const int64_t B  = x.size(0);
    const int64_t C  = x.size(1);
    const int64_t IH = x.size(2);
    const int64_t IW = x.size(3);
    const int64_t KH = w.size(2);
    const int64_t KW = w.size(3);

    const int64_t OH = (IH + 2 * pad_h - dil_h * (KH - 1) - 1) / stride_h + 1;
    const int64_t OW = (IW + 2 * pad_w - dil_w * (KW - 1) - 1) / stride_w + 1;
    TORCH_CHECK(OH > 0 && OW > 0, "computed output size must be positive");

    auto out = at::empty({B, C, OH, OW}, x.options());

    const int64_t total = B * C * OH * OW;
    if (total == 0) return out;

    torch::Tensor bias_c;
    const float* bias_ptr = nullptr;
    bool has_bias = false;
    if (bias.has_value() && bias->defined() && bias->numel() > 0) {
        TORCH_CHECK(bias->scalar_type() == at::kFloat, "bias must be float32");
        bias_c = bias->contiguous();
        bias_ptr = bias_c.data_ptr<float>();
        has_bias = true;
    }

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

    const bool fast =
        stride_h == 1 && stride_w == 1 &&
        pad_h == 0 && pad_w == 0 &&
        dil_h == 1 && dil_w == 1 &&
        OH <= 65535 && (B * C) <= 65535;

    if (fast && KH == 3 && KW == 7) {
        const dim3 block(64);
        const dim3 grid(
            (unsigned)((OW + (int64_t)block.x * TILE_W - 1) / ((int64_t)block.x * TILE_W)),
            (unsigned)OH,
            (unsigned)(B * C));
        const bool vec_ok =
            (IW % 4 == 0) &&
            ((reinterpret_cast<uintptr_t>(x.data_ptr<float>()) & 0xF) == 0);
        dwconv2d_fast_kernel<3, 7><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(),
            w.data_ptr<float>(),
            bias_ptr,
            out.data_ptr<float>(),
            (int)C, (int)IH, (int)IW, (int)OH, (int)OW,
            has_bias, vec_ok);
    } else {
        const int block = 256;
        const int64_t grid64 = (total + block - 1) / block;
        TORCH_CHECK(grid64 <= 0x7fffffff, "grid dimension too large");
        dwconv2d_generic_kernel<<<(unsigned)grid64, block, 0, stream>>>(
            x.data_ptr<float>(),
            w.data_ptr<float>(),
            bias_ptr,
            out.data_ptr<float>(),
            (int)C, (int)IH, (int)IW, (int)OH, (int)OW,
            (int)KH, (int)KW,
            (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w,
            (int)dil_h, (int)dil_w,
            has_bias, (long long)total);
    }

    return out;
}
"""

cpp_source = r"""
#include <optional>
torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    std::optional<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_custom_v1",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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

    Keeps the same parameters (self.conv2d.weight / self.conv2d.bias) as the
    original model so state_dicts are fully compatible, but dispatches the
    forward pass to a hand-written CUDA kernel.
    """
    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:
        w = self.conv2d.weight
        if (
            x.is_cuda
            and w.is_cuda
            and x.dtype == torch.float32
            and w.dtype == torch.float32
        ):
            return _depthwise_conv2d_ext.depthwise_conv2d_forward(
                x,
                w,
                self.conv2d.bias,
                self._stride_h,
                self._stride_w,
                self._padding_h,
                self._padding_w,
                self._dilation_h,
                self._dilation_w,
            )
        return self.conv2d(x)
```