```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.
//
// The operator is purely bandwidth-bound (input is read once, output written
// once), so this kernel is engineered to stream from HBM continuously:
//
//  * No shared memory and no __syncthreads: each thread owns a 4(col) x 4(row)
//    patch of outputs and loads its 6 input rows directly through L1 with
//    16-byte float4 loads (plus two scalar loads for the right halo). The
//    unrolled row loop gives every thread many independent loads in flight,
//    and 256-thread blocks reach high occupancy (no smem limit, low reg use).
//    Halo reuse between neighboring threads/blocks is captured by L1/L2.
//  * Loads only touch memory needed by valid outputs; rows are clamped and
//    columns predicated so edge tiles stay safe, while interior tiles run a
//    single fast vectorized path.
//  * Output stores are streaming (evict-first) float4. Because OW == 2 mod 4
//    (W % 4 == 0), every other output row is 8-byte misaligned; warp neighbors
//    exchange their upper halves with __shfl_up_sync so lanes 1..31 emit
//    aligned float4 stores anyway, with lanes 0/31 patching the leftovers.
namespace {

constexpr int TX = 32;
constexpr int TY = 8;
constexpr int THREADS = TX * TY;      // 256
constexpr int TILE_W = TX * 4;        // 128 output cols per block
constexpr int TILE_H = TY * 4;        // 32 output rows per block

__global__ void __launch_bounds__(THREADS, 6)
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) {

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

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

    const int64_t plane_in  = (int64_t)H * W;
    const int64_t plane_out = (int64_t)OH * OW;
    const int total_z = N * OC;

    const int c0 = tx << 2;   // col base within tile
    const int r0 = ty << 2;   // row base within tile
    const int cx = ow0 + c0;  // global input col base

    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;

        float wr[9];
        {
            const float* __restrict__ wp = w + (int64_t)oc * 9;
            #pragma unroll
            for (int k = 0; k < 9; ++k) {
                wr[k] = __ldg(wp + k);
            }
        }

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

        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;
            }
        }

        // Sweep the 6 input rows that feed this thread's 4 output rows.
        #pragma unroll
        for (int i = 0; i < 6; ++i) {
            const int iy = oh0 + r0 + i;
            const int iyc = iy < H ? iy : H - 1;  // clamp: only feeds invalid outputs
            const float* __restrict__ row = xb + (int64_t)iyc * W;
            float s0, s1, s2, s3, s4, s5;
            if (cx + 3 < W) {
                const float4 v = __ldg(reinterpret_cast<const float4*>(row + cx));
                s0 = v.x; s1 = v.y; s2 = v.z; s3 = v.w;
            } else {
                s0 = (cx + 0 < W) ? row[cx + 0] : 0.0f;
                s1 = (cx + 1 < W) ? row[cx + 1] : 0.0f;
                s2 = (cx + 2 < W) ? row[cx + 2] : 0.0f;
                s3 = (cx + 3 < W) ? row[cx + 3] : 0.0f;
            }
            s4 = (cx + 4 < W) ? row[cx + 4] : 0.0f;
            s5 = (cx + 5 < W) ? row[cx + 5] : 0.0f;

            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                const int ky = i - j;
                if (ky >= 0 && ky < 3) {
                    const float w0 = wr[ky * 3 + 0];
                    const float w1 = wr[ky * 3 + 1];
                    const float 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 = __ldg(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) * plane_out;
        const int ow = ow0 + c0;

        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int oh = oh0 + r0 + j;
            if (oh >= OH) continue;  // warp-uniform (warp == one ty row)
            float* p = ob + (int64_t)oh * OW + ow;
            const float a0 = acc[j][0], a1 = acc[j][1];
            const float a2 = acc[j][2], a3 = acc[j][3];
            // All lanes participate in the shuffle; it is consumed only on the
            // misaligned path below.
            const float pz = __shfl_up_sync(0xffffffffu, a2, 1);
            const float pw = __shfl_up_sync(0xffffffffu, a3, 1);
            if (ow >= OW) continue;
            if ((((uintptr_t)p) & 15) == 0) {
                if (ow + 3 < OW) {
                    __stcs(reinterpret_cast<float4*>(p), make_float4(a0, a1, a2, a3));
                } else {
                    if (ow + 0 < OW) __stcs(p + 0, a0);
                    if (ow + 1 < OW) __stcs(p + 1, a1);
                    if (ow + 2 < OW) __stcs(p + 2, a2);
                    if (ow + 3 < OW) __stcs(p + 3, a3);
                }
            } else {
                // Row base is misaligned by exactly 2 floats (W % 4 == 0
                // implies OW == 2 mod 4). Lanes 1..31 store the aligned float4
                // chunk covering cols ow-2..ow+1 (neighbor's a2,a3 then own
                // a0,a1); lanes 0 and 31 patch the remaining columns.
                if (tx > 0) {
                    if (ow + 1 < OW) {
                        __stcs(reinterpret_cast<float4*>(p - 2), make_float4(pz, pw, a0, a1));
                    } else {
                        if (ow - 2 < OW) __stcs(p - 2, pz);
                        if (ow - 1 < OW) __stcs(p - 1, pw);
                        if (ow + 0 < OW) __stcs(p + 0, a0);
                        if (ow + 1 < OW) __stcs(p + 1, a1);
                    }
                } else {
                    if (ow + 0 < OW) __stcs(p + 0, a0);
                    if (ow + 1 < OW) __stcs(p + 1, a1);
                }
                if (tx == TX - 1) {
                    if (ow + 2 < OW) __stcs(p + 2, a2);
                    if (ow + 3 < OW) __stcs(p + 3, a3);
                }
            }
        }
    }
}

} // 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 int64_t total_z = N * OC;
    if (total_z == 0) {
        return out;
    }

    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>(total_z, 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_v6_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)
```