```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 = (
    "torch::Tensor depthwise_conv3x3_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);"
)

_depthwise_conv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <algorithm>

// Depthwise 3x3 convolution, stride 1, padding 0, float32.
// Each thread block computes a 32x32 output tile for one (batch, out-channel)
// pair. The 34x34 input halo is staged in shared memory and each of the
// 16x16 threads produces a 2x2 patch of outputs.
__global__ void depthwise_conv3x3_s1_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int C, int H, int W,
    int OC, int mult, int OH, int OW,
    bool has_bias) {

    constexpr int TX = 16;
    constexpr int TY = 16;
    constexpr int TILE_W = TX * 2;   // 32 output columns per block
    constexpr int TILE_H = TY * 2;   // 32 output rows per block
    constexpr int SM_W = TILE_W + 2; // 34
    constexpr int SM_H = TILE_H + 2; // 34
    constexpr int PITCH = SM_W + 1;  // 35, avoids shared-memory bank conflicts

    __shared__ float tile[SM_H * PITCH];
    __shared__ float ws[9];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * TX + tx;

    const int ow0 = blockIdx.x * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;

    const int total_z = N * OC;
    for (int z = blockIdx.z; z < total_z; z += gridDim.z) {
        const int n = z / OC;
        const int oc = z - n * OC;
        const int ic = oc / mult;

        if (tid < 9) {
            ws[tid] = w[oc * 9 + tid];
        }

        const float* __restrict__ xb = x + ((int64_t)n * C + ic) * (int64_t)H * W;

        for (int idx = tid; idx < SM_W * SM_H; idx += TX * TY) {
            const int sy = idx / SM_W;
            const int sx = idx - sy * SM_W;
            const int iy = oh0 + sy;
            const int ix = ow0 + sx;
            float v = 0.0f;
            if (iy < H && ix < W) {
                v = __ldg(&xb[(int64_t)iy * W + ix]);
            }
            tile[sy * PITCH + sx] = v;
        }
        __syncthreads();

        const int bx = tx * 2;
        const int by = ty * 2;

        #pragma unroll
        for (int dy = 0; dy < 2; ++dy) {
            const int oh = oh0 + by + dy;
            if (oh >= OH) continue;
            #pragma unroll
            for (int dx = 0; dx < 2; ++dx) {
                const int ow = ow0 + bx + dx;
                if (ow >= OW) continue;
                const float* __restrict__ t = &tile[(by + dy) * PITCH + bx + dx];
                float acc = 0.0f;
                #pragma unroll
                for (int ky = 0; ky < 3; ++ky) {
                    #pragma unroll
                    for (int kx = 0; kx < 3; ++kx) {
                        acc += t[ky * PITCH + kx] * ws[ky * 3 + kx];
                    }
                }
                if (has_bias) {
                    acc += bias[oc];
                }
                out[((int64_t)n * OC + oc) * (int64_t)OH * OW + (int64_t)oh * OW + ow] = acc;
            }
        }
        __syncthreads();
    }
}

torch::Tensor depthwise_conv3x3_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "depthwise_conv3x3_cuda expects CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "depthwise_conv3x3_cuda expects float32 tensors");
    TORCH_CHECK(x.dim() == 4 && weight.dim() == 4, "bad tensor ranks");
    TORCH_CHECK(weight.size(1) == 1 && weight.size(2) == 3 && weight.size(3) == 3,
                "expected depthwise 3x3 weights of shape (OC, 1, 3, 3)");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t H = xc.size(2);
    const int64_t W = xc.size(3);
    const int64_t OC = wc.size(0);
    TORCH_CHECK(OC % C == 0, "out_channels must be a multiple of in_channels for depthwise conv");
    const int64_t mult = OC / C;
    const int64_t OH = H - 2;
    const int64_t OW = W - 2;
    TORCH_CHECK(OH > 0 && OW > 0, "input spatial size must be >= 3");

    auto out = at::empty({N, OC, OH, OW}, xc.options());

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

    dim3 block(16, 16);
    dim3 grid(
        (unsigned int)((OW + 31) / 32),
        (unsigned int)((OH + 31) / 32),
        (unsigned int)std::min<int64_t>(N * OC, 65535));

    auto stream = at::cuda::getCurrentCUDAStream();
    depthwise_conv3x3_s1_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(),
        wc.data_ptr<float>(),
        bias_ptr,
        out.data_ptr<float>(),
        (int)N, (int)C, (int)H, (int)W,
        (int)OC, (int)mult, (int)OH, (int)OW,
        has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_depthwise_conv_ext = load_inline(
    name="depthwise_conv3x3_s1_ext",
    cpp_sources=_depthwise_conv_cpp_source,
    cuda_sources=_depthwise_conv_cuda_source,
    functions=["depthwise_conv3x3_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom fused CUDA kernel for
    the common 3x3 / stride 1 / no-padding float32 case. Falls back to the
    reference nn.Conv2d implementation for every other 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._use_custom_kernel = (kernel_size == 3 and stride == 1 and padding == 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (
            self._use_custom_kernel
            and x.is_cuda
            and x.dtype == torch.float32
            and x.size(2) >= 3
            and x.size(3) >= 3
        ):
            bias = self.conv2d.bias
            if bias is None:
                bias = x.new_empty(0)
            return _depthwise_conv_ext.depthwise_conv3x3_cuda(x, self.conv2d.weight, bias)
        return self.conv2d(x)
```