```
# 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_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define TILE_OUT 32
#define BLOCK_DIM 16
#define SMEM_SIZE (TILE_OUT + 2)

// Fast path: 3x3 kernel, stride 1, padding 0, out_channels == in_channels.
// Each block computes a 32x32 output tile for one (batch, channel) pair.
// The (34x34) input halo tile is staged in shared memory, and each thread
// computes a 2x2 patch of outputs using a 4x4 register tile for data reuse.
__global__ __launch_bounds__(BLOCK_DIM * BLOCK_DIM)
void depthwise_conv2d_3x3_s1_p0_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int B, const int C,
    const int H, const int W,
    const int OH, const int OW,
    const bool has_bias)
{
    __shared__ float tile[SMEM_SIZE][SMEM_SIZE + 1];
    __shared__ float wsh[9];

    const int bc = blockIdx.z;          // b * C + c
    const int c  = bc % C;

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

    if (ty == 0 && tx < 9) {
        wsh[tx] = weight[c * 9 + tx];
    }

    const int h_start = blockIdx.y * TILE_OUT;
    const int w_start = blockIdx.x * TILE_OUT;

    const float* in_base = input + (long)bc * H * W;

    // Stage (TILE_OUT+2) x (TILE_OUT+2) input tile into shared memory.
    for (int i = ty; i < SMEM_SIZE; i += BLOCK_DIM) {
        const int h_in = h_start + i;
        const bool hv = (h_in < H);
        const float* row = in_base + (long)h_in * W;
        for (int j = tx; j < SMEM_SIZE; j += BLOCK_DIM) {
            const int w_in = w_start + j;
            float v = 0.0f;
            if (hv && (w_in < W)) {
                v = row[w_in];
            }
            tile[i][j] = v;
        }
    }
    __syncthreads();

    const int oh0 = blockIdx.y * TILE_OUT + ty * 2;
    const int ow0 = blockIdx.x * TILE_OUT + tx * 2;
    if (oh0 >= OH || ow0 >= OW) return;

    // 4x4 register tile reused across the 2x2 outputs.
    float r[4][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            r[i][j] = tile[ty * 2 + i][tx * 2 + j];
        }
    }

    const float bv = has_bias ? bias[c] : 0.0f;
    float s00 = bv, s01 = bv, s10 = bv, s11 = bv;
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
            const float wv = wsh[kh * 3 + kw];
            s00 += r[kh][kw]         * wv;
            s01 += r[kh][kw + 1]     * wv;
            s10 += r[kh + 1][kw]     * wv;
            s11 += r[kh + 1][kw + 1] * wv;
        }
    }

    float* out = output + ((long)bc * OH + oh0) * OW + ow0;
    const bool wv1 = (ow0 + 1) < OW;
    const bool hv1 = (oh0 + 1) < OH;
    out[0] = s00;
    if (wv1) out[1] = s01;
    if (hv1) {
        out[OW] = s10;
        if (wv1) out[OW + 1] = s11;
    }
}

// Generic fallback: any kernel size / stride / padding / channel multiplier.
__global__ void depthwise_conv2d_general_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int B, const int IC, const int OC,
    const int H, const int W,
    const int OH, const int OW,
    const int KH, const int KW,
    const int SH, const int SW,
    const int PH, const int PW,
    const bool has_bias)
{
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long total = (long)B * OC * OH * OW;
    if (idx >= total) return;

    const int mult = OC / IC;
    const int ow = (int)(idx % OW);
    const int oh = (int)((idx / OW) % OH);
    const int c  = (int)((idx / ((long)OW * OH)) % OC);
    const int b  = (int)(idx / ((long)OW * OH * OC));
    const int ic = c / mult;

    const int h0 = oh * SH - PH;
    const int w0 = ow * SW - PW;

    const float* in_base = input  + ((long)b * IC + ic) * H * W;
    const float* w_base  = weight + (long)c * KH * KW;

    float sum = has_bias ? bias[c] : 0.0f;
    for (int kh = 0; kh < KH; ++kh) {
        const int h = h0 + kh;
        if (h < 0 || h >= H) continue;
        for (int kw = 0; kw < KW; ++kw) {
            const int wcol = w0 + kw;
            if (wcol < 0 || wcol >= W) continue;
            sum += in_base[(long)h * W + wcol] * w_base[kh * KW + kw];
        }
    }
    output[idx] = sum;
}

torch::Tensor depthwise_conv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t stride_h, int64_t stride_w,
    int64_t pad_h, int64_t pad_w)
{
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "input/weight must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 is supported");

    auto x = input.contiguous();
    auto w = weight.contiguous();

    const int B  = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int H  = (int)x.size(2);
    const int W  = (int)x.size(3);
    const int OC = (int)w.size(0);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);
    TORCH_CHECK(IC > 0 && OC % IC == 0, "invalid depthwise convolution shapes");

    const int OH = (int)((H + 2 * pad_h - KH) / stride_h + 1);
    const int OW = (int)((W + 2 * pad_w - KW) / stride_w + 1);

    auto out = at::empty({B, OC, OH, OW}, x.options());
    const long total = (long)B * OC * OH * OW;
    if (total <= 0) return out;

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (KH == 3 && KW == 3 && stride_h == 1 && stride_w == 1 &&
        pad_h == 0 && pad_w == 0 && OC == IC && (long)B * OC <= 65535) {
        const dim3 block(BLOCK_DIM, BLOCK_DIM);
        const dim3 grid((OW + TILE_OUT - 1) / TILE_OUT,
                        (OH + TILE_OUT - 1) / TILE_OUT,
                        B * OC);
        depthwise_conv2d_3x3_s1_p0_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
            B, OC, H, W, OH, OW, has_bias);
    } else {
        const int block = 256;
        const long grid = (total + block - 1) / block;
        depthwise_conv2d_general_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
            B, IC, OC, H, W, OH, OW,
            KH, KW,
            (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w,
            has_bias);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_depthwise_conv_cpp_src = (
    "torch::Tensor depthwise_conv2d_forward(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride_h, int64_t stride_w, int64_t pad_h, int64_t pad_w);"
)

_depthwise_conv = load_inline(
    name="depthwise_conv2d_custom",
    cpp_sources=_depthwise_conv_cpp_src,
    cuda_sources=_depthwise_conv_cuda_src,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)

_NO_BIAS = torch.empty(0)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution accelerated with a custom CUDA kernel.

    Uses a shared-memory tiled kernel (each thread computes a 2x2 output
    patch from a 4x4 register tile) for the common 3x3 / stride 1 / pad 0
    case, with a generic depthwise kernel as fallback for other configs.
    The nn.Conv2d module is kept as a member so initialization and the
    state_dict are identical to the reference Model.
    """
    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:
        conv = self.conv2d
        if (x.is_cuda and x.dtype == torch.float32
                and conv.groups == conv.in_channels
                and conv.dilation == (1, 1)):
            bias = conv.bias if conv.bias is not None else _NO_BIAS
            return _depthwise_conv.depthwise_conv2d_forward(
                x, conv.weight, bias,
                conv.stride[0], conv.stride[1],
                conv.padding[0], conv.padding[1])
        return F.conv2d(x, conv.weight, conv.bias, conv.stride,
                        conv.padding, conv.dilation, conv.groups)
```