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

depthwise_conv_cpp_source = """
#include <torch/extension.h>
torch::Tensor depthwise_conv2d_cuda(torch::Tensor input, torch::Tensor weight, c10::optional<torch::Tensor> bias, int64_t stride, int64_t padding);
"""

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

#define TILE 16

// ---------------------------------------------------------------------------
// Optimized shared-memory tiled kernel for the common case:
//   kernel_size = 3, stride = 1, channel multiplier = 1 (out_channels == in_channels)
// Each block computes a TILE x TILE output tile for one (batch, channel) pair.
// The (TILE+2) x (TILE+2) input patch and the 3x3 weights are cached in
// shared memory, so each input element is read from global memory ~once.
// ---------------------------------------------------------------------------
__global__ void depthwise_conv2d_k3s1_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,
    const int padding, const bool has_bias)
{
    const int SMEM = TILE + 2;

    __shared__ float smem[TILE + 2][TILE + 2];
    __shared__ float sweight[9];

    const int nc = blockIdx.z;              // n * channels + c
    const int c = nc % channels;

    // Cache the 3x3 kernel weights for this channel.
    if (threadIdx.y == 0 && threadIdx.x < 9) {
        sweight[threadIdx.x] = weight[c * 9 + threadIdx.x];
    }

    const int h_out = blockIdx.y * TILE + threadIdx.y;
    const int w_out = blockIdx.x * TILE + threadIdx.x;

    const int h_in_base = blockIdx.y * TILE - padding;
    const int w_in_base = blockIdx.x * TILE - padding;

    const float* in_nc = input + (long)nc * height_in * width_in;

    // Cooperatively load the input tile (zero-filled outside the image).
    for (int i = threadIdx.y; i < SMEM; i += TILE) {
        const int h_in = h_in_base + i;
        const bool row_ok = (h_in >= 0 && h_in < height_in);
        for (int j = threadIdx.x; j < SMEM; j += TILE) {
            const int w_in = w_in_base + j;
            float val = 0.0f;
            if (row_ok && w_in >= 0 && w_in < width_in) {
                val = in_nc[h_in * width_in + w_in];
            }
            smem[i][j] = val;
        }
    }
    __syncthreads();

    if (h_out < height_out && w_out < width_out) {
        float sum = 0.0f;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                sum += smem[threadIdx.y + kh][threadIdx.x + kw] * sweight[kh * 3 + kw];
            }
        }
        if (has_bias) sum += bias[c];
        output[(long)nc * height_out * width_out + h_out * width_out + w_out] = sum;
    }
}

// ---------------------------------------------------------------------------
// Generic fallback: supports any kernel size, stride, padding and multiplier.
// One thread per output element.
// ---------------------------------------------------------------------------
__global__ void depthwise_conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int in_channels, const int multiplier,
    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 bool has_bias, const long total)
{
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int out_channels = in_channels * multiplier;
    const int w_out = idx % width_out;
    const int h_out = (idx / width_out) % height_out;
    const int o = (idx / ((long)width_out * height_out)) % out_channels;
    const int n = idx / ((long)width_out * height_out * out_channels);
    const int c = o / multiplier;

    const int h_in_base = h_out * stride - padding;
    const int w_in_base = w_out * stride - padding;

    const float* in_nc = input + ((long)n * in_channels + c) * height_in * width_in;
    const float* w_o = weight + (long)o * kernel_size * kernel_size;

    float sum = 0.0f;
    for (int kh = 0; kh < kernel_size; ++kh) {
        const int h_in = h_in_base + kh;
        if (h_in < 0 || h_in >= height_in) continue;
        const float* in_row = in_nc + (long)h_in * width_in;
        const float* w_row = w_o + kh * kernel_size;
        for (int kw = 0; kw < kernel_size; ++kw) {
            const int w_in = w_in_base + kw;
            if (w_in < 0 || w_in >= width_in) continue;
            sum += in_row[w_in] * w_row[kw];
        }
    }
    if (has_bias) sum += bias[o];
    output[idx] = sum;
}

torch::Tensor depthwise_conv2d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<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 4D (N, C, H, W)");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");

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

    const int batch = input_c.size(0);
    const int in_channels = input_c.size(1);
    const int height_in = input_c.size(2);
    const int width_in = input_c.size(3);

    const int out_channels = weight_c.size(0);
    const int kernel_size = weight_c.size(2);
    const int multiplier = out_channels / in_channels;

    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, out_channels, height_out, width_out}, input_c.options());

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

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

    if (kernel_size == 3 && stride == 1 && multiplier == 1 && (long)batch * in_channels <= 65535) {
        dim3 block(TILE, TILE);
        dim3 grid((width_out + TILE - 1) / TILE,
                  (height_out + TILE - 1) / TILE,
                  batch * in_channels);
        depthwise_conv2d_k3s1_kernel<<<grid, block, 0, stream>>>(
            input_c.data_ptr<float>(),
            weight_c.data_ptr<float>(),
            bias_ptr,
            output.data_ptr<float>(),
            in_channels, height_in, width_in,
            height_out, width_out,
            (int)padding, has_bias);
    } else {
        const long total = (long)batch * out_channels * height_out * width_out;
        const int block = 256;
        const long grid = (total + block - 1) / block;
        depthwise_conv2d_generic_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            input_c.data_ptr<float>(),
            weight_c.data_ptr<float>(),
            bias_ptr,
            output.data_ptr<float>(),
            in_channels, multiplier,
            height_in, width_in,
            height_out, width_out,
            kernel_size, (int)stride, (int)padding,
            has_bias, total);
    }

    return output;
}
"""

depthwise_conv = load_inline(
    name="depthwise_conv2d_ext",
    cpp_sources=depthwise_conv_cpp_source,
    cuda_sources=depthwise_conv_cuda_source,
    functions=["depthwise_conv2d_cuda"],
    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 specialized for the common
    kernel_size=3 / stride=1 / multiplier=1 case, with a fully generic
    depthwise kernel as fallback for other configurations.
    """
    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:
        if x.is_cuda and x.dtype == torch.float32:
            return depthwise_conv.depthwise_conv2d_cuda(
                x, self.conv2d.weight, self.conv2d.bias, self.stride, self.padding
            )
        return self.conv2d(x)
```