```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 (3x3, stride 1).
# Strategy: one thread block computes a 32x32 output tile for a single
# (batch, channel) pair. The corresponding 34x34 input tile (halo included)
# is staged in shared memory so each input element is fetched from global
# memory essentially once, making the kernel DRAM-bandwidth bound (optimal
# for a memory-bound depthwise convolution).
depthwise_conv2d_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 32

__global__ void depthwise_conv2d_3x3_s1_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int C_in,
    const int C_out,
    const int H,
    const int W,
    const int H_out,
    const int W_out,
    const int pad,
    const int channel_multiplier,
    const bool has_bias)
{
    __shared__ float tile[TILE_H + 2][TILE_W + 2];

    const int bc    = blockIdx.z;          // in [0, N * C_out)
    const int c_out = bc % C_out;
    const int b     = bc / C_out;
    const int c_in  = c_out / channel_multiplier;

    const int h0 = blockIdx.y * TILE_H;
    const int w0 = blockIdx.x * TILE_W;

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

    const float* __restrict__ x_bc =
        x + (((long)b * C_in + c_in) * (long)H) * (long)W;

    const int tid        = ty * TILE_W + tx;
    const int nthreads   = TILE_W * TILE_H;
    const int tile_elems = (TILE_H + 2) * (TILE_W + 2);

    // Stage the input tile (with halo) into shared memory. Out-of-bounds
    // locations contribute zeros (implicit padding).
    for (int i = tid; i < tile_elems; i += nthreads) {
        const int th = i / (TILE_W + 2);
        const int tw = i - th * (TILE_W + 2);
        const int ih = h0 + th - pad;
        const int iw = w0 + tw - pad;
        float v = 0.0f;
        if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) {
            v = x_bc[(long)ih * W + iw];
        }
        tile[th][tw] = v;
    }
    __syncthreads();

    const int h_out = h0 + ty;
    const int w_out = w0 + tx;
    if (h_out < H_out && w_out < W_out) {
        const float* __restrict__ w_c = w + (long)c_out * 9;
        float acc = has_bias ? bias[c_out] : 0.0f;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                acc = fmaf(tile[ty + kh][tx + kw], w_c[kh * 3 + kw], acc);
            }
        }
        out[((long)bc * H_out + h_out) * W_out + w_out] = acc;
    }
}

torch::Tensor depthwise_conv2d_forward(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    int64_t stride,
    int64_t padding)
{
    const int64_t N     = x.size(0);
    const int64_t C_in  = x.size(1);
    const int64_t H     = x.size(2);
    const int64_t W     = x.size(3);
    const int64_t C_out = weight.size(0);
    const int64_t K     = weight.size(2);

    const bool has_bias = bias_opt.has_value() && bias_opt->defined();

    const bool supported =
        x.is_cuda() && weight.is_cuda() &&
        x.scalar_type() == at::kFloat &&
        weight.scalar_type() == at::kFloat &&
        K == 3 && weight.size(3) == 3 &&
        weight.size(1) == 1 &&
        stride == 1 &&
        C_in > 0 && (C_out % C_in) == 0 &&
        (!has_bias || bias_opt->scalar_type() == at::kFloat);

    // Graceful fallback to the ATen implementation for configurations the
    // custom kernel does not cover.
    if (!supported) {
        torch::Tensor bias_t = has_bias ? *bias_opt : torch::Tensor();
        return at::conv2d(x, weight, bias_t, {stride, stride},
                          {padding, padding}, {1, 1}, C_in);
    }

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    torch::Tensor b_c;
    const float* bias_ptr = nullptr;
    if (has_bias) {
        b_c = bias_opt->contiguous();
        bias_ptr = b_c.data_ptr<float>();
    }

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

    auto out = torch::empty({N, C_out, H_out, W_out}, x.options());
    if (H_out <= 0 || W_out <= 0 || N == 0 || C_out == 0) {
        return out;
    }

    dim3 block(TILE_W, TILE_H);
    dim3 grid(
        (unsigned)((W_out + TILE_W - 1) / TILE_W),
        (unsigned)((H_out + TILE_H - 1) / TILE_H),
        (unsigned)(N * C_out));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    depthwise_conv2d_3x3_s1_kernel<<<grid, block, 0, stream>>>(
        x_c.data_ptr<float>(),
        w_c.data_ptr<float>(),
        bias_ptr,
        out.data_ptr<float>(),
        (int)C_in, (int)C_out, (int)H, (int)W,
        (int)H_out, (int)W_out,
        (int)padding,
        (int)(C_out / C_in),
        has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

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

# Compile the inline CUDA extension.
_depthwise_conv2d_ext = 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_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom shared-memory-tiled
    CUDA kernel (specialized for 3x3 kernels with stride 1, with an automatic
    fallback to the ATen implementation 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,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _depthwise_conv2d_ext.depthwise_conv2d_forward(
            x,
            self.conv2d.weight,
            self.conv2d.bias,
            self.conv2d.stride[0],
            self.conv2d.padding[0],
        )
```