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

_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>
#include <cstdint>

// Depthwise 3x3 convolution, stride 1, padding 0, float32.
//
// Each thread block computes a 128x32 output tile for one (batch, channel)
// pair. Block = (32, 8) threads; each thread produces a 4(col) x 4(row)
// patch of outputs. The 130x34 input halo is staged in shared memory through
// aligned float4 global loads. Compute reads shared memory via conflict-free
// float4 accesses and keeps the 3x3 weights in registers, so the inner loop
// is almost pure FMA. Output stores use float4; when a row base is not
// 16B-aligned (OW is 2 mod 4, so exactly every other row), warp neighbours
// exchange halves with __shfl_up_sync so all vectorized stores stay aligned.
namespace {

constexpr int TX = 32;
constexpr int TY = 8;
constexpr int TILE_W = TX * 4;   // 128 output cols per block
constexpr int TILE_H = TY * 4;   // 32 output rows per block
constexpr int SM_W = TILE_W + 2; // 130
constexpr int SM_H = TILE_H + 2; // 34
constexpr int PITCH = 132;       // SM_W rounded up to a multiple of 4
constexpr int VPR = PITCH / 4;   // 33 float4 per shared row

__global__ void __launch_bounds__(TX * TY)
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) {

    __shared__ __align__(16) 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;

        // Stage the SM_H x PITCH shared tile (last 2 columns are pad).
        #pragma unroll
        for (int idx = tid; idx < SM_H * VPR; idx += TX * TY) {
            const int sy = idx / VPR;
            const int c4 = idx - sy * VPR;
            const int iy = oh0 + sy;
            const int ix = ow0 + (c4 << 2);
            float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            if (iy < H) {
                const float* row = xb + (int64_t)iy * W;
                if (ix + 3 < W) {
                    v = *reinterpret_cast<const float4*>(row + ix);
                } else {
                    if (ix + 0 < W) v.x = row[ix + 0];
                    if (ix + 1 < W) v.y = row[ix + 1];
                    if (ix + 2 < W) v.z = row[ix + 2];
                    if (ix + 3 < W) v.w = row[ix + 3];
                }
            }
            *reinterpret_cast<float4*>(&tile[sy * PITCH + (c4 << 2)]) = v;
        }
        __syncthreads();

        float wr[9];
        #pragma unroll
        for (int k = 0; k < 9; ++k) {
            wr[k] = ws[k];
        }

        const int c0 = tx << 2;   // shared col base
        const int r0 = ty << 2;   // shared row base (uses rows r0..r0+5)

        float acc[4][4];
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            #pragma unroll
            for (int c = 0; c < 4; ++c) {
                acc[j][c] = 0.0f;
            }
        }

        #pragma unroll
        for (int i = 0; i < 6; ++i) {
            const float* row = &tile[(r0 + i) * PITCH + c0];
            const float4 a = *reinterpret_cast<const float4*>(row);
            const float4 b = *reinterpret_cast<const float4*>(row + 4);
            const float s0 = a.x, s1 = a.y, s2 = a.z, s3 = a.w, s4 = b.x, s5 = b.y;
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                if (i >= j && i <= j + 2) {
                    const int ky = i - j;
                    const float w0 = wr[ky * 3 + 0], w1 = wr[ky * 3 + 1], w2 = wr[ky * 3 + 2];
                    acc[j][0] += s0 * w0 + s1 * w1 + s2 * w2;
                    acc[j][1] += s1 * w0 + s2 * w1 + s3 * w2;
                    acc[j][2] += s2 * w0 + s3 * w1 + s4 * w2;
                    acc[j][3] += s3 * w0 + s4 * w1 + s5 * w2;
                }
            }
        }

        if (has_bias) {
            const float bv = bias[oc];
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                #pragma unroll
                for (int c = 0; c < 4; ++c) {
                    acc[j][c] += bv;
                }
            }
        }

        float* __restrict__ ob = out + ((int64_t)n * OC + oc) * (int64_t)OH * OW;
        const int ow = ow0 + c0;
        const bool interior = (ow0 + TILE_W <= OW);

        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int oh = oh0 + r0 + j;
            if (oh >= OH) continue;
            float* p = ob + (int64_t)oh * OW + ow;
            if (interior) {
                if ((((uintptr_t)p) & 15) == 0) {
                    *reinterpret_cast<float4*>(p) =
                        make_float4(acc[j][0], acc[j][1], acc[j][2], acc[j][3]);
                } else {
                    // Row base is misaligned by exactly 2 floats (W % 4 == 0
                    // implies OW == 2 mod 4). Shift every lane's (a2, a3)
                    // one lane up so lanes 1..31 write aligned float4 chunks;
                    // lanes 0 and 31 patch the two remaining columns.
                    const float a0 = acc[j][0], a1 = acc[j][1];
                    const float a2 = acc[j][2], a3 = acc[j][3];
                    const float pz = __shfl_up_sync(0xffffffffu, a2, 1);
                    const float pw = __shfl_up_sync(0xffffffffu, a3, 1);
                    if (tx > 0) {
                        *reinterpret_cast<float4*>(p - 2) = make_float4(pz, pw, a0, a1);
                    } else {
                        p[0] = a0;
                        p[1] = a1;
                    }
                    if (tx == TX - 1) {
                        p[2] = a2;
                        p[3] = a3;
                    }
                }
            } else {
                #pragma unroll
                for (int c = 0; c < 4; ++c) {
                    if (ow + c < OW) p[c] = acc[j][c];
                }
            }
        }
        __syncthreads();
    }
}

} // namespace

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");
    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");
    TORCH_CHECK(W % 4 == 0, "input width must be a multiple of 4");
    TORCH_CHECK((OH + TILE_H - 1) / TILE_H <= 65535, "input height too large");
    TORCH_CHECK((reinterpret_cast<uintptr_t>(xc.data_ptr<float>()) & 15) == 0,
                "input data pointer must be 16-byte aligned");

    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(TX, TY);
    dim3 grid(
        (unsigned int)((OW + TILE_W - 1) / TILE_W),
        (unsigned int)((OH + TILE_H - 1) / TILE_H),
        (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_v3_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.dim() == 4
            and x.size(2) >= 3
            and x.size(3) >= 4
            and (x.size(3) & 3) == 0
            and (x.data_ptr() & 15) == 0
            and ((x.size(2) - 2 + 31) // 32) <= 65535
        ):
            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)
```