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

// Vectorized global-store helper; p must be 16B aligned, N a multiple of 4.
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);
    }
};

// FMA update of every output row that input row `r` contributes to.
template <int KH, int KW, int OW, int OH>
__device__ __forceinline__ void fma_seg_row(
    const float* __restrict__ seg,
    const float* __restrict__ w,
    float (* __restrict__ acc)[OW],
    const int r)
{
    #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] = fmaf(seg[j + dx], wv, acc[dy][dx]);
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Register-blocked DIRECT kernel (float, stride == 1, dilation == 1).
// Unlike a shared-memory staged design, each thread streams its (KH+OH-1)
// input row segments straight from global memory (L1-cached, fully
// vectorized 16B loads on the interior fast path) into registers and updates
// all OH output rows it contributes to.  No shared memory, no __syncthreads,
// no fill phase.  Live registers: OH*OW accumulators + (KW+OW-1) segment +
// KH*KW weights.  Edge threads take a guarded scalar path.
// 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_direct_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 W = KW + OW - 1;   // segment registers per row
    constexpr int VH = KH + OH - 1;  // segment rows per thread

    const long bc = blockIdx.z;
    const int c = (int)(bc % channels);

    // Weights -> registers (uniform broadcast loads; latency hidden by the
    // input row loads 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 ow = blockIdx.x * TILE_OW + (int)threadIdx.x * OW;
    const int oh = blockIdx.y * TILE_OH + (int)threadIdx.y * OH;
    const int iw0 = ow - pad_w;
    const int ih0 = oh - pad_h;

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

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

    const bool vec_ok =
        (((uintptr_t)in_base & (uintptr_t)15) == 0) &&
        ((in_w & 3) == 0) && ((iw0 & 3) == 0) &&
        (ih0 >= 0) && (iw0 >= 0) &&
        (ih0 + VH <= in_h) && (iw0 + W <= in_w);

    if (vec_ok) {
        const float* p = in_base + (long)ih0 * in_w + iw0;
        #pragma unroll
        for (int r = 0; r < VH; ++r) {
            float seg[W];
            LdsA<W>::run(p, seg);
            fma_seg_row<KH, KW, OW, OH>(seg, w, acc, r);
            p += in_w;
        }
    } else {
        #pragma unroll
        for (int r = 0; r < VH; ++r) {
            const int ih = ih0 + r;
            const bool row_ok = (ih >= 0) && (ih < in_h);
            const float* p = in_base + (long)ih * in_w + iw0;
            float seg[W];
            #pragma unroll
            for (int t = 0; t < W; ++t) {
                const int iw = iw0 + t;
                seg[t] = (row_ok && iw >= 0 && iw < in_w) ? p[t] : 0.f;
            }
            fma_seg_row<KH, KW, OW, OH>(seg, w, acc, r);
        }
    }

    float* out_base = output + bc * (long)out_h * out_w;
    const bool full_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_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. Fallback for sizes not covered by the direct 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;
}

// ---------------------------------------------------------------------------
// Direct register-blocked float path. Returns true when launched.
// ---------------------------------------------------------------------------
static bool try_launch_direct_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 block(32, 8);

#define TRY_DIRECT(KHV, KWV, OWV, OHV)                                             \
    if (kh == KHV && kw == KWV) {                                                  \
        const dim3 grid(                                                           \
            (unsigned)((out_w + 32 * OWV - 1) / (32 * OWV)),                       \
            (unsigned)((out_h + 8 * OHV - 1) / (8 * OHV)),                         \
            (unsigned)bc_total);                                                   \
        if (grid.y > 65535) return false;                                          \
        depthwise_conv2d_direct_kernel<KHV, KWV, OWV, OHV>                         \
            <<<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 true;                                                               \
    }

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

// ---------------------------------------------------------------------------
// Launch helper (fallback paths)
// ---------------------------------------------------------------------------
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: direct 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_direct_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_direct_ext_v6",
    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 direct-read 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,
        )
```