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

# Custom CUDA kernel for depthwise 2D convolution.
# Each thread block computes a 16x16 tile of outputs for one (batch, channel)
# pair. The corresponding (halo'd) input tile is staged once in shared memory
# and reused for all taps, while the small per-channel filter is cached in
# registers. Compile-time kernel-size templating fully unrolls the tap loops.
depthwise_conv2d_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_W 16
#define BLOCK_H 16

template <int K>
__global__ void depthwise_conv2d_tiled_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int in_channels,
    const int out_channels,
    const int channel_multiplier,
    const int height_in,
    const int width_in,
    const int stride,
    const int padding,
    const int height_out,
    const int width_out)
{
    extern __shared__ float tile[];

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

    const int w_out = blockIdx.x * BLOCK_W + tx;
    const int h_out = blockIdx.y * BLOCK_H + ty;

    const int bc = blockIdx.z;
    const int b = bc / out_channels;
    const int c = bc % out_channels;
    const int ic = c / channel_multiplier;

    const int tile_in_w = (BLOCK_W - 1) * stride + K;
    const int tile_in_h = (BLOCK_H - 1) * stride + K;
    const int tile_ld = tile_in_w + 1;  // padded to reduce bank conflicts

    const int h_in_start = blockIdx.y * BLOCK_H * stride - padding;
    const int w_in_start = blockIdx.x * BLOCK_W * stride - padding;

    const float* input_plane = input + ((int64_t)b * in_channels + ic) * (int64_t)height_in * width_in;

    // Cooperatively stage the input tile (zero-padded outside the image).
    for (int y = ty; y < tile_in_h; y += BLOCK_H) {
        const int h_in = h_in_start + y;
        const bool row_ok = (h_in >= 0) && (h_in < height_in);
        const int64_t row_off = (int64_t)h_in * width_in;
        for (int x = tx; x < tile_in_w; x += BLOCK_W) {
            const int w_in = w_in_start + x;
            float val = 0.0f;
            if (row_ok && w_in >= 0 && w_in < width_in) {
                val = input_plane[row_off + w_in];
            }
            tile[y * tile_ld + x] = val;
        }
    }
    __syncthreads();

    if (w_out < width_out && h_out < height_out) {
        // Cache the per-channel filter in registers.
        float w[K * K];
        const float* wp = weight + (int64_t)c * K * K;
        #pragma unroll
        for (int i = 0; i < K * K; ++i) {
            w[i] = wp[i];
        }

        float sum = (bias != nullptr) ? bias[c] : 0.0f;

        const int base_y = ty * stride;
        const int base_x = tx * stride;

        #pragma unroll
        for (int kh = 0; kh < K; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < K; ++kw) {
                sum = fmaf(tile[(base_y + kh) * tile_ld + (base_x + kw)], w[kh * K + kw], sum);
            }
        }

        output[((int64_t)bc * height_out + h_out) * width_out + w_out] = sum;
    }
}

// Generic fallback for uncommon kernel sizes.
__global__ void depthwise_conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int in_channels,
    const int out_channels,
    const int channel_multiplier,
    const int height_in,
    const int width_in,
    const int kernel_size,
    const int stride,
    const int padding,
    const int height_out,
    const int width_out)
{
    extern __shared__ float tile[];

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

    const int w_out = blockIdx.x * BLOCK_W + tx;
    const int h_out = blockIdx.y * BLOCK_H + ty;

    const int bc = blockIdx.z;
    const int b = bc / out_channels;
    const int c = bc % out_channels;
    const int ic = c / channel_multiplier;

    const int tile_in_w = (BLOCK_W - 1) * stride + kernel_size;
    const int tile_in_h = (BLOCK_H - 1) * stride + kernel_size;
    const int tile_ld = tile_in_w + 1;

    const int h_in_start = blockIdx.y * BLOCK_H * stride - padding;
    const int w_in_start = blockIdx.x * BLOCK_W * stride - padding;

    const float* input_plane = input + ((int64_t)b * in_channels + ic) * (int64_t)height_in * width_in;

    for (int y = ty; y < tile_in_h; y += BLOCK_H) {
        const int h_in = h_in_start + y;
        const bool row_ok = (h_in >= 0) && (h_in < height_in);
        const int64_t row_off = (int64_t)h_in * width_in;
        for (int x = tx; x < tile_in_w; x += BLOCK_W) {
            const int w_in = w_in_start + x;
            float val = 0.0f;
            if (row_ok && w_in >= 0 && w_in < width_in) {
                val = input_plane[row_off + w_in];
            }
            tile[y * tile_ld + x] = val;
        }
    }
    __syncthreads();

    if (w_out < width_out && h_out < height_out) {
        const float* wp = weight + (int64_t)c * kernel_size * kernel_size;
        float sum = (bias != nullptr) ? bias[c] : 0.0f;

        const int base_y = ty * stride;
        const int base_x = tx * stride;

        for (int kh = 0; kh < kernel_size; ++kh) {
            for (int kw = 0; kw < kernel_size; ++kw) {
                sum = fmaf(tile[(base_y + kh) * tile_ld + (base_x + kw)], wp[kh * kernel_size + kw], sum);
            }
        }

        output[((int64_t)bc * height_out + h_out) * width_out + w_out] = sum;
    }
}

