```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 <cuda_runtime.h>
#include <cstdint>

#define BLOCK_X 64
#define BLOCK_Y 8
#define UNROLL_W 4

// Tiled depthwise 2D convolution kernel (NCHW, groups == channels).
// Each block computes a (BLOCK_Y) x (BLOCK_X * UNROLL_W) output tile for one
// (batch, channel) pair. The required input halo region is staged in shared
// memory. Each thread computes UNROLL_W outputs strided by BLOCK_X so that
// shared-memory reads and global stores are fully coalesced / conflict-free.
template <int KH_, int KW_>
__global__ void depthwise_conv2d_tiled_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    const int C, const int H, const int W,
    const int HO, const int WO,
    const int KHr, const int KWr,
    const int SH, const int SW,
    const int PH, const int PW,
    const int DH, const int DW,
    const int tilesW, const int tilesH,
    const int smemH, const int smemW,
    const bool has_bias)
{
    const int KH = (KH_ > 0) ? KH_ : KHr;
    const int KW = (KW_ > 0) ? KW_ : KWr;

    extern __shared__ float smem[];

    const int linear = blockIdx.x;
    const int tw = linear % tilesW;
    const int tmp = linear / tilesW;
    const int th = tmp % tilesH;
    const int bc = tmp / tilesH;
    const int n = bc / C;
    const int c = bc % C;

    const float* __restrict__ xp = x + ((int64_t)n * C + c) * (int64_t)H * W;
    const float* __restrict__ wp = w + (int64_t)c * KH * KW;
    float* __restrict__ yp = y + ((int64_t)n * C + c) * (int64_t)HO * WO;

    const int tileW = BLOCK_X * UNROLL_W;
    const int out_x0 = tw * tileW;
    const int out_y0 = th * BLOCK_Y;
    const int in_x0 = out_x0 * SW - PW;
    const int in_y0 = out_y0 * SH - PH;

    // Stage the input tile (with halo) into shared memory.
    const int tid = threadIdx.y * BLOCK_X + threadIdx.x;
    const int total = smemH * smemW;
    for (int idx = tid; idx < total; idx += BLOCK_X * BLOCK_Y) {
        const int r = idx / smemW;
        const int col = idx - r * smemW;
        const int gx = in_x0 + col;
        const int gy = in_y0 + r;
        float v = 0.0f;
        if (gx >= 0 && gx < W && gy >= 0 && gy < H) {
            v = xp[(int64_t)gy * W + gx];
        }
        smem[idx] = v;
    }
    __syncthreads();

    const int oy = out_y0 + threadIdx.y;
    if (oy >= HO) return;

    float acc[UNROLL_W];
    const float biasv = has_bias ? b[c] : 0.0f;
    #pragma unroll
    for (int u = 0; u < UNROLL_W; ++u) acc[u] = biasv;

    const int sy = threadIdx.y * SH;
    const int sx = threadIdx.x * SW;

    #pragma unroll
    for (int i = 0; i < KH; ++i) {
        const int row = (sy + i * DH) * smemW;
        #pragma unroll
        for (int j = 0; j < KW; ++j) {
            const float wv = wp[i * KW + j];
            const int base = row + sx + j * DW;
            #pragma unroll
            for (int u = 0; u < UNROLL_W; ++u) {
                acc[u] = fmaf(smem[base + u * (BLOCK_X * SW)], wv, acc[u]);
            }
        }
    }

    #pragma unroll
    for (int u = 0; u < UNROLL_W; ++u) {
        const int xcol = out_x0 + threadIdx.x + u * BLOCK_X;
        if (xcol < WO) {
            yp[(int64_t)oy * WO + xcol] = acc[u];
        }
    }
}

// Simple fallback for configurations whose shared-memory tile would be too large.
__global__ void depthwise_conv2d_naive_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    const int64_t total,
    const int C, const int H, const int W,
    const int HO, const int WO,
    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 int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int wo = (int)(idx % WO);
    const int ho = (int)((idx / WO) % HO);
    const int64_t nclin = idx / ((int64_t)HO * WO);
    const int c = (int)(nclin % C);
    const int n = (int)(nclin / C);

    const float* __restrict__ xp = x + ((int64_t)n * C + c) * (int64_t)H * W;
    const float* __restrict__ wp = w + (int64_t)c * KH * KW;

    float acc = has_bias ? b[c] : 0.0f;
    for (int i = 0; i < KH; ++i) {
        const int ih = ho * SH + i * DH - PH;
        if (ih < 0 || ih >= H) continue;
        for (int j = 0; j < KW; ++j) {
            const int iw = wo * SW + j * DW - PW;
            if (iw < 0 || iw >= W) continue;
            acc = fmaf(xp[(int64_t)ih * W + iw], wp[i * KW + j], acc);
        }
    }
    y[idx] = acc;
}

torch::Tensor depthwise_conv2d_cuda(
    torch::Tensor x,
    torch::Tensor w,
    c10::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(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.dim() == 4, "x must be a 4D NCHW tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only float32 is supported");

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int B = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int H = (int)xc.size(2);
    const int W = (int)xc.size(3);
    const int KH = (int)wc.size(2);
    const int KW = (int)wc.size(3);
    const int SH = (int)stride_h, SW = (int)stride_w;
    const int PH = (int)pad_h,   PW = (int)pad_w;
    const int DH = (int)dil_h,   DW = (int)dil_w;

    const int HO = (int)((H + 2 * PH - DH * (int64_t)(KH - 1) - 1) / SH + 1);
    const int WO = (int)((W + 2 * PW - DW * (int64_t)(KW - 1) - 1) / SW + 1);

    auto y = at::empty({B, C, HO, WO}, xc.options());
    if (HO <= 0 || WO <= 0) return y;

    torch::Tensor bc;
    const float* bptr = nullptr;
    bool has_bias = false;
    if (bias.has_value() && bias->defined()) {
        bc = bias->contiguous();
        bptr = bc.data_ptr<float>();
        has_bias = true;
    }

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

    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    float* yp = y.data_ptr<float>();

    const int tileW = BLOCK_X * UNROLL_W;
    const int smemW = (tileW - 1) * SW + (KW - 1) * DW + 1;
    const int smemH = (BLOCK_Y - 1) * SH + (KH - 1) * DH + 1;
    const size_t smem_bytes = (size_t)smemH * (size_t)smemW * sizeof(float);
    const int tilesW = (WO + tileW - 1) / tileW;
    const int tilesH = (HO + BLOCK_Y - 1) / BLOCK_Y;
    const int64_t total_blocks = (int64_t)B * C * tilesH * tilesW;

    if (smem_bytes <= 48 * 1024 && total_blocks < 2000000000LL) {
        const dim3 block(BLOCK_X, BLOCK_Y);
        const dim3 grid((unsigned int)total_blocks);

        #define LAUNCH_TILED(KHv, KWv) \
            depthwise_conv2d_tiled_kernel<KHv, KWv><<<grid, block, smem_bytes, stream>>>( \
                xp, wp, bptr, yp, C, H, W, HO, WO, KH, KW, \
                SH, SW, PH, PW, DH, DW, tilesW, tilesH, smemH, smemW, 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                         { LAUNCH_TILED(0, 0); }

        #undef LAUNCH_TILED
    } else {
        const int64_t total = (int64_t)B * C * HO * WO;
        const int tpb = 256;
        const int64_t blocks = (total + tpb - 1) / tpb;
        depthwise_conv2d_naive_kernel<<<(unsigned int)blocks, tpb, 0, stream>>>(
            xp, wp, bptr, yp, total, C, H, W, HO, WO, KH, KW,
            SH, SW, PH, PW, DH, DW, has_bias);
    }

    return y;
}
"""

_depthwise_conv2d_cpp_src = (
    "torch::Tensor depthwise_conv2d_cuda(torch::Tensor x, torch::Tensor w, "
    "c10::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_ext",
    cpp_sources=_depthwise_conv2d_cpp_src,
    cuda_sources=_depthwise_conv2d_cuda_src,
    functions=["depthwise_conv2d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class _DepthwiseConv2dFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, stride, padding, dilation, groups):
        ctx.stride = stride
        ctx.padding = padding
        ctx.dilation = dilation
        ctx.groups = groups
        ctx.save_for_backward(x, weight, bias)
        return _depthwise_conv2d_ext.depthwise_conv2d_cuda(
            x, weight, bias,
            stride[0], stride[1],
            padding[0], padding[1],
            dilation[0], dilation[1],
        )

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, bias = ctx.saved_tensors
        grad_output = grad_output.contiguous()
        grad_input = torch.nn.grad.conv2d_input(
            x.shape, weight, grad_output, ctx.stride, ctx.padding, ctx.dilation, ctx.groups
        )
        grad_weight = torch.nn.grad.conv2d_weight(
            x, weight.shape, grad_output, ctx.stride, ctx.padding, ctx.dilation, ctx.groups
        )
        grad_bias = grad_output.sum((0, 2, 3)) if bias is not None else None
        return grad_input, grad_weight, grad_bias, None, None, None, None


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution with asymmetric input and asymmetric kernel,
    accelerated by a custom shared-memory tiled 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 = (stride_h, stride_w)
        self._padding = (padding_h, padding_w)
        self._dilation = (dilation_h, dilation_w)
        self._groups = in_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32:
            return _DepthwiseConv2dFn.apply(
                x, self.conv2d.weight, self.conv2d.bias,
                self._stride, self._padding, self._dilation, self._groups,
            )
        return self.conv2d(x)
```