```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_depthwise_conv2d_cpp_src = (
    "#include <torch/extension.h>\n"
    "torch::Tensor depthwise_conv2d_forward("
    "const torch::Tensor& input, "
    "const torch::Tensor& weight, "
    "const c10::optional<torch::Tensor>& bias, "
    "const int64_t stride_h, const int64_t stride_w, "
    "const int64_t pad_h, const int64_t pad_w, "
    "const int64_t dil_h, const int64_t dil_w);"
)

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

#define TILE_W 32
#define TILE_H 16

// ---------------------------------------------------------------------------
// Vectorized shared-memory load helpers (C++11-compatible template structs).
// LdsA: pointer must be 16B aligned. LdsB: pointer must be 8B aligned.
// ---------------------------------------------------------------------------
template <int N>
struct LdsA {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        *reinterpret_cast<float4*>(v) = *reinterpret_cast<const float4*>(p);
        LdsA<N - 4>::run(p + 4, v + 4);
    }
};
template <>
struct LdsA<0> {
    static __device__ __forceinline__ void run(const float*, float*) {}
};
template <>
struct LdsA<2> {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        *reinterpret_cast<float2*>(v) = *reinterpret_cast<const float2*>(p);
    }
};

template <int N>
struct LdsB {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        *reinterpret_cast<float2*>(v) = *reinterpret_cast<const float2*>(p);
        LdsB<N - 2>::run(p + 2, v + 2);
    }
};
template <>
struct LdsB<0> {
    static __device__ __forceinline__ void run(const float*, float*) {}
};
template <>
struct LdsB<1> {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        v[0] = p[0];
    }
};

// Window loader: pick widest safe vector width based on per-thread output
// block width OW (OW==4 => column offset is a multiple of 4 floats => 16B).
template <int OW, int W>
struct WinLoader {
    static __device__ __forceinline__ void run(const float* p, float* v) {
        LdsB<W>::run(p, v);
    }
};
template <int W>
struct WinLoader<4, W> {
    static __device__ __forceinline__ void run(const float* p, float* v) {
        LdsA<W>::run(p, v);
    }
};

// Vectorized global-store helper; p must be N*4-byte aligned.
template <int N>
struct StgV {
    static __device__ __forceinline__ void run(float* __restrict__ p, const float* __restrict__ v) {
        *reinterpret_cast<float4*>(p) = *reinterpret_cast<const float4*>(v);
        StgV<N - 4>::run(p + 4, v + 4);
    }
};
template <>
struct StgV<0> {
    static __device__ __forceinline__ void run(float*, const float*) {}
};
template <>
struct StgV<2> {
    static __device__ __forceinline__ void run(float* __restrict__ p, const float* __restrict__ v) {
        *reinterpret_cast<float2*>(p) = *reinterpret_cast<const float2*>(v);
    }
};

// ---------------------------------------------------------------------------
// Register-blocked kernel (float, stride == 1, dilation == 1).
// Each block computes a (32*OW) x (8*OH) output tile; each thread computes
// OW x OH outputs. Key differences vs a full-window design:
//   * Weights are loaded once per thread straight from global memory (uniform
//     broadcast) into registers -- no shared copy, no extra sync.
//   * The compute loop streams ONE input row segment through registers at a
//     time and updates every output row it contributes to. Live registers:
//     OH*OW accumulators + (KW+OW-1) segment + KH*KW weights, instead of the
//     full (KH+OH-1)*(KW+OW-1) window. The lower register pressure allows
//     __launch_bounds__(256, 4) => 4 resident blocks/SM on 64K-reg GPUs.
//   * Interior tiles are filled with vectorized float4 global loads.
//   * Stores use float4 when 16B-aligned, float2 pairs when 8B-aligned.
// Grid: (ceil(out_w/(32*OW)), ceil(out_h/(8*OH)), batch*channels)
// Block: (32, 8)
// ---------------------------------------------------------------------------
template <int KH, int KW, int OW, int OH>
__global__ __launch_bounds__(256, 4) void depthwise_conv2d_reg_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 pad_h, const int pad_w,
    const bool has_bias)
{
    constexpr int TX = 32;
    constexpr int TY = 8;
    constexpr int TILE_OW = TX * OW;
    constexpr int TILE_OH = TY * OH;
    constexpr int TILE_IW = TILE_OW + KW - 1;
    constexpr int TILE_IH = TILE_OH + KH - 1;
    constexpr int TILE_IWP = (TILE_IW + 3) & ~3;   // pad rows to 16B multiples
    constexpr int W = KW + OW - 1;                 // segment registers per row
    constexpr int VH = KH + OH - 1;                // segment register rows
    constexpr int NT = TX * TY;
    constexpr int NLOAD = TILE_IH * TILE_IW;
    constexpr int NVEC = TILE_IH * (TILE_IWP / 4);

    __shared__ __align__(16) float s_in[TILE_IH][TILE_IWP];

    const long bc = blockIdx.z;
    const int c = (int)(bc % channels);
    const int tid = threadIdx.y * TX + threadIdx.x;

    // Weights -> registers (uniform broadcast loads, latency hidden behind
    // the tile fill below).
    float w[KH * KW];
    {
        const float* wp = weight + (long)c * (KH * KW);
        #pragma unroll
        for (int i = 0; i < KH * KW; ++i) w[i] = __ldg(wp + i);
    }

    const int ow0 = blockIdx.x * TILE_OW;
    const int oh0 = blockIdx.y * TILE_OH;
    const int ih0 = oh0 - pad_h;
    const int iw0 = ow0 - pad_w;

    const float* in_base = input + bc * (long)in_h * in_w;

    if (ih0 >= 0 && iw0 >= 0 && ih0 + TILE_IH <= in_h && iw0 + TILE_IW <= in_w) {
        if (iw0 + TILE_IWP <= in_w && (in_w & 3) == 0 && (iw0 & 3) == 0) {
            // Fully interior and 16B aligned: float4 fill (also covers the
            // padded columns, which are real in-bounds data here).
            #pragma unroll
            for (int k = 0; k < (NVEC + NT - 1) / NT; ++k) {
                const int i = tid + k * NT;
                if (i < NVEC) {
                    const int r = i / (TILE_IWP / 4);
                    const int cq = i - r * (TILE_IWP / 4);
                    *reinterpret_cast<float4*>(&s_in[r][cq * 4]) =
                        *reinterpret_cast<const float4*>(
                            in_base + (long)(ih0 + r) * in_w + (iw0 + cq * 4));
                }
            }
        } else {
            #pragma unroll
            for (int k = 0; k < (NLOAD + NT - 1) / NT; ++k) {
                const int i = tid + k * NT;
                if (i < NLOAD) {
                    const int r = i / TILE_IW;
                    const int col = i - r * TILE_IW;
                    s_in[r][col] = in_base[(long)(ih0 + r) * in_w + (iw0 + col)];
                }
            }
        }
    } else {
        #pragma unroll
        for (int k = 0; k < (NLOAD + NT - 1) / NT; ++k) {
            const int i = tid + k * NT;
            if (i < NLOAD) {
                const int r = i / TILE_IW;
                const int col = i - r * TILE_IW;
                const int ih = ih0 + r;
                const int iw = iw0 + col;
                float v = 0.f;
                if (ih >= 0 && ih < in_h && iw >= 0 && iw < in_w) {
                    v = in_base[(long)ih * in_w + iw];
                }
                s_in[r][col] = v;
            }
        }
    }
    __syncthreads();

    const int rb = threadIdx.y * OH;
    const int cb = threadIdx.x * OW;

    const float b = has_bias ? __ldg(bias + c) : 0.f;
    float acc[OH][OW];
    #pragma unroll
    for (int dy = 0; dy < OH; ++dy) {
        #pragma unroll
        for (int dx = 0; dx < OW; ++dx) acc[dy][dx] = b;
    }

    // Row-sliding window accumulation.
    #pragma unroll
    for (int r = 0; r < VH; ++r) {
        float seg[W];
        WinLoader<OW, W>::run(&s_in[rb + r][cb], seg);
        #pragma unroll
        for (int dy = 0; dy < OH; ++dy) {
            const int i = r - dy;
            if (i >= 0 && i < KH) {
                #pragma unroll
                for (int j = 0; j < KW; ++j) {
                    const float wv = w[i * KW + j];
                    #pragma unroll
                    for (int dx = 0; dx < OW; ++dx) {
                        acc[dy][dx] += seg[j + dx] * wv;
                    }
                }
            }
        }
    }

    const int ow = ow0 + cb;
    const int oh = oh0 + rb;
    float* out_base = output + bc * (long)out_h * out_w;
    const bool full_tile_x = (ow + OW <= out_w);

    #pragma unroll
    for (int dy = 0; dy < OH; ++dy) {
        const int y = oh + dy;
        if (y >= out_h) continue;
        float* out_row = out_base + (long)y * out_w + ow;
        if (full_tile_x) {
            const uintptr_t a = reinterpret_cast<uintptr_t>(out_row);
            if ((a & (uintptr_t)15) == 0) {
                StgV<OW>::run(out_row, acc[dy]);
            } else if ((a & (uintptr_t)7) == 0) {
                #pragma unroll
                for (int dx = 0; dx < OW; dx += 2) {
                    *reinterpret_cast<float2*>(out_row + dx) =
                        *reinterpret_cast<const float2*>(&acc[dy][dx]);
                }
            } else {
                #pragma unroll
                for (int dx = 0; dx < OW; ++dx) out_row[dx] = acc[dy][dx];
            }
        } else {
            #pragma unroll
            for (int dx = 0; dx < OW; ++dx) {
                if (ow + dx < out_w) out_row[dx] = acc[dy][dx];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Shared-memory tile kernel: templated on kernel size (fully unrolled),
// valid for stride == 1 and dilation == 1 (any padding). One output pixel
// per thread. Used as fallback for sizes not covered by the reg kernel.
// ---------------------------------------------------------------------------
template <int KH, int KW, typename scalar_t>
__global__ void depthwise_conv2d_fast_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int pad_h, const int pad_w,
    const bool has_bias)
{
    constexpr int TILE_IN_H = TILE_H + KH - 1;
    constexpr int TILE_IN_W = TILE_W + KW - 1;

    __shared__ scalar_t s_w[KH * KW];
    __shared__ scalar_t s_in[TILE_IN_H][TILE_IN_W];

    const long bc = blockIdx.x;
    const int c = (int)(bc % channels);
    const int tid = threadIdx.y * TILE_W + threadIdx.x;
    constexpr int nthreads = TILE_W * TILE_H;

    if (tid < KH * KW) {
        s_w[tid] = weight[(long)c * KH * KW + tid];
    }

    const int ow0 = blockIdx.z * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int ih0 = oh0 - pad_h;
    const int iw0 = ow0 - pad_w;

    const scalar_t* in_base = input + bc * (long)in_h * in_w;

    #pragma unroll 2
    for (int i = tid; i < TILE_IN_H * TILE_IN_W; i += nthreads) {
        const int r = i / TILE_IN_W;
        const int col = i - r * TILE_IN_W;
        const int ih = ih0 + r;
        const int iw = iw0 + col;
        scalar_t v = scalar_t(0);
        if (ih >= 0 && ih < in_h && iw >= 0 && iw < in_w) {
            v = in_base[(long)ih * in_w + iw];
        }
        s_in[r][col] = v;
    }
    __syncthreads();

    const int ow = ow0 + threadIdx.x;
    const int oh = oh0 + threadIdx.y;
    if (ow < out_w && oh < out_h) {
        scalar_t sum = has_bias ? bias[c] : scalar_t(0);
        #pragma unroll
        for (int i = 0; i < KH; ++i) {
            #pragma unroll
            for (int j = 0; j < KW; ++j) {
                sum += s_in[threadIdx.y + i][threadIdx.x + j] * s_w[i * KW + j];
            }
        }
        output[(bc * (long)out_h + oh) * out_w + ow] = sum;
    }
}

// ---------------------------------------------------------------------------
// Generic kernel: arbitrary kernel size / stride / dilation / padding, using
// dynamically-sized shared memory tile.
// ---------------------------------------------------------------------------
template <typename scalar_t>
__global__ void depthwise_conv2d_generic_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __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 bool has_bias)
{
    extern __shared__ char smem_raw[];
    const int tile_in_w = (TILE_W - 1) * stride_w + (kw - 1) * dil_w + 1;
    const int tile_in_h = (TILE_H - 1) * stride_h + (kh - 1) * dil_h + 1;
    scalar_t* s_w = reinterpret_cast<scalar_t*>(smem_raw);
    scalar_t* s_in = s_w + kh * kw;

    const long bc = blockIdx.x;
    const int c = (int)(bc % channels);
    const int tid = threadIdx.y * TILE_W + threadIdx.x;
    const int nthreads = TILE_W * TILE_H;

    for (int i = tid; i < kh * kw; i += nthreads) {
        s_w[i] = weight[(long)c * kh * kw + i];
    }

    const int ow0 = blockIdx.z * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int ih0 = oh0 * stride_h - pad_h;
    const int iw0 = ow0 * stride_w - pad_w;

    const scalar_t* in_base = input + bc * (long)in_h * in_w;

    for (int i = tid; i < tile_in_h * tile_in_w; i += nthreads) {
        const int r = i / tile_in_w;
        const int col = i - r * tile_in_w;
        const int ih = ih0 + r;
        const int iw = iw0 + col;
        scalar_t v = scalar_t(0);
        if (ih >= 0 && ih < in_h && iw >= 0 && iw < in_w) {
            v = in_base[(long)ih * in_w + iw];
        }
        s_in[r * tile_in_w + col] = v;
    }
    __syncthreads();

    const int ow = ow0 + threadIdx.x;
    const int oh = oh0 + threadIdx.y;
    if (ow < out_w && oh < out_h) {
        scalar_t sum = has_bias ? bias[c] : scalar_t(0);
        const int r0 = threadIdx.y * stride_h;
        const int c0 = threadIdx.x * stride_w;
        for (int i = 0; i < kh; ++i) {
            const int r = r0 + i * dil_h;
            for (int j = 0; j < kw; ++j) {
                sum += s_in[r * tile_in_w + (c0 + j * dil_w)] * s_w[i * kw + j];
            }
        }
        output[(bc * (long)out_h + oh) * out_w + ow] = sum;
    }
}

// ---------------------------------------------------------------------------
// Naive fallback (used only when the shared-memory tile would be too large).
// ---------------------------------------------------------------------------
template <typename scalar_t>
__global__ void depthwise_conv2d_naive_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __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 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 % out_w);
    const int oh = (int)((idx / out_w) % out_h);
    const long bc = idx / ((long)out_h * out_w);
    const int c = (int)(bc % channels);

    const scalar_t* in_base = input + bc * (long)in_h * in_w;
    const scalar_t* w_base = weight + (long)c * kh * kw;

    scalar_t sum = has_bias ? bias[c] : scalar_t(0);
    for (int i = 0; i < kh; ++i) {
        const int ih = oh * stride_h - pad_h + i * dil_h;
        if (ih < 0 || ih >= in_h) continue;
        for (int j = 0; j < kw; ++j) {
            const int iw = ow * stride_w - pad_w + j * dil_w;
            if (iw < 0 || iw >= in_w) continue;
            sum += in_base[(long)ih * in_w + iw] * w_base[i * kw + j];
        }
    }
    output[idx] = sum;
}

// ---------------------------------------------------------------------------
// Register-blocked float path. Returns true when the kernel was launched.
// ---------------------------------------------------------------------------
static bool try_launch_reg_float(
    const float* input,
    const float* weight,
    const float* bias,
    float* output,
    const int batch, 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 pad_h, const int pad_w,
    cudaStream_t stream)
{
    const long bc_total = (long)batch * channels;
    if (bc_total > 65535) return false;
    const bool has_bias = (bias != nullptr);
    const dim3 rblock(32, 8);

#define TRY_REG(KHV, KWV, OWV, OHV)                                                \
    if (kh == KHV && kw == KWV) {                                                  \
        const dim3 rgrid(                                                          \
            (unsigned)((out_w + 32 * OWV - 1) / (32 * OWV)),                       \
            (unsigned)((out_h + 8 * OHV - 1) / (8 * OHV)),                         \
            (unsigned)bc_total);                                                   \
        if (rgrid.y > 65535) return false;                                         \
        depthwise_conv2d_reg_kernel<KHV, KWV, OWV, OHV>                            \
            <<<rgrid, rblock, 0, stream>>>(                                        \
                input, weight, bias, output, channels, in_h, in_w,                 \
                out_h, out_w, pad_h, pad_w, has_bias);                             \
        C10_CUDA_KERNEL_LAUNCH_CHECK();                                            \
        return true;                                                               \
    }

    TRY_REG(1, 3, 4, 4)
    TRY_REG(3, 1, 4, 4)
    TRY_REG(1, 5, 4, 4)
    TRY_REG(5, 1, 4, 4)
    TRY_REG(1, 7, 4, 4)
    TRY_REG(7, 1, 4, 4)
    TRY_REG(3, 3, 4, 4)
    TRY_REG(3, 5, 4, 4)
    TRY_REG(5, 3, 4, 4)
    TRY_REG(3, 7, 4, 4)
    TRY_REG(7, 3, 4, 4)
    TRY_REG(5, 5, 4, 2)
#undef TRY_REG
    return false;
}

// ---------------------------------------------------------------------------
// Launch helper
// ---------------------------------------------------------------------------
template <typename scalar_t>
static void launch_depthwise_conv2d(
    const scalar_t* input,
    const scalar_t* weight,
    const scalar_t* bias,
    scalar_t* output,
    const int batch, 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,
    cudaStream_t stream)
{
    const dim3 block(TILE_W, TILE_H);
    const dim3 grid(
        (unsigned int)(batch * channels),
        (unsigned int)((out_h + TILE_H - 1) / TILE_H),
        (unsigned int)((out_w + TILE_W - 1) / TILE_W));
    const bool has_bias = (bias != nullptr);

#define TRY_FAST(KHV, KWV)                                                          \
    if (kh == KHV && kw == KWV) {                                                   \
        depthwise_conv2d_fast_kernel<KHV, KWV, scalar_t>                            \
            <<<grid, block, 0, stream>>>(                                           \
                input, weight, bias, output, channels, in_h, in_w,                  \
                out_h, out_w, pad_h, pad_w, has_bias);                              \
        C10_CUDA_KERNEL_LAUNCH_CHECK();                                             \
        return;                                                                     \
    }

    if (stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1) {
        TRY_FAST(1, 1)
        TRY_FAST(1, 3)
        TRY_FAST(3, 1)
        TRY_FAST(1, 5)
        TRY_FAST(5, 1)
        TRY_FAST(1, 7)
        TRY_FAST(7, 1)
        TRY_FAST(3, 3)
        TRY_FAST(3, 5)
        TRY_FAST(5, 3)
        TRY_FAST(3, 7)
        TRY_FAST(7, 3)
        TRY_FAST(5, 5)
        TRY_FAST(5, 7)
        TRY_FAST(7, 5)
        TRY_FAST(7, 7)
    }
#undef TRY_FAST

    const int tile_in_w = (TILE_W - 1) * stride_w + (kw - 1) * dil_w + 1;
    const int tile_in_h = (TILE_H - 1) * stride_h + (kh - 1) * dil_h + 1;
    const size_t shmem = (size_t)(kh * kw + tile_in_h * tile_in_w) * sizeof(scalar_t);
    if (shmem <= 48 * 1024) {
        depthwise_conv2d_generic_kernel<scalar_t><<<grid, block, shmem, stream>>>(
            input, weight, bias, output, channels, in_h, in_w, out_h, out_w,
            kh, kw, stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, has_bias);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        const long total = (long)batch * channels * out_h * out_w;
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        depthwise_conv2d_naive_kernel<scalar_t>
            <<<(unsigned int)blocks, threads, 0, stream>>>(
                input, weight, bias, output, channels, in_h, in_w, out_h, out_w,
                kh, kw, stride_h, stride_w, pad_h, pad_w, dil_h, dil_w,
                has_bias, total);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
}

// ---------------------------------------------------------------------------
// Exposed op
// ---------------------------------------------------------------------------
torch::Tensor depthwise_conv2d_forward(
    const torch::Tensor& input,
    const torch::Tensor& weight,
    const c10::optional<torch::Tensor>& bias,
    const int64_t stride_h, const int64_t stride_w,
    const int64_t pad_h, const int64_t pad_w,
    const int64_t dil_h, const int64_t dil_w)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
    TORCH_CHECK(weight.dim() == 4, "weight must be 4D (C, 1, KH, KW)");

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

    auto input_c = input.contiguous();
    auto weight_c = weight.contiguous();

    const int batch = (int)input_c.size(0);
    const int channels = (int)input_c.size(1);
    const int in_h = (int)input_c.size(2);
    const int in_w = (int)input_c.size(3);
    const int kh = (int)weight_c.size(2);
    const int kw = (int)weight_c.size(3);
    TORCH_CHECK(weight_c.size(0) == channels && weight_c.size(1) == 1,
                "expected depthwise weight of shape (C, 1, KH, KW)");

    const int out_h = (int)((in_h + 2 * pad_h - dil_h * (kh - 1) - 1) / stride_h + 1);
    const int out_w = (int)((in_w + 2 * pad_w - dil_w * (kw - 1) - 1) / stride_w + 1);
    TORCH_CHECK(out_h > 0 && out_w > 0, "invalid output size");

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

    const bool has_bias = bias.has_value() && bias->defined();
    torch::Tensor bias_c;
    if (has_bias) {
        bias_c = bias->contiguous();
    }

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

    // Highest-performance path: register-blocked float kernel.
    if (input_c.scalar_type() == at::kFloat &&
        stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1) {
        if (try_launch_reg_float(
                input_c.data_ptr<float>(),
                weight_c.data_ptr<float>(),
                has_bias ? bias_c.data_ptr<float>() : nullptr,
                output.data_ptr<float>(),
                batch, channels, in_h, in_w, out_h, out_w,
                kh, kw, (int)pad_h, (int)pad_w, stream)) {
            return output;
        }
    }

    AT_DISPATCH_FLOATING_TYPES(input_c.scalar_type(), "depthwise_conv2d_forward", [&] {
        launch_depthwise_conv2d<scalar_t>(
            input_c.data_ptr<scalar_t>(),
            weight_c.data_ptr<scalar_t>(),
            has_bias ? bias_c.data_ptr<scalar_t>() : nullptr,
            output.data_ptr<scalar_t>(),
            batch, channels, in_h, in_w, out_h, out_w,
            kh, kw,
            (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w,
            (int)dil_h, (int)dil_w,
            stream);
    });

    return output;
}
"""

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


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution with asymmetric input and asymmetric kernel,
    accelerated with a custom register-blocked row-sliding-window 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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and (x.dtype == torch.float32 or x.dtype == torch.float64):
            return _depthwise_conv2d_ext.depthwise_conv2d_forward(
                x,
                self.conv2d.weight,
                self.conv2d.bias,
                self.conv2d.stride[0], self.conv2d.stride[1],
                self.conv2d.padding[0], self.conv2d.padding[1],
                self.conv2d.dilation[0], self.conv2d.dilation[1],
            )
        return F.conv2d(
            x,
            self.conv2d.weight,
            self.conv2d.bias,
            self.conv2d.stride,
            self.conv2d.padding,
            self.conv2d.dilation,
            self.conv2d.groups,
        )
```