```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 implementation of depthwise 2D convolution.
# Fast path: shared-memory tiled kernel specialized for stride=1, dilation=1
# with a compile-time 3x7 kernel (fully unrolled), each thread block producing
# a 32 x 32 output tile (each thread computes 2 vertically separated outputs).
# Generic path: one thread per output element supporting any kernel size,
# stride, padding and dilation.
depthwise_conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr int TILE_W = 32;
constexpr int TILE_H = 16;

// Specialized kernel for stride == 1 and dilation == 1 with compile-time
// kernel dimensions. Padding is handled via bounds-checked tile loads.
template <int KH, int KW, int ROWS>
__global__ void depthwise_conv2d_s1d1_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 OUT_TILE_H = TILE_H * ROWS;
    constexpr int TILE_IN_W = TILE_W + KW - 1;
    constexpr int TILE_IN_H = OUT_TILE_H + KH - 1;

    __shared__ float tile[TILE_IN_H][TILE_IN_W];
    __shared__ float wgt[KH * KW];

    const int nc = blockIdx.z;          // n * C + c
    const int c = nc % C;

    const int out_w0 = blockIdx.x * TILE_W;
    const int out_h0 = blockIdx.y * OUT_TILE_H;
    const int in_w0 = out_w0 - padW;
    const int in_h0 = out_h0 - padH;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * TILE_W + tx;
    constexpr int NT = TILE_W * TILE_H;

    const float* in_base = input + (size_t)nc * (size_t)H * (size_t)W;

    // Load this channel's kernel weights into shared memory.
    if (tid < KH * KW) {
        wgt[tid] = weight[(size_t)c * (size_t)(KH * KW) + tid];
    }

    // Load the input tile (with halo) into shared memory, zero-padding
    // out-of-bounds locations.
    #pragma unroll
    for (int idx = tid; idx < TILE_IN_H * TILE_IN_W; idx += NT) {
        const int lh = idx / TILE_IN_W;
        const int lw = idx - lh * TILE_IN_W;
        const int ih = in_h0 + lh;
        const int iw = in_w0 + lw;
        float v = 0.0f;
        if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
            v = in_base[(size_t)ih * (size_t)W + (size_t)iw];
        }
        tile[lh][lw] = v;
    }
    __syncthreads();

    const int ow = out_w0 + tx;
    if (ow < outW) {
        #pragma unroll
        for (int r = 0; r < ROWS; ++r) {
            const int ty_r = ty + r * TILE_H;
            const int oh = out_h0 + ty_r;
            if (oh < outH) {
                float acc = has_bias ? bias[c] : 0.0f;
                #pragma unroll
                for (int kh = 0; kh < KH; ++kh) {
                    #pragma unroll
                    for (int kw = 0; kw < KW; ++kw) {
                        acc = fmaf(tile[ty_r + kh][tx + kw], wgt[kh * KW + kw], acc);
                    }
                }
                output[((size_t)nc * (size_t)outH + (size_t)oh) * (size_t)outW + (size_t)ow] = acc;
            }
        }
    }
}

// Generic fallback kernel: one thread per output element, supports arbitrary
// kernel size, stride, padding and dilation.
__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 strideH, const int strideW,
    const int padH, const int padW,
    const int dilH, const int dilW,
    const bool has_bias,
    const long total)
{
    const long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int ow = (int)(idx % outW);
    const int oh = (int)((idx / outW) % outH);
    const long nc = idx / ((long)outH * (long)outW);
    const int c = (int)(nc % C);

    const float* in_base = input + (size_t)nc * (size_t)H * (size_t)W;
    const float* w_base = weight + (size_t)c * (size_t)(KH * KW);

    float acc = has_bias ? bias[c] : 0.0f;
    for (int kh = 0; kh < KH; ++kh) {
        const int ih = oh * strideH - padH + kh * dilH;
        if (ih < 0 || ih >= H) continue;
        const float* in_row = in_base + (size_t)ih * (size_t)W;
        const float* w_row = w_base + kh * KW;
        for (int kw = 0; kw < KW; ++kw) {
            const int iw = ow * strideW - padW + kw * dilW;
            if (iw >= 0 && iw < W) {
                acc = fmaf(__ldg(in_row + iw), __ldg(w_row + kw), acc);
            }
        }
    }
    output[idx] = acc;
}

} // namespace

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 and 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 and weight must be 4-D");

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

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);

    const int outH = (int)((H + 2 * pad_h - dil_h * (KH - 1) - 1) / stride_h + 1);
    const int outW = (int)((W + 2 * pad_w - dil_w * (KW - 1) - 1) / stride_w + 1);
    TORCH_CHECK(outH > 0 && outW > 0, "output size must be positive");

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

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

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

    if (stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1 &&
        KH == 3 && KW == 7 && (int64_t)N * (int64_t)C <= 65535) {
        constexpr int ROWS = 2;
        dim3 block(TILE_W, TILE_H);
        dim3 grid((outW + TILE_W - 1) / TILE_W,
                  (outH + TILE_H * ROWS - 1) / (TILE_H * ROWS),
                  N * C);
        depthwise_conv2d_s1d1_kernel<3, 7, ROWS><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
            C, H, W, outH, outW, (int)pad_h, (int)pad_w, has_bias);
    } else {
        const long total = (long)N * (long)C * (long)outH * (long)outW;
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        depthwise_conv2d_generic_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
            C, H, W, outH, outW, KH, KW,
            (int)stride_h, (int)stride_w, (int)pad_h, (int)pad_w,
            (int)dil_h, (int)dil_w, has_bias, total);
    }

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));
    return out;
}
"""

depthwise_conv2d_cpp_source = (
    "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);"
)

# Compile the inline CUDA extension.
depthwise_conv2d_ext = load_inline(
    name="depthwise_conv2d_asym_ext",
    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 kernel, accelerated with a custom
    CUDA kernel (shared-memory tiled, fully unrolled fast path for the 3x7 /
    stride-1 / dilation-1 configuration, plus a generic 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

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