```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernels for depthwise 2D convolution (v5).
#
# This workload is DRAM-bandwidth bound (read the whole input once, write the
# whole output once), so the design goals are:
#   * high occupancy to hide DRAM latency: the primary fast path uses
#     512-thread blocks (32x16 threads, 4x1 outputs/thread) with filter taps
#     staged in shared memory (broadcast, conflict-free) to keep register
#     usage under 64 regs/thread -> 2 resident blocks/SM (1024 threads/SM,
#     2x the occupancy of the previous 256-thread design);
#   * the (TILE_H+KH-1) x padded(TILE_W+KW-1) input halo tile is staged in
#     shared memory with aligned float4 global/shared accesses whenever the
#     tensor layout allows (W % 4 == 0 and 16B-aligned rows), so each input
#     element is fetched from DRAM ~once (halo re-reads are served by L2);
#   * each thread pulls its shared-memory region into registers exactly once
#     with aligned float4 loads (conflict free);
#   * output stores are streaming stores (__stcs, evict-first) so the 500+MB
#     output stream does not evict input halo lines from L2, and are
#     vectorized (float4 when outW % 4 == 0, float2 when outW % 2 == 0,
#     scalar otherwise).
# A secondary tiled config (256 threads, 4x2 outputs/thread) handles taller
# kernels, and a fully general kernel handles any kernel size, stride,
# padding and dilation.
dwconv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define DW_TX 32   // threads along width
#define DW_RW 4    // output columns per thread
#define DW_TILE_W (DW_TX * DW_RW)   // 128 output columns per block

template <int KH, int KW, int TY, int RH, int MINB>
__global__ void __launch_bounds__(DW_TX * TY, MINB)
dwconv2d_tile_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int C, int H, int W, int outH, int outW,
    int tilesX, int tilesPerSlice, int vecOK)
{
    constexpr int NT     = DW_TX * TY;
    constexpr int TILE_H = TY * RH;
    // Register-block width rounded up to a float4 multiple; SH_W is padded so
    // that every per-thread register-block row is 16B aligned in shared memory.
    constexpr int REG_W  = (((DW_RW + KW - 1) + 3) >> 2) << 2;
    constexpr int SH_W   = (DW_TILE_W - DW_RW) + REG_W;
    constexpr int SH_H   = TILE_H + KH - 1;

    __shared__ __align__(16) float sh[SH_H * SH_W];
    __shared__ float wsh[KH * KW];

    const int tile = blockIdx.x % tilesPerSlice;
    const int nc   = blockIdx.x / tilesPerSlice;
    const int ty   = tile / tilesX;
    const int tx   = tile % tilesX;
    const int n    = nc / C;
    const int c    = nc % C;

    const int ow0 = tx * DW_TILE_W;
    const int oh0 = ty * TILE_H;
    const int tid = threadIdx.y * DW_TX + threadIdx.x;

    // Stage this channel's filter taps in shared memory (broadcast reads).
    if (tid < KH * KW) wsh[tid] = weight[(long)c * (KH * KW) + tid];

    // Stage the input halo tile in shared memory (zero outside the image).
    const float* in_base = input + ((long)n * C + c) * H * W;
    if (vecOK) {
        constexpr int V4ROW = SH_W / 4;
        constexpr int V4TOT = SH_H * V4ROW;
        #pragma unroll 2
        for (int idx = tid; idx < V4TOT; idx += NT) {
            const int sy = idx / V4ROW;
            const int sx = (idx - sy * V4ROW) * 4;
            const int gy = oh0 + sy;
            const int gx = ow0 + sx;
            float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            if (gy < H) {
                const float* gptr = in_base + (long)gy * W + gx;
                if (gx + 3 < W) {
                    v = *reinterpret_cast<const float4*>(gptr);
                } else {
                    v.x = (gx + 0 < W) ? gptr[0] : 0.0f;
                    v.y = (gx + 1 < W) ? gptr[1] : 0.0f;
                    v.z = (gx + 2 < W) ? gptr[2] : 0.0f;
                    v.w = (gx + 3 < W) ? gptr[3] : 0.0f;
                }
            }
            *reinterpret_cast<float4*>(sh + sy * SH_W + sx) = v;
        }
    } else {
        #pragma unroll 2
        for (int idx = tid; idx < SH_H * SH_W; idx += NT) {
            const int sy = idx / SH_W;
            const int sx = idx - sy * SH_W;
            const int gy = oh0 + sy;
            const int gx = ow0 + sx;
            sh[idx] = (gy < H && gx < W) ? in_base[(long)gy * W + gx] : 0.0f;
        }
    }
    __syncthreads();

    const int sy0 = threadIdx.y * RH;
    const int sx0 = threadIdx.x * DW_RW;

    // Pull this thread's shared-memory region into registers exactly once
    // using aligned float4 loads (conflict free).
    float reg[RH + KH - 1][REG_W];
    #pragma unroll
    for (int i = 0; i < RH + KH - 1; ++i) {
        const float* srow = sh + (sy0 + i) * SH_W + sx0;
        #pragma unroll
        for (int j = 0; j < REG_W / 4; ++j) {
            const float4 v = *reinterpret_cast<const float4*>(srow + 4 * j);
            reg[i][4 * j + 0] = v.x;
            reg[i][4 * j + 1] = v.y;
            reg[i][4 * j + 2] = v.z;
            reg[i][4 * j + 3] = v.w;
        }
    }

    const float b0 = (bias != nullptr) ? bias[c] : 0.0f;
    float acc[RH][DW_RW];
    #pragma unroll
    for (int i = 0; i < RH; ++i)
        #pragma unroll
        for (int j = 0; j < DW_RW; ++j)
            acc[i][j] = b0;

    #pragma unroll
    for (int kh = 0; kh < KH; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < KW; ++kw) {
            const float wv = wsh[kh * KW + kw];
            #pragma unroll
            for (int i = 0; i < RH; ++i) {
                #pragma unroll
                for (int j = 0; j < DW_RW; ++j) {
                    acc[i][j] = fmaf(reg[i + kh][j + kw], wv, acc[i][j]);
                }
            }
        }
    }

    const int ow = ow0 + sx0;
    const int oh = oh0 + sy0;
    float* out_base = output + ((long)n * C + c) * outH * outW;
    // ow is always a multiple of DW_RW (4) and out is a fresh 16B-aligned
    // allocation, so stores can be safely vectorized based on outW divisibility.
    // Streaming (evict-first) stores keep L2 available for input halo reuse.
    if ((outW & 3) == 0) {
        #pragma unroll
        for (int i = 0; i < RH; ++i) {
            const int oy = oh + i;
            if (oy < outH && ow < outW) {
                __stcs(reinterpret_cast<float4*>(out_base + (long)oy * outW + ow),
                       make_float4(acc[i][0], acc[i][1], acc[i][2], acc[i][3]));
            }
        }
    } else if ((outW & 1) == 0) {
        #pragma unroll
        for (int i = 0; i < RH; ++i) {
            const int oy = oh + i;
            if (oy < outH) {
                float* orow = out_base + (long)oy * outW + ow;
                if (ow < outW) {
                    __stcs(reinterpret_cast<float2*>(orow), make_float2(acc[i][0], acc[i][1]));
                }
                if (ow + 2 < outW) {
                    __stcs(reinterpret_cast<float2*>(orow + 2), make_float2(acc[i][2], acc[i][3]));
                }
            }
        }
    } else {
        #pragma unroll
        for (int i = 0; i < RH; ++i) {
            const int oy = oh + i;
            if (oy < outH) {
                float* orow = out_base + (long)oy * outW + ow;
                #pragma unroll
                for (int j = 0; j < DW_RW; ++j) {
                    if (ow + j < outW) __stcs(orow + j, acc[i][j]);
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// General depthwise kernel: any kernel size / stride / padding / dilation.
// ---------------------------------------------------------------------------
__global__ void dwconv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int C, int H, int W, int outH, int outW,
    int KH, int KW, int strideH, int strideW, int padH, int padW, int dilH, int dilW)
{
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long total = (long)N * C * outH * outW;
    if (idx >= total) return;

    const int ow = (int)(idx % outW);
    const int oh = (int)((idx / outW) % outH);
    const int c  = (int)((idx / ((long)outW * outH)) % C);
    const int n  = (int)(idx / ((long)outW * outH * C));

    const float* in_base = input + ((long)n * C + c) * H * W;
    const float* w_base  = weight + (long)c * KH * KW;

    float sum = (bias != nullptr) ? bias[c] : 0.0f;
    for (int kh = 0; kh < KH; ++kh) {
        const int ih = oh * strideH - padH + kh * dilH;
        if (ih < 0 || ih >= H) continue;
        const float* in_row = in_base + (long)ih * W;
        const float* w_row  = w_base + kh * KW;
        for (int kw = 0; kw < KW; ++kw) {
            const int iw = ow * strideW - padW + kw * dilW;
            if (iw < 0 || iw >= W) continue;
            sum = fmaf(in_row[iw], w_row[kw], sum);
        }
    }
    output[idx] = sum;
}

template <int KH, int KW, int TY, int RH, int MINB>
static void launch_tile(
    const float* input, const float* weight, const float* bias, float* output,
    int N, int C, int H, int W, int outH, int outW, int vecOK, cudaStream_t stream)
{
    constexpr int TILE_H = TY * RH;
    const int tilesX = (outW + DW_TILE_W - 1) / DW_TILE_W;
    const int tilesY = (outH + TILE_H - 1) / TILE_H;
    const int tilesPerSlice = tilesX * tilesY;
    const long blocks = (long)tilesPerSlice * N * C;
    dim3 grid((unsigned int)blocks);
    dim3 block(DW_TX, TY);
    dwconv2d_tile_kernel<KH, KW, TY, RH, MINB><<<grid, block, 0, stream>>>(
        input, weight, bias, output, N, C, H, W, outH, outW, tilesX, tilesPerSlice, vecOK);
}

torch::Tensor dwconv2d_forward(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t strideH, int64_t strideW,
    int64_t padH, int64_t padW,
    int64_t dilH, int64_t dilW)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");

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

    const int N  = (int)x.size(0);
    const int C  = (int)x.size(1);
    const int H  = (int)x.size(2);
    const int W  = (int)x.size(3);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);

    const int outH = (int)((H + 2 * padH - dilH * (KH - 1) - 1) / strideH + 1);
    const int outW = (int)((W + 2 * padW - dilW * (KW - 1) - 1) / strideW + 1);

    auto out = torch::empty({N, C, outH, outW}, x.options());
    if (out.numel() == 0) return out;

    const float* in_ptr = x.data_ptr<float>();
    const float* w_ptr  = w.data_ptr<float>();
    float* out_ptr      = out.data_ptr<float>();

    torch::Tensor b;
    const float* b_ptr = nullptr;
    if (bias.has_value() && bias->defined()) {
        b = bias->contiguous();
        b_ptr = b.data_ptr<float>();
    }

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

    const bool fast = (strideH == 1 && strideW == 1 && padH == 0 && padW == 0 &&
                       dilH == 1 && dilW == 1 && H >= KH && W >= KW);

    if (fast) {
        // 16B-aligned base, 16B-aligned rows and per-slice offsets allow
        // fully vectorized float4 staging.
        const int vecOK = (((reinterpret_cast<uintptr_t>(in_ptr)) & 15) == 0 &&
                           (W & 3) == 0 &&
                           (((long)H * W) & 3) == 0) ? 1 : 0;
        // High-occupancy config: 512 threads (32x16), 4x1 outputs/thread.
        if (KH == 3 && KW == 7) { launch_tile<3, 7, 16, 1, 2>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 3 && KW == 3) { launch_tile<3, 3, 16, 1, 2>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 1 && KW == 7) { launch_tile<1, 7, 16, 1, 2>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 1 && KW == 3) { launch_tile<1, 3, 16, 1, 2>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 3 && KW == 1) { launch_tile<3, 1, 16, 1, 2>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 1 && KW == 1) { launch_tile<1, 1, 16, 1, 2>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        // Taller kernels: 256 threads (32x8), 4x2 outputs/thread.
        if (KH == 7 && KW == 3) { launch_tile<7, 3, 8, 2, 1>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 5 && KW == 5) { launch_tile<5, 5, 8, 2, 1>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
        if (KH == 7 && KW == 7) { launch_tile<7, 7, 8, 2, 1>(in_ptr, w_ptr, b_ptr, out_ptr, N, C, H, W, outH, outW, vecOK, stream); return out; }
    }

    const long total = (long)N * C * outH * outW;
    const int threads = 256;
    const long blocks = (total + threads - 1) / threads;
    dwconv2d_generic_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        in_ptr, w_ptr, b_ptr, out_ptr,
        N, C, H, W, outH, outW,
        KH, KW, (int)strideH, (int)strideW, (int)padH, (int)padW, (int)dilH, (int)dilW);
    return out;
}
"""

dwconv2d_cpp_source = (
    "torch::Tensor dwconv2d_forward(torch::Tensor input, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias, int64_t strideH, int64_t strideW, "
    "int64_t padH, int64_t padW, int64_t dilH, int64_t dilW);"
)

# Compile the inline CUDA extension
dwconv2d_ext = load_inline(
    name="dwconv2d_ext_v5",
    cpp_sources=dwconv2d_cpp_source,
    cuda_sources=dwconv2d_cuda_source,
    functions=["dwconv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)

_dwconv2d_forward = dwconv2d_ext.dwconv2d_forward


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

    The nn.Conv2d submodule is kept so that parameter initialization and
    state_dict layout exactly match the reference Model.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size_h: int, kernel_size_w: int, stride_h: int = 1, stride_w: int = 1, padding_h: int = 0, padding_w: int = 0, dilation_h: int = 1, dilation_w: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, (kernel_size_h, kernel_size_w), stride=(stride_h, stride_w), padding=(padding_h, padding_w), dilation=(dilation_h, dilation_w), groups=in_channels, bias=bias)
        self.stride_h = stride_h
        self.stride_w = stride_w
        self.padding_h = padding_h
        self.padding_w = padding_w
        self.dilation_h = dilation_h
        self.dilation_w = dilation_w

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.conv2d.weight
        if x.is_cuda and weight.is_cuda and x.dtype == torch.float32 and weight.dtype == torch.float32:
            return _dwconv2d_forward(
                x, weight, self.conv2d.bias,
                self.stride_h, self.stride_w,
                self.padding_h, self.padding_w,
                self.dilation_h, self.dilation_w)
        return F.conv2d(
            x, weight, self.conv2d.bias,
            stride=(self.stride_h, self.stride_w),
            padding=(self.padding_h, self.padding_w),
            dilation=(self.dilation_h, self.dilation_w),
            groups=self.conv2d.groups)
```