```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 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);
    }
};

// ===========================================================================
// Rolling-window kernel (float, stride == 1, dilation == 1).
// Each thread owns OW=4 consecutive output columns of ONE (batch*channel)
// plane and sweeps down all output rows, keeping a KH x (KW+OW-1) sliding
// window in registers. Every input element is read from global memory at most
// once per thread-neighborhood (halo reuse comes from L1), so DRAM traffic is
// essentially input (read once) + output (written once): the theoretical
// minimum. No shared memory, no __syncthreads. Grid: (ceil(BC/4),
// ceil(out_w/256)); block: 256 threads = 4 planes x 64 column-threads.
// ===========================================================================
template <int KH, int KW, bool VEC4>
__device__ __forceinline__ void roll_load_row(
    const float* __restrict__ in_base,
    const int ih, const int in_h, const int in_w,
    const int iw0t, const bool int_x,
    float* __restrict__ v)
{
    constexpr int W = KW + 3;  // KW + OW - 1 with OW = 4
    if ((unsigned)ih < (unsigned)in_h) {
        if (int_x) {
            const float* p = in_base + (long)ih * in_w + iw0t;
            if (VEC4) {
                LdsA<W>::run(p, v);
            } else {
                LdsB<W>::run(p, v);
            }
        } else {
            const float* rowp = in_base + (long)ih * in_w;
            #pragma unroll
            for (int j = 0; j < W; ++j) {
                const int col = iw0t + j;
                v[j] = ((unsigned)col < (unsigned)in_w) ? __ldg(rowp + col) : 0.f;
            }
        }
    } else {
        #pragma unroll
        for (int j = 0; j < W; ++j) v[j] = 0.f;
    }
}

template <int KH, int KW, bool VEC4>
__global__ __launch_bounds__(256, 3) void depthwise_conv2d_roll_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 int bc_total,
    const bool has_bias)
{
    constexpr int OW = 4;
    constexpr int W = KW + OW - 1;
    constexpr int WT = 64;          // column-threads per plane
    constexpr int G = 256 / WT;     // planes per block

    const int g = threadIdx.x / WT;
    const int tx = threadIdx.x - g * WT;
    const int bc = blockIdx.x * G + g;
    if (bc >= bc_total) return;
    const int c = bc % channels;

    const int ox = blockIdx.y * (WT * OW) + tx * OW;
    if (ox >= out_w) return;
    const int iw0t = ox - pad_w;
    const bool int_x = (iw0t >= 0) && (iw0t + W <= in_w);
    const bool full_out = (ox + OW <= out_w);

    const float* in_base = input + (long)bc * in_h * in_w;
    float* out_base = output + (long)bc * out_h * out_w;

    // Cache this channel's weights in registers (uniform across the warp).
    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 float bv = has_bias ? __ldg(bias + c) : 0.f;

    // Sliding register window: KH rows x W columns.
    float wr[KH][W];
    #pragma unroll
    for (int i = 0; i < KH; ++i) {
        roll_load_row<KH, KW, VEC4>(in_base, i - pad_h, in_h, in_w, iw0t, int_x, wr[i]);
    }

    unsigned par = 0u;  // (oh*out_w + ox) & 3 ; ox % 4 == 0

    for (int oh = 0; oh < out_h; ++oh) {
        float acc[OW];
        #pragma unroll
        for (int dx = 0; dx < OW; ++dx) {
            float s = bv;
            #pragma unroll
            for (int i = 0; i < KH; ++i) {
                #pragma unroll
                for (int j = 0; j < KW; ++j) {
                    s += wr[i][j + dx] * w[i * KW + j];
                }
            }
            acc[dx] = s;
        }

        float* op = out_base + (long)oh * out_w + ox;
        if (full_out) {
            if (par == 0u) {
                __stcs(reinterpret_cast<float4*>(op),
                       make_float4(acc[0], acc[1], acc[2], acc[3]));
            } else if ((par & 1u) == 0u) {
                __stcs(reinterpret_cast<float2*>(op), make_float2(acc[0], acc[1]));
                __stcs(reinterpret_cast<float2*>(op + 2), make_float2(acc[2], acc[3]));
            } else {
                __stcs(op, acc[0]);
                __stcs(op + 1, acc[1]);
                __stcs(op + 2, acc[2]);
                __stcs(op + 3, acc[3]);
            }
        } else {
            #pragma unroll
            for (int dx = 0; dx < OW; ++dx) {
                if (ox + dx < out_w) __stcs(op + dx, acc[dx]);
            }
        }

        // Rotate window up and prefetch the next input row (distance 1).
        #pragma unroll
        for (int i = 0; i < KH - 1; ++i) {
            #pragma unroll
            for (int j = 0; j < W; ++j) {
                wr[i][j] = wr[i + 1][j];
            }
        }
        if (oh + 1 < out_h) {
            roll_load_row<KH, KW, VEC4>(
                in_base, oh + 1 - pad_h + KH - 1, in_h, in_w, iw0t, int_x, wr[KH - 1]);
        }
        par = (par + (unsigned)out_w) & 3u;
    }
}

// ---------------------------------------------------------------------------
// 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. Shared-memory tiled fallback for shapes the rolling
// kernel does not cover.
// ---------------------------------------------------------------------------
template <int KH, int KW, int OW, int OH>
__global__ __launch_bounds__(256, 3) 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;
    constexpr int W = KW + OW - 1;
    constexpr int VH = KH + OH - 1;
    constexpr int NT = TX * TY;
    constexpr int NLOAD = TILE_IH * TILE_IW;

    __shared__ float s_w[KH * KW];
    __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;

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

    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) {
        #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();

    float w[KH * KW];
    #pragma unroll
    for (int i = 0; i < KH * KW; ++i) w[i] = s_w[i];

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

    float v[VH][W];
    #pragma unroll
    for (int i = 0; i < VH; ++i) {
        WinLoader<OW, W>::run(&s_in[rb + i][cb], v[i]);
    }

    const float b = has_bias ? bias[c] : 0.f;
    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 outv[OW];
        #pragma unroll
        for (int dx = 0; dx < OW; ++dx) {
            float acc[KH];
            #pragma unroll
            for (int i = 0; i < KH; ++i) {
                acc[i] = 0.f;
                #pragma unroll
                for (int j = 0; j < KW; ++j) {
                    acc[i] += v[i + dy][j + dx] * w[i * KW + j];
                }
            }
            float sum = b;
            #pragma unroll
            for (int i = 0; i < KH; ++i) sum += acc[i];
            outv[dx] = sum;
        }

        float* out_row = out_base + (long)y * out_w + ow;
        if (full_tile_x &&
            ((reinterpret_cast<uintptr_t>(out_row) & (uintptr_t)(OW * 4 - 1)) == 0)) {
            StgV<OW>::run(out_row, outv);
        } else {
            #pragma unroll
            for (int dx = 0; dx < OW; ++dx) {
                if (ow + dx < out_w) out_row[dx] = outv[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;
}

// ---------------------------------------------------------------------------
// Rolling-window float path. Returns true when the kernel was launched.
// ---------------------------------------------------------------------------
static bool try_launch_roll_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 < 256) return false;  // not enough blocks to fill the GPU
    const bool vec4 = ((in_w & 3) == 0) && ((pad_w & 3) == 0);
    const bool vec2 = ((in_w & 1) == 0) && ((pad_w & 1) == 0);
    if (!vec4 && !vec2) return false;
    const bool has_bias = (bias != nullptr);
    const dim3 block(256);
    const dim3 grid(
        (unsigned)((bc_total + 3) / 4),
        (unsigned)((out_w + 255) / 256));

#define TRY_ROLL(KHV, KWV)                                                        \
    if (kh == KHV && kw == KWV) {                                                 \
        if (vec4) {                                                               \
            depthwise_conv2d_roll_kernel<KHV, KWV, true>                          \
                <<<grid, block, 0, stream>>>(                                     \
                    input, weight, bias, output, channels, in_h, in_w,            \
                    out_h, out_w, pad_h, pad_w, (int)bc_total, has_bias);         \
        } else {                                                                  \
            depthwise_conv2d_roll_kernel<KHV, KWV, false>                         \
                <<<grid, block, 0, stream>>>(                                     \
                    input, weight, bias, output, channels, in_h, in_w,            \
                    out_h, out_w, pad_h, pad_w, (int)bc_total, has_bias);         \
        }                                                                         \
        C10_CUDA_KERNEL_LAUNCH_CHECK();                                           \
        return true;                                                              \
    }

    TRY_ROLL(1, 1)
    TRY_ROLL(1, 3)
    TRY_ROLL(3, 1)
    TRY_ROLL(1, 5)
    TRY_ROLL(5, 1)
    TRY_ROLL(1, 7)
    TRY_ROLL(7, 1)
    TRY_ROLL(3, 3)
    TRY_ROLL(3, 5)
    TRY_ROLL(5, 3)
    TRY_ROLL(3, 7)
    TRY_ROLL(7, 3)
    TRY_ROLL(5, 5)
    TRY_ROLL(5, 7)
    TRY_ROLL(7, 5)
    TRY_ROLL(7, 7)
#undef TRY_ROLL
    return false;
}

// ---------------------------------------------------------------------------
// Register-blocked float tiled path. Returns true when 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, 2)
    TRY_REG(3, 1, 4, 2)
    TRY_REG(1, 5, 4, 2)
    TRY_REG(5, 1, 4, 2)
    TRY_REG(1, 7, 4, 2)
    TRY_REG(7, 1, 4, 2)
    TRY_REG(3, 3, 4, 2)
    TRY_REG(3, 5, 4, 2)
    TRY_REG(5, 3, 4, 2)
    TRY_REG(3, 7, 4, 2)
    TRY_REG(7, 3, 2, 2)
    TRY_REG(5, 5, 2, 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