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

# Shared-memory tiled depthwise 2D convolution kernel.
# Each thread block handles one (batch, channel) pair and produces a
# TILE_W x TILE_H tile of the output. The corresponding input patch
# (including the kernel halo) is staged in shared memory with fully
# coalesced global loads, then every thread computes its output pixel
# from shared memory. The kernel is templated on the kernel size so the
# inner convolution loops are fully unrolled for common sizes (1/3/5/7).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define TILE_W 32
#define TILE_H 16
#define BLOCK_THREADS (TILE_W * TILE_H)

template <int KT>
__global__ void __launch_bounds__(BLOCK_THREADS)
depthwise_conv2d_tiled_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int C,
    const int H,
    const int W,
    const int outH,
    const int outW,
    const int stride,
    const int padding)
{
    extern __shared__ float tile[];

    const int nc = blockIdx.z;          // linearized (batch, channel)
    const int c  = nc % C;              // channel index

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int h_out = blockIdx.y * TILE_H + ty;
    const int w_out = blockIdx.x * TILE_W + tx;

    // Top-left corner (in input coordinates) of the patch that backs this tile.
    const int h_in_base = (int)(blockIdx.y * TILE_H) * stride - padding;
    const int w_in_base = (int)(blockIdx.x * TILE_W) * stride - padding;

    const int smemH = (TILE_H - 1) * stride + KT;
    const int smemW = (TILE_W - 1) * stride + KT;

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

    // Cooperative, coalesced load of the input patch into shared memory.
    const int total = smemH * smemW;
    const int tid = ty * TILE_W + tx;
    for (int idx = tid; idx < total; idx += BLOCK_THREADS) {
        const int i = idx / smemW;
        const int j = idx - i * smemW;
        const int h_in = h_in_base + i;
        const int w_in = w_in_base + j;
        float v = 0.0f;
        if ((unsigned int)h_in < (unsigned int)H && (unsigned int)w_in < (unsigned int)W) {
            v = in_nc[(long long)h_in * W + w_in];
        }
        tile[idx] = v;
    }
    __syncthreads();

    if (h_out < outH && w_out < outW) {
        const float* __restrict__ w_c = weight + (long long)c * KT * KT;
        const int anchor = (ty * stride) * smemW + (tx * stride);
        float acc = 0.0f;
        #pragma unroll
        for (int kh = 0; kh < KT; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < KT; ++kw) {
                acc += tile[anchor + kh * smemW + kw] * w_c[kh * KT + kw];
            }
        }
        if (bias != nullptr) {
            acc += bias[c];
        }
        output[((long long)nc * outH + h_out) * outW + w_out] = acc;
    }
}

// Generic fallback for arbitrary (runtime) kernel sizes.
__global__ void __launch_bounds__(BLOCK_THREADS)
depthwise_conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int C,
    const int H,
    const int W,
    const int outH,
    const int outW,
    const int stride,
    const int padding,
    const int K)
{
    extern __shared__ float tile[];

    const int nc = blockIdx.z;
    const int c  = nc % C;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int h_out = blockIdx.y * TILE_H + ty;
    const int w_out = blockIdx.x * TILE_W + tx;

    const int h_in_base = (int)(blockIdx.y * TILE_H) * stride - padding;
    const int w_in_base = (int)(blockIdx.x * TILE_W) * stride - padding;

    const int smemH = (TILE_H - 1) * stride + K;
    const int smemW = (TILE_W - 1) * stride + K;

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

    const int total = smemH * smemW;
    const int tid = ty * TILE_W + tx;
    for (int idx = tid; idx < total; idx += BLOCK_THREADS) {
        const int i = idx / smemW;
        const int j = idx - i * smemW;
        const int h_in = h_in_base + i;
        const int w_in = w_in_base + j;
        float v = 0.0f;
        if ((unsigned int)h_in < (unsigned int)H && (unsigned int)w_in < (unsigned int)W) {
            v = in_nc[(long long)h_in * W + w_in];
        }
        tile[idx] = v;
    }
    __syncthreads();

    if (h_out < outH && w_out < outW) {
        const float* __restrict__ w_c = weight + (long long)c * K * K;
        const int anchor = (ty * stride) * smemW + (tx * stride);
        float acc = 0.0f;
        for (int kh = 0; kh < K; ++kh) {
            for (int kw = 0; kw < K; ++kw) {
                acc += tile[anchor + kh * smemW + kw] * w_c[kh * K + kw];
            }
        }
        if (bias != nullptr) {
            acc += bias[c];
        }
        output[((long long)nc * outH + h_out) * outW + w_out] = acc;
    }
}

torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride,
    int64_t padding)
{
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "input and weight must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 tensors are supported");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "input and weight must be 4-D");
    TORCH_CHECK(weight.size(2) == weight.size(3), "kernel must be square");
    TORCH_CHECK(weight.size(1) == 1, "weight must be depthwise (in_channels/groups == 1)");
    TORCH_CHECK(weight.size(0) == input.size(1), "only channel multiplier 1 is supported");

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

    const int64_t N = x.size(0);
    const int64_t C = x.size(1);
    const int64_t H = x.size(2);
    const int64_t W = x.size(3);
    const int64_t outC = w.size(0);
    const int64_t K = w.size(2);

    const int64_t outH = (H + 2 * padding - K) / stride + 1;
    const int64_t outW = (W + 2 * padding - K) / stride + 1;

    auto out = torch::empty({N, outC, outH, outW}, x.options());
    if (out.numel() == 0) {
        return out;
    }

    TORCH_CHECK(N * C <= 65535, "batch * channels exceeds grid.z limit");

    const int smemH = (TILE_H - 1) * (int)stride + (int)K;
    const int smemW = (TILE_W - 1) * (int)stride + (int)K;
    const size_t smem_bytes = (size_t)smemH * (size_t)smemW * sizeof(float);
    TORCH_CHECK(smem_bytes <= 48 * 1024, "shared memory requirement too large");

    const dim3 block(TILE_W, TILE_H);
    const dim3 grid(
        (unsigned int)((outW + TILE_W - 1) / TILE_W),
        (unsigned int)((outH + TILE_H - 1) / TILE_H),
        (unsigned int)(N * C));

    const float* bias_ptr = nullptr;
    torch::Tensor b;
    if (bias.has_value() && bias->defined() && bias->numel() > 0) {
        b = bias->contiguous();
        bias_ptr = b.data_ptr<float>();
    }

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

    const float* x_ptr = x.data_ptr<float>();
    const float* w_ptr = w.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();

    switch (K) {
        case 1:
            depthwise_conv2d_tiled_kernel<1><<<grid, block, smem_bytes, stream>>>(
                x_ptr, w_ptr, bias_ptr, out_ptr, (int)C, (int)H, (int)W,
                (int)outH, (int)outW, (int)stride, (int)padding);
            break;
        case 3:
            depthwise_conv2d_tiled_kernel<3><<<grid, block, smem_bytes, stream>>>(
                x_ptr, w_ptr, bias_ptr, out_ptr, (int)C, (int)H, (int)W,
                (int)outH, (int)outW, (int)stride, (int)padding);
            break;
        case 5:
            depthwise_conv2d_tiled_kernel<5><<<grid, block, smem_bytes, stream>>>(
                x_ptr, w_ptr, bias_ptr, out_ptr, (int)C, (int)H, (int)W,
                (int)outH, (int)outW, (int)stride, (int)padding);
            break;
        case 7:
            depthwise_conv2d_tiled_kernel<7><<<grid, block, smem_bytes, stream>>>(
                x_ptr, w_ptr, bias_ptr, out_ptr, (int)C, (int)H, (int)W,
                (int)outH, (int)outW, (int)stride, (int)padding);
            break;
        default:
            depthwise_conv2d_generic_kernel<<<grid, block, smem_bytes, stream>>>(
                x_ptr, w_ptr, bias_ptr, out_ptr, (int)C, (int)H, (int)W,
                (int)outH, (int)outW, (int)stride, (int)padding, (int)K);
            break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

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

_depthwise_conv2d_ext = load_inline(
    name="depthwise_conv2d_tiled_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)

# Keep these in sync with the kernel launch configuration.
_KERNEL_TILE_W = 32
_KERNEL_TILE_H = 16
_MAX_GRID_Z = 65535
_MAX_SMEM_BYTES = 48 * 1024


class _DepthwiseConv2dFunc(torch.autograd.Function):
    """Autograd wrapper so the custom forward kernel also works in training."""

    @staticmethod
    def forward(ctx, x, weight, bias, stride, padding):
        ctx.stride = stride
        ctx.padding = padding
        ctx.has_bias = bias is not None
        ctx.save_for_backward(x, weight, bias)
        return _depthwise_conv2d_ext.depthwise_conv2d_forward(
            x, weight, bias, stride, padding
        )

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, bias = ctx.saved_tensors
        needs = ctx.needs_input_grad
        grad_input = grad_weight = grad_bias = None
        if needs[0] or needs[1] or (needs[2] and ctx.has_bias):
            bias_sizes = [weight.size(0)] if ctx.has_bias else None
            grad_input, grad_weight, grad_bias = torch.ops.aten.convolution_backward(
                grad_output.contiguous(),
                x,
                weight,
                bias_sizes,
                [ctx.stride, ctx.stride],
                [ctx.padding, ctx.padding],
                [1, 1],
                False,
                [0, 0],
                x.size(1),
                [bool(needs[0]), bool(needs[1]), bool(needs[2]) and ctx.has_bias],
            )
        return grad_input, grad_weight, grad_bias, None, None


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom shared-memory tiled
    CUDA kernel. Falls back to the native PyTorch implementation whenever
    the custom kernel does not support the given configuration.
    """

    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 = self.conv2d.stride[0]
        self.padding = self.conv2d.padding[0]

    def _use_custom_kernel(self, x: torch.Tensor, weight: torch.Tensor) -> bool:
        if not (x.is_cuda and weight.is_cuda):
            return False
        if x.dtype != torch.float32 or weight.dtype != torch.float32:
            return False
        if self.conv2d.dilation[0] != 1 or self.conv2d.dilation[1] != 1:
            return False
        # Depthwise with channel multiplier 1 and square kernel.
        if weight.size(1) != 1 or weight.size(0) != x.size(1):
            return False
        if weight.size(2) != weight.size(3):
            return False
        # Kernel launch limits.
        if x.size(0) * x.size(1) > _MAX_GRID_Z:
            return False
        k = weight.size(2)
        s = self.stride
        smem_h = (_KERNEL_TILE_H - 1) * s + k
        smem_w = (_KERNEL_TILE_W - 1) * s + k
        if smem_h * smem_w * 4 > _MAX_SMEM_BYTES:
            return False
        return True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.conv2d.weight
        bias = self.conv2d.bias
        if self._use_custom_kernel(x, weight):
            return _DepthwiseConv2dFunc.apply(x, weight, bias, self.stride, self.padding)
        return F.conv2d(
            x,
            weight,
            bias,
            self.conv2d.stride,
            self.conv2d.padding,
            self.conv2d.dilation,
            self.conv2d.groups,
        )
```