torch::Tensor depthwise_conv2d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t stride,
    int64_t padding)
{
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "input and weight must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(input.dim() == 4, "input must be NCHW");
    TORCH_CHECK(weight.dim() == 4, "weight must be (out_channels, 1, K, K)");
    TORCH_CHECK(weight.size(2) == weight.size(3), "kernel must be square");

    auto input_c = input.contiguous();
    auto weight_c = weight.contiguous();

    const int batch_size = (int)input_c.size(0);
    const int in_channels = (int)input_c.size(1);
    const int height_in = (int)input_c.size(2);
    const int width_in = (int)input_c.size(3);
    const int out_channels = (int)weight_c.size(0);
    const int kernel_size = (int)weight_c.size(2);
    TORCH_CHECK(out_channels % in_channels == 0,
                "out_channels must be a multiple of in_channels for depthwise conv");
    const int channel_multiplier = out_channels / in_channels;

    const int height_out = (int)((height_in + 2 * padding - kernel_size) / stride + 1);
    const int width_out = (int)((width_in + 2 * padding - kernel_size) / stride + 1);

    auto output = torch::empty({batch_size, out_channels, height_out, width_out}, input_c.options());

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

    const int tile_in_w = (BLOCK_W - 1) * (int)stride + kernel_size;
    const int tile_in_h = (BLOCK_H - 1) * (int)stride + kernel_size;
    const int tile_ld = tile_in_w + 1;
    const size_t shared_mem = (size_t)tile_in_h * tile_ld * sizeof(float);

    dim3 block(BLOCK_W, BLOCK_H, 1);
    dim3 grid((unsigned)((width_out + BLOCK_W - 1) / BLOCK_W),
              (unsigned)((height_out + BLOCK_H - 1) / BLOCK_H),
              (unsigned)(batch_size * out_channels));

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

    const float* in_ptr = input_c.data_ptr<float>();
    const float* w_ptr = weight_c.data_ptr<float>();
    float* out_ptr = output.data_ptr<float>();

#define LAUNCH_TILED(KK)                                                          \
    depthwise_conv2d_tiled_kernel<KK><<<grid, block, shared_mem, stream>>>(       \
        in_ptr, w_ptr, bias_ptr, out_ptr,                                         \
        in_channels, out_channels, channel_multiplier,                            \
        height_in, width_in, (int)stride, (int)padding, height_out, width_out)

    switch (kernel_size) {
        case 1: LAUNCH_TILED(1); break;
        case 2: LAUNCH_TILED(2); break;
        case 3: LAUNCH_TILED(3); break;
        case 4: LAUNCH_TILED(4); break;
        case 5: LAUNCH_TILED(5); break;
        case 7: LAUNCH_TILED(7); break;
        default:
            depthwise_conv2d_generic_kernel<<<grid, block, shared_mem, stream>>>(
                in_ptr, w_ptr, bias_ptr, out_ptr,
                in_channels, out_channels, channel_multiplier,
                height_in, width_in, kernel_size, (int)stride, (int)padding,
                height_out, width_out);
            break;
    }
#undef LAUNCH_TILED

    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

depthwise_conv2d_cpp_src = (
    "torch::Tensor depthwise_conv2d_cuda(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride, int64_t padding);"
)

_depthwise_conv2d_mod = load_inline(
    name="depthwise_conv2d_custom",
    cpp_sources=depthwise_conv2d_cpp_src,
    cuda_sources=depthwise_conv2d_cuda_src,
    functions=["depthwise_conv2d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution backed by a custom tiled CUDA kernel.
    Parameters are held in an nn.Conv2d module so initialization and state
    dicts match the reference model exactly.
    """
    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.stride = stride
        self.padding = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (not x.is_cuda) or x.dtype != torch.float32:
            return self.conv2d(x)
        bias = self.conv2d.bias
        if bias is None:
            bias = x.new_empty(0)
        return _depthwise_conv2d_mod.depthwise_conv2d_cuda(
            x, self.conv2d.weight, bias, self.stride, self.padding
        )
```