```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 CUDA kernels for depthwise 2D convolution.
#  - A shared-memory tiled kernel specialized for the common (and given) case:
#    kernel_size=3, stride=1, padding=0, out_channels == in_channels.
#  - A generic depthwise kernel that handles any kernel size / stride / padding
#    (and channel multipliers) as a fallback.
_depthwise_conv_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define TILE_W 32
#define TILE_H 16

// Specialized depthwise conv: 3x3 kernel, stride 1, padding 0, out_channels == in_channels.
// Each block computes a TILE_W x TILE_H tile of the output for one (batch, channel) pair.
__global__ void depthwise_conv2d_3x3_s1_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int channels,
    const int height_in, const int width_in,
    const int height_out, const int width_out)
{
    // (TILE_H+2) x (TILE_W+2) halo tile, row pitch TILE_W+3 to avoid bank conflicts.
    __shared__ float tile[TILE_H + 2][TILE_W + 3];

    const int nc = blockIdx.z;              // combined n * channels + c
    const int c  = nc % channels;
    const long long base = (long long)nc * height_in * width_in;

    const int out_y0 = blockIdx.y * TILE_H;
    const int out_x0 = blockIdx.x * TILE_W;

    const int tid = threadIdx.y * TILE_W + threadIdx.x;
    const int nthreads = TILE_W * TILE_H;

    // Cooperative load of the input tile (with zero-fill outside the image).
    #pragma unroll
    for (int i = tid; i < (TILE_H + 2) * (TILE_W + 2); i += nthreads) {
        const int ty = i / (TILE_W + 2);
        const int tx = i - ty * (TILE_W + 2);
        const int iy = out_y0 + ty;
        const int ix = out_x0 + tx;
        float v = 0.0f;
        if (iy < height_in && ix < width_in) {
            v = __ldg(input + base + (long long)iy * width_in + ix);
        }
        tile[ty][tx] = v;
    }
    __syncthreads();

    const int ox = out_x0 + threadIdx.x;
    const int oy = out_y0 + threadIdx.y;
    if (ox < width_out && oy < height_out) {
        const float* w = weight + c * 9;
        const float w0 = __ldg(w + 0), w1 = __ldg(w + 1), w2 = __ldg(w + 2);
        const float w3 = __ldg(w + 3), w4 = __ldg(w + 4), w5 = __ldg(w + 5);
        const float w6 = __ldg(w + 6), w7 = __ldg(w + 7), w8 = __ldg(w + 8);

        const int ty = threadIdx.y;
        const int tx = threadIdx.x;
        float sum =
              tile[ty    ][tx    ] * w0 + tile[ty    ][tx + 1] * w1 + tile[ty    ][tx + 2] * w2
            + tile[ty + 1][tx    ] * w3 + tile[ty + 1][tx + 1] * w4 + tile[ty + 1][tx + 2] * w5
            + tile[ty + 2][tx    ] * w6 + tile[ty + 2][tx + 1] * w7 + tile[ty + 2][tx + 2] * w8;

        if (bias != nullptr) sum += __ldg(bias + c);

        output[(long long)nc * height_out * width_out + (long long)oy * width_out + ox] = sum;
    }
}

// Generic depthwise conv fallback: any kernel size / stride / padding / channel multiplier.
__global__ void depthwise_conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int batch_size, const int in_channels, const int out_channels,
    const int height_in, const int width_in,
    const int height_out, const int width_out,
    const int kernel_size, const int stride, const int padding)
{
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    const long long total = (long long)batch_size * out_channels * height_out * width_out;
    if (idx >= total) return;

    const int w_out = (int)(idx % width_out);
    long long t = idx / width_out;
    const int h_out = (int)(t % height_out);
    t /= height_out;
    const int c = (int)(t % out_channels);
    const int n = (int)(t / out_channels);

    // Depthwise grouping: output channel c reads input channel c_in.
    const int c_in = (int)(((long long)c * in_channels) / out_channels);

    const float* in_ptr = input + ((long long)n * in_channels + c_in) * height_in * width_in;
    const float* w_ptr = weight + (long long)c * kernel_size * kernel_size;

    const int h_start = h_out * stride - padding;
    const int w_start = w_out * stride - padding;

    float sum = 0.0f;
    for (int kh = 0; kh < kernel_size; ++kh) {
        const int h_in = h_start + kh;
        if (h_in < 0 || h_in >= height_in) continue;
        const float* in_row = in_ptr + (long long)h_in * width_in;
        const float* w_row = w_ptr + kh * kernel_size;
        for (int kw = 0; kw < kernel_size; ++kw) {
            const int w_in = w_start + kw;
            if (w_in < 0 || w_in >= width_in) continue;
            sum += __ldg(in_row + w_in) * __ldg(w_row + kw);
        }
    }
    if (bias != nullptr) sum += __ldg(bias + c);
    output[idx] = sum;
}

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() && weight.is_cuda(), "input/weight must be CUDA tensors");
    TORCH_CHECK(input.dim() == 4, "input must be NCHW");
    TORCH_CHECK(input.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(weight.size(2) == weight.size(3), "kernel must be square");

    const at::cuda::CUDAGuard device_guard(input.device());

    auto input_c = input.contiguous();
    auto weight_c = weight.contiguous();

    const int batch_size  = (int)input_c.size(0);
    const int in_channels = (int)input_c.size(1);
    const int height_in   = (int)input_c.size(2);
    const int width_in    = (int)input_c.size(3);
    const int out_channels = (int)weight_c.size(0);
    const int kernel_size  = (int)weight_c.size(2);

    const int height_out = (height_in + 2 * (int)padding - kernel_size) / (int)stride + 1;
    const int width_out  = (width_in  + 2 * (int)padding - kernel_size) / (int)stride + 1;

    auto output = torch::empty({batch_size, out_channels, height_out, width_out},
                               input_c.options());

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

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

    const float* in_ptr = input_c.data_ptr<float>();
    const float* w_ptr  = weight_c.data_ptr<float>();
    float* out_ptr      = output.data_ptr<float>();

    if (kernel_size == 3 && stride == 1 && padding == 0 &&
        out_channels == in_channels &&
        (long long)batch_size * in_channels <= 65535) {
        dim3 block(TILE_W, TILE_H);
        dim3 grid((width_out + TILE_W - 1) / TILE_W,
                  (height_out + TILE_H - 1) / TILE_H,
                  batch_size * in_channels);
        depthwise_conv2d_3x3_s1_kernel<<<grid, block, 0, stream>>>(
            in_ptr, w_ptr, bias_ptr, out_ptr,
            in_channels, height_in, width_in, height_out, width_out);
    } else {
        const long long total = (long long)batch_size * out_channels * height_out * width_out;
        const int block = 256;
        const long long grid = (total + block - 1) / block;
        depthwise_conv2d_generic_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            in_ptr, w_ptr, bias_ptr, out_ptr,
            batch_size, in_channels, out_channels,
            height_in, width_in, height_out, width_out,
            kernel_size, (int)stride, (int)padding);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

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

_depthwise_conv_ext = load_inline(
    name="depthwise_conv2d_ext",
    cpp_sources=_depthwise_conv_cpp_src,
    cuda_sources=_depthwise_conv_cuda_src,
    functions=["depthwise_conv2d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    The parameters are stored in an nn.Conv2d module so the state_dict is fully
    compatible with the original Model, but the forward pass on CUDA runs a
    hand-written shared-memory tiled depthwise convolution kernel.
    """
    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:
        weight = self.conv2d.weight
        bias = self.conv2d.bias
        if x.is_cuda and x.dtype == torch.float32 and weight.dtype == torch.float32:
            if bias is None:
                bias = torch.empty(0, device=x.device, dtype=torch.float32)
            return _depthwise_conv_ext.depthwise_conv2d_cuda(
                x, weight, bias,
                self.conv2d.stride[0], self.conv2d.padding[0])
        return self.conv2d(x)
```