```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

# Custom shared-memory tiled depthwise 2D convolution CUDA kernel.
# Each thread block computes a TILE_W x TILE_H tile of the output for one
# (batch, channel) pair. The corresponding input tile (including the kernel
# halo) is staged in shared memory with coalesced global loads, so each input
# element is read from global memory roughly once instead of KH*KW times.
# The kernel is templated on the kernel size so the common square sizes
# (1x1, 3x3, 5x5, 7x7) get fully unrolled inner loops, with a generic
# fallback for arbitrary kernel sizes, strides, and padding. Bias addition
# is fused into the kernel.
depthwise_conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <climits>

#define TILE_W 32
#define TILE_H 16
#define MAX_SHMEM_BYTES (48 * 1024)

template <int KH_T, int KW_T>
__global__ void depthwise_conv2d_forward_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int batch,
    const int in_channels,
    const int out_channels,
    const int H,
    const int W,
    const int outH,
    const int outW,
    const int KH,
    const int KW,
    const int stride,
    const int pad,
    const int tilesX,
    const int tilesY)
{
    const int khN = (KH_T > 0) ? KH_T : KH;
    const int kwN = (KW_T > 0) ? KW_T : KW;

    const int shW = (TILE_W - 1) * stride + kwN;
    const int shH = (TILE_H - 1) * stride + khN;

    extern __shared__ float tile[];

    // Decode the 1D block index into (batch, channel, tile_y, tile_x).
    const int tile_x = blockIdx.x % tilesX;
    int rem = blockIdx.x / tilesX;
    const int tile_y = rem % tilesY;
    const int bc = rem / tilesY;
    const int c = bc % out_channels;
    const int b = bc / out_channels;
    const int depth_mult = out_channels / in_channels;
    const int ic = c / depth_mult;

    const int out_x = tile_x * TILE_W + threadIdx.x;
    const int out_y = tile_y * TILE_H + threadIdx.y;

    // Top-left corner (in input coordinates) of the shared-memory tile.
    const int in_x0 = tile_x * TILE_W * stride - pad;
    const int in_y0 = tile_y * TILE_H * stride - pad;

    const float* in_base = input + ((int64_t)b * in_channels + ic) * (int64_t)H * (int64_t)W;

    // Stage the input tile (with halo) into shared memory. Out-of-bounds
    // elements are zero-filled, which also implements zero padding.
    const int tid = threadIdx.y * TILE_W + threadIdx.x;
    const int nthreads = TILE_W * TILE_H;
    const int shSize = shH * shW;
    for (int i = tid; i < shSize; i += nthreads) {
        const int sy = i / shW;
        const int sx = i - sy * shW;
        const int iy = in_y0 + sy;
        const int ix = in_x0 + sx;
        float v = 0.0f;
        if ((unsigned int)iy < (unsigned int)H && (unsigned int)ix < (unsigned int)W) {
            v = in_base[(int64_t)iy * W + ix];
        }
        tile[i] = v;
    }
    __syncthreads();

    if (out_x < outW && out_y < outH) {
        const float* w_c = weight + (int64_t)c * khN * kwN;
        float acc = (bias != nullptr) ? bias[c] : 0.0f;
        const int ty = threadIdx.y * stride;
        const int tx = threadIdx.x * stride;
        #pragma unroll
        for (int kh = 0; kh < khN; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < kwN; ++kw) {
                acc = fmaf(tile[(ty + kh) * shW + (tx + kw)], w_c[kh * kwN + kw], acc);
            }
        }
        output[((int64_t)bc * outH + out_y) * (int64_t)outW + out_x] = acc;
    }
}

torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    int64_t stride,
    int64_t padding)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
    TORCH_CHECK(weight.dim() == 4 && weight.size(1) == 1, "weight must be (O, 1, kH, kW)");
    TORCH_CHECK(stride >= 1 && padding >= 0, "invalid stride/padding");

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

    const int batch = (int)x.size(0);
    const int in_channels = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int out_channels = (int)w.size(0);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);

    TORCH_CHECK(in_channels > 0 && out_channels % in_channels == 0,
                "out_channels must be a positive multiple of in_channels");

    const int outH = (int)((H + 2 * padding - KH) / stride + 1);
    const int outW = (int)((W + 2 * padding - KW) / stride + 1);
    TORCH_CHECK(outH > 0 && outW > 0, "output spatial size must be positive");

    auto out = at::empty({batch, out_channels, outH, outW}, x.options());

    const int tilesX = (outW + TILE_W - 1) / TILE_W;
    const int tilesY = (outH + TILE_H - 1) / TILE_H;
    const int64_t totalBlocks = (int64_t)batch * out_channels * tilesX * tilesY;
    TORCH_CHECK(totalBlocks < (int64_t)INT_MAX, "grid too large");

    const int shW = (TILE_W - 1) * (int)stride + KW;
    const int shH = (TILE_H - 1) * (int)stride + KH;
    const size_t shmem = (size_t)shW * (size_t)shH * sizeof(float);
    TORCH_CHECK(shmem <= MAX_SHMEM_BYTES, "shared memory requirement too large");

    const float* in_ptr = x.data_ptr<float>();
    const float* w_ptr = w.data_ptr<float>();
    const float* b_ptr = nullptr;
    torch::Tensor bias_t;
    if (bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0) {
        bias_t = bias_opt->contiguous();
        b_ptr = bias_t.data_ptr<float>();
    }
    float* out_ptr = out.data_ptr<float>();

    const dim3 block(TILE_W, TILE_H);
    const dim3 grid((unsigned int)totalBlocks);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    #define LAUNCH_KERNEL(KHT, KWT)                                                 \
        depthwise_conv2d_forward_kernel<KHT, KWT><<<grid, block, shmem, stream>>>(  \
            in_ptr, w_ptr, b_ptr, out_ptr, batch, in_channels, out_channels, H, W,  \
            outH, outW, KH, KW, (int)stride, (int)padding, tilesX, tilesY)

    if (KH == 3 && KW == 3) {
        LAUNCH_KERNEL(3, 3);
    } else if (KH == 1 && KW == 1) {
        LAUNCH_KERNEL(1, 1);
    } else if (KH == 5 && KW == 5) {
        LAUNCH_KERNEL(5, 5);
    } else if (KH == 7 && KW == 7) {
        LAUNCH_KERNEL(7, 7);
    } else {
        LAUNCH_KERNEL(0, 0);
    }
    #undef LAUNCH_KERNEL

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

depthwise_conv2d_cpp_source = (
    "torch::Tensor depthwise_conv2d_forward(torch::Tensor input, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias_opt, int64_t stride, int64_t padding);"
)

# Compile the inline CUDA extension.
depthwise_conv2d_mod = 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 shared-memory tiled CUDA kernel.

    Falls back to the standard F.conv2d implementation for configurations the
    custom kernel does not support (non-CUDA tensors, non-float32 dtypes, or
    tile shared-memory requirements beyond the 48 KB static limit).
    """
    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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv2d
        w = conv.weight
        if x.is_cuda and x.dtype == torch.float32 and conv.groups == conv.in_channels:
            kh, kw = w.shape[2], w.shape[3]
            sh_w = 31 * self.stride + kw   # shared-memory tile width
            sh_h = 15 * self.stride + kh   # shared-memory tile height
            if sh_w * sh_h * 4 <= 48 * 1024:
                return depthwise_conv2d_mod.depthwise_conv2d_forward(
                    x, w, conv.bias, self.stride, self.padding
                )
        return F.conv2d(x, w, conv.bias, stride=conv.stride, padding=conv.padding, groups=conv.groups)
```