```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 kernel: fused, tiled depthwise 2D convolution specialized for a
# 3x7 kernel, stride 1, zero padding and unit dilation (the configuration used
# by this model). Each thread block computes a 64x16 output tile for one
# (batch, channel) slice using a shared-memory input tile, register-cached
# weights, a register sliding window over the 7-tap width, vectorized float4
# input loads and float2 output stores.
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

namespace {

constexpr int KH = 3;
constexpr int KW = 7;
constexpr int TILE_W = 64;             // output columns per block
constexpr int TILE_H = 16;             // output rows per block
constexpr int THREADS_X = TILE_W / 2;  // each thread produces 2 adjacent outputs
constexpr int NTHREADS = THREADS_X * TILE_H;
constexpr int SH_W = TILE_W + KW - 1;  // 70
constexpr int SH_H = TILE_H + KH - 1;  // 18
constexpr int QUAD_W = SH_W / 4;       // 17
constexpr int REM_W = SH_W % 4;        // 2

__global__ __launch_bounds__(NTHREADS)
void depthwise_conv2d_3x7_s1_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    const int C,
    const int H,
    const int W,
    const int H_out,
    const int W_out)
{
    __shared__ float tile[SH_H][SH_W];

    const int nc = blockIdx.z;
    const int c = nc % C;
    const float* __restrict__ in_ptr = input + (long long)nc * H * W;

    const int x0 = blockIdx.x * TILE_W;
    const int y0 = blockIdx.y * TILE_H;

    // Cache this channel's 21 weights in registers.
    float w[KH * KW];
    const float* __restrict__ wc = weight + c * (KH * KW);
#pragma unroll
    for (int k = 0; k < KH * KW; ++k) {
        w[k] = __ldg(wc + k);
    }

    const int tid = threadIdx.y * THREADS_X + threadIdx.x;

    // Cooperative load of the (18 x 70) input tile into shared memory.
    if ((W & 3) == 0) {
        // Every row stays 16B aligned when W % 4 == 0, so use float4 loads.
        if (tid < SH_H * QUAD_W) {
            const int sy = tid / QUAD_W;
            const int sx = (tid - sy * QUAD_W) * 4;
            const int gx = x0 + sx;
            const int gy = y0 + sy;
            float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            if (gy < H) {
                const float* row = in_ptr + (long long)gy * W;
                if (gx + 3 < W) {
                    v = *reinterpret_cast<const float4*>(row + gx);
                } else {
                    if (gx + 0 < W) v.x = __ldg(row + gx + 0);
                    if (gx + 1 < W) v.y = __ldg(row + gx + 1);
                    if (gx + 2 < W) v.z = __ldg(row + gx + 2);
                    if (gx + 3 < W) v.w = __ldg(row + gx + 3);
                }
            }
            float* trow = tile[sy] + sx;
            trow[0] = v.x;
            trow[1] = v.y;
            trow[2] = v.z;
            trow[3] = v.w;
        } else if (tid < SH_H * QUAD_W + SH_H * REM_W) {
            const int t2 = tid - SH_H * QUAD_W;
            const int sy = t2 / REM_W;
            const int sx = QUAD_W * 4 + (t2 - sy * REM_W);
            const int gx = x0 + sx;
            const int gy = y0 + sy;
            float v = 0.0f;
            if (gy < H && gx < W) {
                v = __ldg(in_ptr + (long long)gy * W + gx);
            }
            tile[sy][sx] = v;
        }
    } else {
        for (int i = tid; i < SH_H * SH_W; i += NTHREADS) {
            const int sy = i / SH_W;
            const int sx = i - sy * SH_W;
            const int gy = y0 + sy;
            const int gx = x0 + sx;
            float v = 0.0f;
            if (gy < H && gx < W) {
                v = __ldg(in_ptr + (long long)gy * W + gx);
            }
            tile[sy][sx] = v;
        }
    }
    __syncthreads();

    const int ox = x0 + threadIdx.x * 2;
    const int oy = y0 + threadIdx.y;
    if (ox >= W_out || oy >= H_out) {
        return;
    }

    const int sx = threadIdx.x * 2;
    const int sy = threadIdx.y;

    // Register sliding window over the 7-tap width; two adjacent outputs.
    float acc0 = 0.0f;
    float acc1 = 0.0f;
#pragma unroll
    for (int kh = 0; kh < KH; ++kh) {
        float v[KW + 1];
#pragma unroll
        for (int i = 0; i < KW + 1; ++i) {
            v[i] = tile[sy + kh][sx + i];
        }
#pragma unroll
        for (int kw = 0; kw < KW; ++kw) {
            acc0 = fmaf(v[kw],     w[kh * KW + kw], acc0);
            acc1 = fmaf(v[kw + 1], w[kh * KW + kw], acc1);
        }
    }

    float* __restrict__ out_ptr = output + ((long long)nc * H_out + oy) * W_out;
    if ((W_out & 1) == 0) {
        // With even W_out, (nc*H_out + oy)*W_out + ox is always even and
        // ox + 1 < W_out always holds, so one aligned 8B store suffices.
        *reinterpret_cast<float2*>(out_ptr + ox) = make_float2(acc0, acc1);
    } else {
        out_ptr[ox] = acc0;
        if (ox + 1 < W_out) {
            out_ptr[ox + 1] = acc1;
        }
    }
}

}  // namespace

torch::Tensor depthwise_conv2d_forward(torch::Tensor input, torch::Tensor weight) {
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "input and weight must be CUDA tensors");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "input/weight must be 4-D");
    TORCH_CHECK(input.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");

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

    const int N = static_cast<int>(x.size(0));
    const int C = static_cast<int>(x.size(1));
    const int H = static_cast<int>(x.size(2));
    const int W = static_cast<int>(x.size(3));
    TORCH_CHECK(wgt.size(0) == C && wgt.size(1) == 1 && wgt.size(2) == KH && wgt.size(3) == KW,
                "expected depthwise weights of shape (C, 1, 3, 7)");

    const int H_out = H - KH + 1;
    const int W_out = W - KW + 1;
    TORCH_CHECK(H_out > 0 && W_out > 0, "input is smaller than the 3x7 kernel");

    auto out = at::empty({N, C, H_out, W_out}, x.options());

    dim3 block(THREADS_X, TILE_H, 1);
    dim3 grid((W_out + TILE_W - 1) / TILE_W, (H_out + TILE_H - 1) / TILE_H, N * C);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    depthwise_conv2d_3x7_s1_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        wgt.data_ptr<float>(),
        out.data_ptr<float>(),
        C, H, W, H_out, W_out);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_cpp_source = "torch::Tensor depthwise_conv2d_forward(torch::Tensor input, torch::Tensor weight);"

_depthwise_ext = load_inline(
    name="depthwise_conv2d_3x7_ext",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["depthwise_conv2d_forward"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _DepthwiseConv2d3x7Fn(torch.autograd.Function):
    """Autograd wrapper: custom CUDA forward, PyTorch-computed backward."""

    @staticmethod
    def forward(ctx, x, weight):
        ctx.save_for_backward(x, weight)
        return _depthwise_ext.depthwise_conv2d_forward(x, weight)

    @staticmethod
    def backward(ctx, grad_output):
        x, weight = ctx.saved_tensors
        with torch.enable_grad():
            x_r = x.detach().requires_grad_(True)
            w_r = weight.detach().requires_grad_(True)
            out = F.conv2d(x_r, w_r, None, (1, 1), (0, 0), (1, 1), x.shape[1])
            grad_input, grad_weight = torch.autograd.grad(
                out, (x_r, w_r), grad_output.contiguous())
        return grad_input, grad_weight


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution with an asymmetric 3x7 kernel, accelerated by a
    custom shared-memory tiled CUDA kernel. Falls back to the standard
    nn.Conv2d implementation for any configuration the kernel does not cover.
    """
    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._use_custom = (
            kernel_size_h == 3 and kernel_size_w == 7 and
            stride_h == 1 and stride_w == 1 and
            padding_h == 0 and padding_w == 0 and
            dilation_h == 1 and dilation_w == 1
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (self._use_custom
                and x.is_cuda
                and x.dtype == torch.float32
                and x.dim() == 4
                and x.shape[1] == self.conv2d.weight.shape[0]
                and x.shape[2] >= 3 and x.shape[3] >= 7
                and x.shape[0] * x.shape[1] <= 65535):
            return _DepthwiseConv2d3x7Fn.apply(x, self.conv2d.weight)
        return self.conv2d(x)
```