```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_cpp_source = """
torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t strideH, int64_t strideW,
    int64_t padH, int64_t padW,
    int64_t dilH, int64_t dilW);
"""

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

#define TILE_H 16
#define TILE_W 32
#define THREADS (TILE_H * TILE_W)

// Tiled shared-memory depthwise conv (specialized for stride=1, dilation=1).
// Each block computes a TILE_H x TILE_W output tile for one (batch, channel) slice.
template <int KH, int KW>
__global__ void __launch_bounds__(THREADS) depthwise_conv2d_tiled_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int C, const int H, const int W,
    const int outH, const int outW,
    const int padH, const int padW,
    const bool has_bias)
{
    constexpr int SH = TILE_H + KH - 1;
    constexpr int SW = TILE_W + KW - 1;

    __shared__ float tile[SH * SW];
    __shared__ float wgt[KH * KW];
    __shared__ float bias_val;

    const int nc  = blockIdx.z;              // n * C + c
    const int c   = nc % C;
    const int h0  = blockIdx.y * TILE_H;
    const int w0  = blockIdx.x * TILE_W;
    const int tid = threadIdx.x;

    if (tid < KH * KW) {
        wgt[tid] = weight[static_cast<long>(c) * (KH * KW) + tid];
    }
    if (tid == 0) {
        bias_val = has_bias ? bias[c] : 0.0f;
    }

    const float* in_nc = input + static_cast<long>(nc) * H * W;
    const int h_base = h0 - padH;
    const int w_base = w0 - padW;

    // Cooperative, coalesced load of the input tile (zero-filled out of bounds).
    for (int i = tid; i < SH * SW; i += THREADS) {
        const int sh = i / SW;
        const int sw = i - sh * SW;
        const int h_in = h_base + sh;
        const int w_in = w_base + sw;
        float v = 0.0f;
        if (static_cast<unsigned>(h_in) < static_cast<unsigned>(H) &&
            static_cast<unsigned>(w_in) < static_cast<unsigned>(W)) {
            v = in_nc[h_in * W + w_in];
        }
        tile[i] = v;
    }
    __syncthreads();

    const int ty = tid / TILE_W;
    const int tx = tid - ty * TILE_W;
    const int h_out = h0 + ty;
    const int w_out = w0 + tx;

    if (h_out < outH && w_out < outW) {
        // Cache the (broadcast) weights in registers to cut shared-memory traffic.
        float wr[KH * KW];
        #pragma unroll
        for (int i = 0; i < KH * KW; ++i) {
            wr[i] = wgt[i];
        }
        float acc = bias_val;
        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                acc = fmaf(tile[(ty + kh) * SW + tx + kw], wr[kh * KW + kw], acc);
            }
        }
        output[(static_cast<long>(nc) * outH + h_out) * outW + w_out] = acc;
    }
}

// Fully generic fallback: one thread per output element.
__global__ void depthwise_conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int C, const int H, const int W,
    const int outH, const int outW,
    const int kH, const int kW,
    const int sH, const int sW,
    const int pH, const int pW,
    const int dH, const int dW,
    const bool has_bias,
    const long total)
{
    const long idx = blockIdx.x * static_cast<long>(blockDim.x) + threadIdx.x;
    if (idx >= total) return;

    const long nc = idx / (static_cast<long>(outH) * outW);
    const long within = idx - nc * static_cast<long>(outH) * outW;
    const int h_out = static_cast<int>(within / outW);
    const int w_out = static_cast<int>(within - static_cast<long>(h_out) * outW);
    const int c = static_cast<int>(nc % C);

    const float* in_c = input + nc * static_cast<long>(H) * W;
    const float* w_c = weight + static_cast<long>(c) * kH * kW;

    float acc = has_bias ? bias[c] : 0.0f;
    for (int kh = 0; kh < kH; ++kh) {
        const int h_in = h_out * sH + kh * dH - pH;
        if (h_in < 0 || h_in >= H) continue;
        const float* in_row = in_c + static_cast<long>(h_in) * W;
        const float* w_row = w_c + kh * kW;
        for (int kw = 0; kw < kW; ++kw) {
            const int w_in = w_out * sW + kw * dW - pW;
            if (w_in < 0 || w_in >= W) continue;
            acc = fmaf(in_row[w_in], w_row[kw], acc);
        }
    }
    output[idx] = acc;
}

torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t strideH, int64_t strideW,
    int64_t padH, int64_t padW,
    int64_t dilH, int64_t dilW)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(input.dim() == 4, "input must be 4D NCHW");
    TORCH_CHECK(weight.dim() == 4, "weight must be 4D");

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

    auto x = input.contiguous();
    auto w = weight.contiguous();
    const bool has_bias = bias.defined() && bias.numel() > 0;
    torch::Tensor b;
    if (has_bias) {
        b = bias.contiguous();
    }

    const int N = static_cast<int>(x.size(0));
    const int C = static_cast<int>(x.size(1));
    const int H = static_cast<int>(x.size(2));
    const int W = static_cast<int>(x.size(3));
    const int kH = static_cast<int>(w.size(2));
    const int kW = static_cast<int>(w.size(3));

    const int outH = static_cast<int>((H + 2 * padH - dilH * (kH - 1) - 1) / strideH + 1);
    const int outW = static_cast<int>((W + 2 * padW - dilW * (kW - 1) - 1) / strideW + 1);

    auto out = at::empty({N, C, outH, outW}, x.options());

    const float* xptr = x.data_ptr<float>();
    const float* wptr = w.data_ptr<float>();
    const float* bptr = has_bias ? b.data_ptr<float>() : nullptr;
    float* optr = out.data_ptr<float>();

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

    const bool fast_ok = strideH == 1 && strideW == 1 && dilH == 1 && dilW == 1 &&
        static_cast<long>(N) * C <= 65535 &&
        (outH + TILE_H - 1) / TILE_H <= 65535;

    if (fast_ok) {
        dim3 grid((outW + TILE_W - 1) / TILE_W,
                  (outH + TILE_H - 1) / TILE_H,
                  N * C);
        dim3 block(THREADS);
        bool launched = true;
        #define LAUNCH_TILED(KH_, KW_)                                                   \
            depthwise_conv2d_tiled_kernel<KH_, KW_><<<grid, block, 0, stream>>>(         \
                xptr, wptr, bptr, optr, C, H, W, outH, outW,                             \
                static_cast<int>(padH), static_cast<int>(padW), has_bias)
        if      (kH == 3 && kW == 7) { LAUNCH_TILED(3, 7); }
        else if (kH == 7 && kW == 3) { LAUNCH_TILED(7, 3); }
        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 == 1 && kW == 3) { LAUNCH_TILED(1, 3); }
        else if (kH == 3 && kW == 1) { LAUNCH_TILED(3, 1); }
        else if (kH == 1 && kW == 5) { LAUNCH_TILED(1, 5); }
        else if (kH == 5 && kW == 1) { LAUNCH_TILED(5, 1); }
        else if (kH == 1 && kW == 7) { LAUNCH_TILED(1, 7); }
        else if (kH == 7 && kW == 1) { LAUNCH_TILED(7, 1); }
        else if (kH == 3 && kW == 5) { LAUNCH_TILED(3, 5); }
        else if (kH == 5 && kW == 3) { LAUNCH_TILED(5, 3); }
        else { launched = false; }
        #undef LAUNCH_TILED
        if (launched) {
            C10_CUDA_KERNEL_LAUNCH_CHECK();
            return out;
        }
    }

    const long total = static_cast<long>(N) * C * outH * outW;
    const int threads = 256;
    const long blocks = (total + threads - 1) / threads;
    depthwise_conv2d_generic_kernel<<<static_cast<unsigned>(blocks), threads, 0, stream>>>(
        xptr, wptr, bptr, optr, C, H, W, outH, outW, kH, kW,
        static_cast<int>(strideH), static_cast<int>(strideW),
        static_cast<int>(padH), static_cast<int>(padW),
        static_cast<int>(dilH), static_cast<int>(dilW),
        has_bias, total);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_depthwise_conv2d = load_inline(
    name="depthwise_conv2d_custom",
    cpp_sources=_depthwise_conv2d_cpp_source,
    cuda_sources=_depthwise_conv2d_cuda_source,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution with asymmetric input/kernel, accelerated by a
    custom tiled shared-memory CUDA kernel (with a generic CUDA fallback).
    """
    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.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.conv2d.weight
        if (x.is_cuda and x.dtype == torch.float32 and x.dim() == 4
                and weight.is_cuda and weight.dtype == torch.float32):
            bias = self.conv2d.bias
            if bias is None:
                bias = self._empty_bias
            return _depthwise_conv2d.depthwise_conv2d_forward(
                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)
```