```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 kernels for depthwise 2D convolution.
#
# Two execution paths are provided:
#  1) A fast path for stride==1, float32 and common square kernel sizes
#     (1/3/5/7). Each thread computes 4 horizontally adjacent outputs using
#     a register sliding window fed by vectorized float4 shared-memory loads.
#  2) A generic tiled shared-memory kernel supporting arbitrary stride,
#     padding, kernel size, depthwise multiplier, optional bias and
#     float32/float64.
depthwise_conv2d_cpp_source = (
    "torch::Tensor depthwise_conv2d_forward(torch::Tensor x, torch::Tensor w, "
    "c10::optional<torch::Tensor> bias, int64_t stride, int64_t padding);"
)

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

namespace {

constexpr int TILE_X = 32;
constexpr int TILE_Y = 8;

// ---------------- generic tiled kernel: one thread per output ---------------
template <typename scalar_t, int K>
__global__ void dwconv2d_tiled_kernel(
    const scalar_t* __restrict__ x,
    const scalar_t* __restrict__ w,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ out,
    const int C_in, const int C_out, const int mult,
    const int H, const int W,
    const int H_out, const int W_out,
    const int stride, const int padding)
{
    const int tile_in_w = (TILE_X - 1) * stride + K;
    const int tile_in_h = (TILE_Y - 1) * stride + K;

    extern __shared__ __align__(16) unsigned char smem_raw[];
    scalar_t* s_in = reinterpret_cast<scalar_t*>(smem_raw);
    scalar_t* s_w  = s_in + tile_in_h * tile_in_w;

    const int tx  = threadIdx.x;
    const int ty  = threadIdx.y;
    const int tid = ty * TILE_X + tx;

    const int ow = blockIdx.x * TILE_X + tx;
    const int oh = blockIdx.y * TILE_Y + ty;

    const int nc = blockIdx.z;
    const int n  = nc / C_out;
    const int o  = nc - n * C_out;
    const int c  = o / mult;

    for (int i = tid; i < K * K; i += TILE_X * TILE_Y)
        s_w[i] = w[(int64_t)o * K * K + i];

    const int in_h0 = blockIdx.y * TILE_Y * stride - padding;
    const int in_w0 = blockIdx.x * TILE_X * stride - padding;

    const scalar_t* x_nc = x + ((int64_t)n * C_in + c) * H * W;
    const int elems = tile_in_h * tile_in_w;
    for (int i = tid; i < elems; i += TILE_X * TILE_Y) {
        const int r  = i / tile_in_w;
        const int q  = i - r * tile_in_w;
        const int ih = in_h0 + r;
        const int iw = in_w0 + q;
        scalar_t v = scalar_t(0);
        if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W)
            v = x_nc[(int64_t)ih * W + iw];
        s_in[i] = v;
    }
    __syncthreads();

    if (oh < H_out && ow < W_out) {
        scalar_t acc = bias ? bias[o] : scalar_t(0);
        const int bh = ty * stride;
        const int bw = tx * stride;
        #pragma unroll
        for (int kh = 0; kh < K; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < K; ++kw)
                acc += s_in[(bh + kh) * tile_in_w + bw + kw] * s_w[kh * K + kw];
        }
        out[((int64_t)nc * H_out + oh) * W_out + ow] = acc;
    }
}

// ------------- generic runtime-K kernel (any kernel size) -------------------
template <typename scalar_t>
__global__ void dwconv2d_generic_kernel(
    const scalar_t* __restrict__ x,
    const scalar_t* __restrict__ w,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ out,
    const int C_in, const int C_out, const int mult,
    const int H, const int W,
    const int H_out, const int W_out,
    const int K, const int stride, const int padding)
{
    const int tile_in_w = (TILE_X - 1) * stride + K;
    const int tile_in_h = (TILE_Y - 1) * stride + K;

    extern __shared__ __align__(16) unsigned char smem_raw[];
    scalar_t* s_in = reinterpret_cast<scalar_t*>(smem_raw);
    scalar_t* s_w  = s_in + tile_in_h * tile_in_w;

    const int tx  = threadIdx.x;
    const int ty  = threadIdx.y;
    const int tid = ty * TILE_X + tx;

    const int ow = blockIdx.x * TILE_X + tx;
    const int oh = blockIdx.y * TILE_Y + ty;

    const int nc = blockIdx.z;
    const int n  = nc / C_out;
    const int o  = nc - n * C_out;
    const int c  = o / mult;

    for (int i = tid; i < K * K; i += TILE_X * TILE_Y)
        s_w[i] = w[(int64_t)o * K * K + i];

    const int in_h0 = blockIdx.y * TILE_Y * stride - padding;
    const int in_w0 = blockIdx.x * TILE_X * stride - padding;

    const scalar_t* x_nc = x + ((int64_t)n * C_in + c) * H * W;
    const int elems = tile_in_h * tile_in_w;
    for (int i = tid; i < elems; i += TILE_X * TILE_Y) {
        const int r  = i / tile_in_w;
        const int q  = i - r * tile_in_w;
        const int ih = in_h0 + r;
        const int iw = in_w0 + q;
        scalar_t v = scalar_t(0);
        if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W)
            v = x_nc[(int64_t)ih * W + iw];
        s_in[i] = v;
    }
    __syncthreads();

    if (oh < H_out && ow < W_out) {
        scalar_t acc = bias ? bias[o] : scalar_t(0);
        const int bh = ty * stride;
        const int bw = tx * stride;
        for (int kh = 0; kh < K; ++kh)
            for (int kw = 0; kw < K; ++kw)
                acc += s_in[(bh + kh) * tile_in_w + bw + kw] * s_w[kh * K + kw];
        out[((int64_t)nc * H_out + oh) * W_out + ow] = acc;
    }
}

// ---- fast path: stride==1, float32; 4 outputs/thread with float4 window ----
template <int K>
__global__ void dwconv2d_vec_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int C_in, const int C_out, const int mult,
    const int H, const int W,
    const int H_out, const int W_out,
    const int padding)
{
    constexpr int OUT_X   = 4;
    constexpr int TILE_OX = TILE_X * OUT_X;   // 128 outputs per block row
    constexpr int TILE_OY = TILE_Y;           // 8 output rows per block
    constexpr int WIN     = K + OUT_X - 1;    // sliding window width
    constexpr int NVEC    = (WIN + 3) / 4;    // float4 loads per row

    const int tile_in_w = (TILE_OX + K - 1 + 3) & ~3;  // multiple of 4
    const int tile_in_h = TILE_OY + K - 1;

    extern __shared__ __align__(16) unsigned char smem_raw[];
    float* s_in = reinterpret_cast<float*>(smem_raw);
    float* s_w  = s_in + tile_in_h * tile_in_w;

    const int tx  = threadIdx.x;
    const int ty  = threadIdx.y;
    const int tid = ty * TILE_X + tx;

    const int nc = blockIdx.z;
    const int n  = nc / C_out;
    const int o  = nc - n * C_out;
    const int c  = o / mult;

    for (int i = tid; i < K * K; i += TILE_X * TILE_Y)
        s_w[i] = w[(int64_t)o * K * K + i];

    const int in_h0 = blockIdx.y * TILE_OY - padding;
    const int in_w0 = blockIdx.x * TILE_OX - padding;

    const float* x_nc = x + ((int64_t)n * C_in + c) * H * W;
    const int elems = tile_in_h * tile_in_w;
    for (int i = tid; i < elems; i += TILE_X * TILE_Y) {
        const int r  = i / tile_in_w;
        const int q  = i - r * tile_in_w;
        const int ih = in_h0 + r;
        const int iw = in_w0 + q;
        float v = 0.f;
        if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W)
            v = x_nc[(int64_t)ih * W + iw];
        s_in[i] = v;
    }
    __syncthreads();

    const int oh  = blockIdx.y * TILE_OY + ty;
    const int ow0 = blockIdx.x * TILE_OX + tx * OUT_X;

    if (oh < H_out && ow0 < W_out) {
        float acc[OUT_X];
        const float b0 = bias ? bias[o] : 0.f;
        #pragma unroll
        for (int i = 0; i < OUT_X; ++i) acc[i] = b0;

        #pragma unroll
        for (int kh = 0; kh < K; ++kh) {
            // tile_in_w is a multiple of 4, so rows are 16B aligned
            const float4* r4 = reinterpret_cast<const float4*>(
                s_in + (ty + kh) * tile_in_w) + tx;
            float v[NVEC * 4];
            #pragma unroll
            for (int qv = 0; qv < NVEC; ++qv) {
                const float4 t = r4[qv];
                v[4*qv+0] = t.x; v[4*qv+1] = t.y;
                v[4*qv+2] = t.z; v[4*qv+3] = t.w;
            }
            #pragma unroll
            for (int kw = 0; kw < K; ++kw) {
                const float wv = s_w[kh * K + kw];
                #pragma unroll
                for (int i = 0; i < OUT_X; ++i)
                    acc[i] += v[kw + i] * wv;
            }
        }

        float* orow = out + ((int64_t)nc * H_out + oh) * W_out;
        #pragma unroll
        for (int i = 0; i < OUT_X; ++i)
            if (ow0 + i < W_out) orow[ow0 + i] = acc[i];
    }
}

}  // namespace

torch::Tensor depthwise_conv2d_forward(torch::Tensor x, torch::Tensor w,
                                       c10::optional<torch::Tensor> bias,
                                       int64_t stride, int64_t padding)
{
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "x and w must be CUDA tensors");
    TORCH_CHECK(x.dim() == 4 && w.dim() == 4, "x and w must be 4D");
    TORCH_CHECK(x.scalar_type() == w.scalar_type(), "dtype mismatch");

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

    const int64_t N     = xc.size(0);
    const int64_t C_in  = xc.size(1);
    const int64_t H     = xc.size(2);
    const int64_t W     = xc.size(3);
    const int64_t C_out = wc.size(0);
    const int64_t K     = wc.size(2);
    TORCH_CHECK(wc.size(1) == 1 && wc.size(3) == K,
                "expected depthwise square-kernel weights");
    TORCH_CHECK(C_in > 0 && C_out % C_in == 0,
                "out_channels must be a multiple of in_channels");
    const int64_t mult = C_out / C_in;

    const int64_t H_out = (H + 2 * padding - K) / stride + 1;
    const int64_t W_out = (W + 2 * padding - K) / stride + 1;
    TORCH_CHECK(H_out >= 0 && W_out >= 0, "invalid output size");

    auto out = torch::empty({N, C_out, H_out, W_out}, xc.options());
    if (out.numel() == 0) return out;

    // Fall back to cuDNN for exotic configurations.
    const int64_t zblocks = N * C_out;
    const int64_t tw = (TILE_X - 1) * stride + K;
    const int64_t th = (TILE_Y - 1) * stride + K;
    const int64_t smem_bytes = (tw * th + K * K) * xc.element_size();
    if (zblocks > 65535 || smem_bytes > 44 * 1024) {
        return at::conv2d(xc, wc, bias, {stride, stride},
                          {padding, padding}, {1, 1}, C_in);
    }

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const bool templatedK = (K == 1 || K == 3 || K == 5 || K == 7);

    // Fast path: stride 1, float32, common kernel sizes.
    if (xc.scalar_type() == at::kFloat && stride == 1 && templatedK) {
        constexpr int TILE_OX = TILE_X * 4;
        constexpr int TILE_OY = TILE_Y;
        const int tile_in_w = (TILE_OX + (int)K - 1 + 3) & ~3;
        const int tile_in_h = TILE_OY + (int)K - 1;
        const size_t smem =
            (size_t)(tile_in_w * tile_in_h + (int)(K * K)) * sizeof(float);
        dim3 block(TILE_X, TILE_Y);
        dim3 grid((unsigned)((W_out + TILE_OX - 1) / TILE_OX),
                  (unsigned)((H_out + TILE_OY - 1) / TILE_OY),
                  (unsigned)zblocks);
        const float* xp = xc.data_ptr<float>();
        const float* wp = wc.data_ptr<float>();
        const float* bp = has_bias ? bc.data_ptr<float>() : nullptr;
        float* op = out.data_ptr<float>();
        #define LAUNCH_VEC(KK)                                                  \
            dwconv2d_vec_kernel<KK><<<grid, block, smem, stream>>>(             \
                xp, wp, bp, op, (int)C_in, (int)C_out, (int)mult, (int)H,      \
                (int)W, (int)H_out, (int)W_out, (int)padding)
        switch (K) {
            case 1: LAUNCH_VEC(1); break;
            case 3: LAUNCH_VEC(3); break;
            case 5: LAUNCH_VEC(5); break;
            case 7: LAUNCH_VEC(7); break;
        }
        #undef LAUNCH_VEC
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return out;
    }

    // Generic path.
    AT_DISPATCH_FLOATING_TYPES(xc.scalar_type(), "depthwise_conv2d_forward", [&] {
        const int tile_in_w = (int)tw;
        const int tile_in_h = (int)th;
        const size_t smem =
            (size_t)(tile_in_w * tile_in_h + (int)(K * K)) * sizeof(scalar_t);
        dim3 block(TILE_X, TILE_Y);
        dim3 grid((unsigned)((W_out + TILE_X - 1) / TILE_X),
                  (unsigned)((H_out + TILE_Y - 1) / TILE_Y),
                  (unsigned)zblocks);
        const scalar_t* xp = xc.data_ptr<scalar_t>();
        const scalar_t* wp = wc.data_ptr<scalar_t>();
        const scalar_t* bp = has_bias ? bc.data_ptr<scalar_t>() : nullptr;
        scalar_t* op = out.data_ptr<scalar_t>();
        #define LAUNCH_T(KK)                                                    \
            dwconv2d_tiled_kernel<scalar_t, KK><<<grid, block, smem, stream>>>( \
                xp, wp, bp, op, (int)C_in, (int)C_out, (int)mult, (int)H,      \
                (int)W, (int)H_out, (int)W_out, (int)stride, (int)padding)
        switch (K) {
            case 1: LAUNCH_T(1); break;
            case 3: LAUNCH_T(3); break;
            case 5: LAUNCH_T(5); break;
            case 7: LAUNCH_T(7); break;
            default:
                dwconv2d_generic_kernel<scalar_t>
                    <<<grid, block, smem, stream>>>(
                        xp, wp, bp, op, (int)C_in, (int)C_out, (int)mult,
                        (int)H, (int)W, (int)H_out, (int)W_out, (int)K,
                        (int)stride, (int)padding);
        }
        #undef LAUNCH_T
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    });

    return out;
}
"""

depthwise_conv2d_ext = load_inline(
    name="depthwise_conv2d_ext",
    cpp_sources=depthwise_conv2d_cpp_source,
    cuda_sources=depthwise_conv2d_cuda_source,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution backed by a custom CUDA kernel.

    Uses a shared-memory tiled kernel with vectorized sliding-window loads
    (fast path: stride==1, float32, kernel size 1/3/5/7) and falls back to the
    native PyTorch implementation whenever the custom kernel does not apply
    (e.g. CPU tensors or unsupported dtypes).
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(
            in_channels, out_channels,
            kernel_size=(kernel_size, kernel_size),
            stride=stride, padding=padding,
            groups=in_channels, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (x.is_cuda and x.dim() == 4
                and x.dtype in (torch.float32, torch.float64)
                and x.dtype == self.conv2d.weight.dtype):
            return depthwise_conv2d_ext.depthwise_conv2d_forward(
                x, self.conv2d.weight, self.conv2d.bias,
                self.conv2d.stride[0], self.conv2d.padding[0])
        return self.conv2d(x)
```