```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.
//
// Memory-streaming design: the op is bandwidth-bound (it reads the input once
// and writes the output once), so the kernel is built to keep HBM busy
// continuously. Each block owns a fixed 128x64 output tile and loops over
// several (batch, channel) planes. The 130x66 input halo of the *next* plane
// is prefetched into a second shared-memory buffer with 16-byte cp.async
// copies (zero-filling out-of-range tails) while the current plane is being
// computed and stored, so loads, compute and stores of consecutive planes
// overlap. Input streams are marked L2 evict-first (they are read once),
// output stores are streaming (evict-first) float4; when an output row base
// is not 16B-aligned (OW == 2 mod 4, i.e. 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 = 16;
constexpr int THREADS = TX * TY;               // 512
constexpr int TILE_W = TX * 4;                 // 128 output cols per block
constexpr int TILE_H = TY * 4;                 // 64 output rows per block
constexpr int SM_H = TILE_H + 2;               // 66
constexpr int PITCH = TILE_W + 4;              // 132 (SM_W = 130 rounded to /4)
constexpr int VPR = PITCH / 4;                 // 33 float4 per shared row
constexpr int TILE_FLOATS = SM_H * PITCH;      // 8712
constexpr int SMEM_BYTES = 2 * TILE_FLOATS * (int)sizeof(float);  // 69696
constexpr int FILL_ITERS = (SM_H * VPR + THREADS - 1) / THREADS;  // 5

__global__ void __launch_bounds__(THREADS)
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,
    int zBlocks, bool has_bias) {

    extern __shared__ __align__(16) float smem[];
    float* tile0 = smem;
    float* tile1 = smem + TILE_FLOATS;

    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 int64_t plane_in = (int64_t)H * W;
    const int64_t plane_out = (int64_t)OH * OW;
    const int total_z = N * OC;

    int z = blockIdx.z;
    if (z >= total_z) return;

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    // Inputs are streamed exactly once: mark them evict-first in L2.
    unsigned long long pol;
    asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(pol));

    // Asynchronously stage the shared tile for plane `zz` into buffer `buf`.
    auto fill = [&](float* buf, int zz) {
        const int n = zz / OC;
        const int oc = zz - n * OC;
        const int ic = oc / mult;
        const float* __restrict__ xb = x + ((int64_t)n * C + ic) * plane_in;
        #pragma unroll
        for (int it = 0; it < FILL_ITERS; ++it) {
            const int idx = tid + it * THREADS;
            if (idx < SM_H * VPR) {
                const int sy = idx / VPR;
                const int c4 = idx - sy * VPR;
                const int iy = oh0 + sy;
                const int ix = ow0 + (c4 << 2);
                const float* src = xb + (int64_t)iy * W + ix;
                int nbytes = 0;
                if (iy < H) {
                    const int rem = W - ix;  // floats remaining in this row
                    nbytes = rem >= 4 ? 16 : (rem > 0 ? rem << 2 : 0);
                }
                if (nbytes == 0) {
                    src = xb;  // safe address; nothing is read anyway
                }
                const unsigned dst =
                    (unsigned)__cvta_generic_to_shared(&buf[sy * PITCH + (c4 << 2)]);
                asm volatile(
                    "cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2, %3;\n" ::
                        "r"(dst), "l"(src), "r"(nbytes), "l"(pol));
            }
        }
        asm volatile("cp.async.commit_group;\n" ::);
    };
#else
    auto fill = [&](float* buf, int zz) {
        const int n = zz / OC;
        const int oc = zz - n * OC;
        const int ic = oc / mult;
        const float* __restrict__ xb = x + ((int64_t)n * C + ic) * plane_in;
        #pragma unroll
        for (int it = 0; it < FILL_ITERS; ++it) {
            const int idx = tid + it * THREADS;
            if (idx < SM_H * VPR) {
                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*>(&buf[sy * PITCH + (c4 << 2)]) = v;
            }
        }
    };
#endif

    // Prologue: stage the first plane.
    fill(tile0, z);
    float* cur = tile0;
    float* nxt = tile1;

    for (; z < total_z; z += zBlocks) {
        const int n = z / OC;
        const int oc = z - n * OC;

        // Issue the weight/bias loads early; their latency hides under the
        // cp.async wait below.
        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);
            }
        }
        float bv = 0.0f;
        if (has_bias) {
            bv = __ldg(bias + oc);
        }

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
        asm volatile("cp.async.wait_group 0;\n" ::);
#endif
        // All threads are done reading `nxt` (previous plane) and the tile in
        // `cur` is fully staged.
        __syncthreads();

        const int zn = z + zBlocks;
        if (zn < total_z) {
            fill(nxt, zn);  // overlaps with the compute/store below
        }

        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 = &cur[(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) {
            #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 must take part in the shuffle; it is used 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). Shift every lane's (a2, a3) one lane
                // up so lanes 1..31 write aligned float4 chunks; 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);
                }
            }
        }

        float* tmp = cur;
        cur = nxt;
        nxt = tmp;
    }
}

} // 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>();
    }

    // Pick the z extent of the grid: keep it a multiple of OC when possible so
    // every z-slice a block visits shares the same channel, and give each
    // block several planes so the prefetch pipeline amortizes.
    int64_t groups = 1;
    while (groups < N && groups < 8 && OC * (groups * 2) <= 65535) {
        groups *= 2;
    }
    int64_t zBlocks;
    if (OC <= 65535) {
        zBlocks = OC * groups;  // multiple of OC, <= 65535
    } else {
        zBlocks = std::min<int64_t>(total_z, 65535);
    }
    zBlocks = std::min<int64_t>(zBlocks, total_z);

    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(depthwise_conv3x3_s1_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             SMEM_BYTES);
        attr_set = true;
    }

    dim3 block(TX, TY);
    dim3 grid(
        (unsigned int)((OW + TILE_W - 1) / TILE_W),
        (unsigned int)((OH + TILE_H - 1) / TILE_H),
        (unsigned int)zBlocks);

    auto stream = at::cuda::getCurrentCUDAStream();
    depthwise_conv3x3_s1_kernel<<<grid, block, SMEM_BYTES, 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,
        (int)zBlocks,
        has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_depthwise_conv_ext = load_inline(
    name="depthwise_conv3x3_s1_v5_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 + 63) // 64) <= 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)
```