```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 kernel for depthwise 2D convolution.
# Strategy: shared-memory tiled depthwise conv. Each block handles a 64x32
# output tile of one (batch, channel) slice. 512 threads, each thread computes
# a 2x2 group of outputs. The kernel weights are preloaded into registers and
# the input tile (with halo) is staged in shared memory. Access patterns are
# arranged so both global loads/stores and shared-memory reads are coalesced /
# bank-conflict free. Kernel size is a template parameter (1..7) so the inner
# loops are fully unrolled; a generic fallback kernel handles any other case.
depthwise_conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BX 32
#define BY 16
#define TW (BX * 2)
#define TH (BY * 2)

template <int K>
__global__ void dwconv2d_tiled_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int Cin, const int Cout,
    const int H, const int W,
    const int outH, const int outW,
    const int stride, const int pad)
{
    const int c  = blockIdx.z % Cout;
    const int n  = blockIdx.z / Cout;
    const int ic = c / (Cout / Cin);
    const long nc_in = (long)n * Cin + ic;

    const int shW = (TW - 1) * stride + K;
    const int shH = (TH - 1) * stride + K;

    extern __shared__ float tile[];

    // Preload this channel's kernel into registers.
    float wreg[K * K];
    const float* wp = weight + (long)c * K * K;
    #pragma unroll
    for (int i = 0; i < K * K; ++i) wreg[i] = wp[i];

    const int in_x0 = blockIdx.x * TW * stride - pad;
    const int in_y0 = blockIdx.y * TH * stride - pad;

    const float* in_nc = input + nc_in * (long)H * W;

    const int tid = threadIdx.y * BX + threadIdx.x;
    const int nthreads = BX * BY;
    const int tile_size = shW * shH;

    // Stage the input tile (including halo) into shared memory, zero padding
    // out-of-bounds elements.
    for (int i = tid; i < tile_size; i += nthreads) {
        const int ty = i / shW;
        const int tx = i - ty * shW;
        const int ix = in_x0 + tx;
        const int iy = in_y0 + ty;
        float v = 0.0f;
        if ((unsigned)ix < (unsigned)W && (unsigned)iy < (unsigned)H)
            v = in_nc[(long)iy * W + ix];
        tile[i] = v;
    }
    __syncthreads();

    // Each thread computes a 2x2 set of outputs located at (tx, ty),
    // (tx+BX, ty), (tx, ty+BY), (tx+BX, ty+BY) within the tile.
    const int sxA = threadIdx.x * stride;
    const int sxB = (threadIdx.x + BX) * stride;
    const int syA = threadIdx.y * stride;
    const int syB = (threadIdx.y + BY) * stride;

    float acc00 = 0.f, acc01 = 0.f, acc10 = 0.f, acc11 = 0.f;

    #pragma unroll
    for (int ky = 0; ky < K; ++ky) {
        const float* rowA = tile + (syA + ky) * shW;
        const float* rowB = tile + (syB + ky) * shW;
        #pragma unroll
        for (int kx = 0; kx < K; ++kx) {
            const float wv = wreg[ky * K + kx];
            acc00 += rowA[sxA + kx] * wv;
            acc01 += rowA[sxB + kx] * wv;
            acc10 += rowB[sxA + kx] * wv;
            acc11 += rowB[sxB + kx] * wv;
        }
    }

    if (bias != nullptr) {
        const float b = bias[c];
        acc00 += b; acc01 += b; acc10 += b; acc11 += b;
    }

    const int ox0 = blockIdx.x * TW + threadIdx.x;
    const int ox1 = ox0 + BX;
    const int oy0 = blockIdx.y * TH + threadIdx.y;
    const int oy1 = oy0 + BY;

    float* out_nc = output + ((long)n * Cout + c) * (long)outH * outW;
    if (oy0 < outH) {
        if (ox0 < outW) out_nc[(long)oy0 * outW + ox0] = acc00;
        if (ox1 < outW) out_nc[(long)oy0 * outW + ox1] = acc01;
    }
    if (oy1 < outH) {
        if (ox0 < outW) out_nc[(long)oy1 * outW + ox0] = acc10;
        if (ox1 < outW) out_nc[(long)oy1 * outW + ox1] = acc11;
    }
}

// Generic fallback (any kernel size / stride / padding / size).
__global__ void dwconv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int Cin, const int Cout,
    const int H, const int W,
    const int outH, const int outW,
    const int K, const int stride, const int pad,
    const long total)
{
    const long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int ox = idx % outW;
    const int oy = (idx / outW) % outH;
    const long nco = idx / ((long)outW * outH);
    const int c  = nco % Cout;
    const int n  = nco / Cout;
    const int ic = c / (Cout / Cin);

    const float* wp = weight + (long)c * K * K;
    const float* in_nc = input + ((long)n * Cin + ic) * (long)H * W;

    float acc = 0.f;
    for (int ky = 0; ky < K; ++ky) {
        const int iy = oy * stride - pad + ky;
        if (iy < 0 || iy >= H) continue;
        for (int kx = 0; kx < K; ++kx) {
            const int ix = ox * stride - pad + kx;
            if (ix < 0 || ix >= W) continue;
            acc += in_nc[(long)iy * W + ix] * wp[ky * K + kx];
        }
    }
    if (bias != nullptr) acc += bias[c];
    output[idx] = acc;
}

torch::Tensor depthwise_conv2d_cuda(torch::Tensor input, torch::Tensor weight,
                                    torch::Tensor bias, int64_t stride,
                                    int64_t padding)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 4, "input must be NCHW");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 supported");

    auto x = input.contiguous();
    auto w = weight.contiguous();

    const int N    = x.size(0);
    const int Cin  = x.size(1);
    const int H    = x.size(2);
    const int W    = x.size(3);
    const int Cout = w.size(0);
    const int K    = w.size(2);

    const int s = (int)stride;
    const int p = (int)padding;

    const int outH = (H + 2 * p - K) / s + 1;
    const int outW = (W + 2 * p - K) / s + 1;

    auto out = torch::empty({N, Cout, outH, outW}, x.options());

    const float* bias_ptr =
        (bias.defined() && bias.numel() > 0) ? bias.data_ptr<float>() : nullptr;

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

    const long NC = (long)N * Cout;
    const long total = NC * (long)outH * outW;
    if (total == 0) return out;

    const int shW = (TW - 1) * s + K;
    const int shH = (TH - 1) * s + K;
    const size_t smem = (size_t)shW * shH * sizeof(float);

    if (K >= 1 && K <= 7 && NC <= 65535 && smem <= 44 * 1024) {
        dim3 block(BX, BY);
        dim3 grid((outW + TW - 1) / TW, (outH + TH - 1) / TH, (unsigned)NC);
        const float* xp = x.data_ptr<float>();
        const float* wp = w.data_ptr<float>();
        float* op = out.data_ptr<float>();
        switch (K) {
            case 1: dwconv2d_tiled_kernel<1><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
            case 2: dwconv2d_tiled_kernel<2><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
            case 3: dwconv2d_tiled_kernel<3><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
            case 4: dwconv2d_tiled_kernel<4><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
            case 5: dwconv2d_tiled_kernel<5><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
            case 6: dwconv2d_tiled_kernel<6><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
            case 7: dwconv2d_tiled_kernel<7><<<grid, block, smem, stream>>>(xp, wp, bias_ptr, op, Cin, Cout, H, W, outH, outW, s, p); break;
        }
    } else {
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        dwconv2d_generic_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), Cin, Cout, H, W, outH, outW, K, s, p, total);
    }
    return out;
}
"""

depthwise_conv2d_cpp_source = (
    "torch::Tensor depthwise_conv2d_cuda(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride, int64_t padding);"
)

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


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom shared-memory tiled
    CUDA kernel (falls back to nn.Conv2d for non-CUDA / non-float32 inputs).
    """
    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)
        self.stride = stride
        self.padding = padding
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32 and self.conv2d.weight.dtype == torch.float32:
            bias = self.conv2d.bias
            if bias is None:
                bias = self._empty_bias
            return _depthwise_conv2d.depthwise_conv2d_cuda(
                x, self.conv2d.weight, bias, self.stride, self.padding)
        return self.conv2d(x)
```