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


def _pair(v):
    if isinstance(v, (tuple, list)):
        return int(v[0]), int(v[1])
    return int(v), int(v)


_CPP_SRC = "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias);"

_TEMPLATE = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <stdint.h>

#define IN_C {IN_C}
#define KH {KH}
#define KW {KW}
#define STRIDE_H {SH}
#define STRIDE_W {SW}
#define PAD_H {PH}
#define PAD_W {PW}
#define DIL_H {DH}
#define DIL_W {DW}

#define TILE_W 128
#define TILE_H 32
#define THREADS_X 16
#define THREADS_Y 16
#define OUT_W 8
#define OUT_H 2
#define NTHREADS 256

#define SMEM_W ((TILE_W - 1) * STRIDE_W + (KW - 1) * DIL_W + 1)
#define SMEM_H ((TILE_H - 1) * STRIDE_H + (KH - 1) * DIL_H + 1)
#define SMEM_W_PAD ((SMEM_W + 3) & ~3)
#define SMEM_W_Q (SMEM_W / 4)
#define SMEM_W_TAIL (SMEM_W % 4)
#define SW_TOTAL (IN_C * KH * KW)
#define RLEN ((OUT_W - 1) * STRIDE_W + (KW - 1) * DIL_W + 1)

template <bool HAS_BIAS>
__global__ void __launch_bounds__(NTHREADS, 2)
direct_conv2d_kernel(const float* __restrict__ x,
                     const float* __restrict__ w,
                     const float* __restrict__ bias,
                     float* __restrict__ y,
                     const int H, const int W, const int OC,
                     const int H_out, const int W_out)
{
    __shared__ float s_w[SW_TOTAL];
    __shared__ float s_x[SMEM_H][SMEM_W_PAD];

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

    const int oc     = blockIdx.z % OC;
    const int n      = blockIdx.z / OC;
    const int tile_x = blockIdx.x * TILE_W;
    const int tile_y = blockIdx.y * TILE_H;

    // ---- load this output channel's weights into shared memory ----
    {
        const float* w_oc = w + (long)oc * (long)SW_TOTAL;
        const int nv4 = SW_TOTAL / 4;
        if ((((uintptr_t)w_oc) & 15) == 0) {
            const float4* w4 = reinterpret_cast<const float4*>(w_oc);
            for (int i = tid; i < nv4; i += NTHREADS) {
                reinterpret_cast<float4*>(s_w)[i] = w4[i];
            }
        } else {
            for (int i = tid; i < nv4 * 4; i += NTHREADS) s_w[i] = w_oc[i];
        }
        for (int i = nv4 * 4 + tid; i < SW_TOTAL; i += NTHREADS) s_w[i] = w_oc[i];
    }

    const long chan_stride = (long)H * (long)W;
    const float* x_base = x + (long)n * (long)IN_C * chan_stride;

    float acc[OUT_H][OUT_W];
#pragma unroll
    for (int i = 0; i < OUT_H; ++i)
#pragma unroll
        for (int j = 0; j < OUT_W; ++j) acc[i][j] = 0.f;

    const int base_c = tx * OUT_W;
    const int base_r = ty * OUT_H;
    const int in_r0  = base_r * STRIDE_H;
    const int in_c0  = base_c * STRIDE_W;

    // ---- loop over input channels ----
    for (int ic = 0; ic < IN_C; ++ic) {
        const float* x_c = x_base + (long)ic * chan_stride;

        // load input tile (with halo) into shared memory
#pragma unroll 1
        for (int idx = tid; idx < SMEM_H * SMEM_W_Q; idx += NTHREADS) {
            const int r  = idx / SMEM_W_Q;
            const int q  = idx - r * SMEM_W_Q;
            const int c0 = q * 4;
            const int gr  = tile_y * STRIDE_H - PAD_H + r;
            const int gc0 = tile_x * STRIDE_W - PAD_W + c0;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gr >= 0 && gr < H) {
                const float* p = x_c + (long)gr * W + gc0;
                if (gc0 >= 0 && gc0 + 3 < W && ((((uintptr_t)p) & 15) == 0)) {
                    v = *reinterpret_cast<const float4*>(p);
                } else {
                    float t0 = (gc0     >= 0 && gc0     < W) ? p[0] : 0.f;
                    float t1 = (gc0 + 1 >= 0 && gc0 + 1 < W) ? p[1] : 0.f;
                    float t2 = (gc0 + 2 >= 0 && gc0 + 2 < W) ? p[2] : 0.f;
                    float t3 = (gc0 + 3 >= 0 && gc0 + 3 < W) ? p[3] : 0.f;
                    v = make_float4(t0, t1, t2, t3);
                }
            }
            *reinterpret_cast<float4*>(&s_x[r][c0]) = v;
        }
#if SMEM_W_TAIL != 0
        for (int idx = tid; idx < SMEM_H * SMEM_W_TAIL; idx += NTHREADS) {
            const int r = idx / SMEM_W_TAIL;
            const int t = idx - r * SMEM_W_TAIL;
            const int c = SMEM_W_Q * 4 + t;
            const int gr = tile_y * STRIDE_H - PAD_H + r;
            const int gc = tile_x * STRIDE_W - PAD_W + c;
            float v = 0.f;
            if (gr >= 0 && gr < H && gc >= 0 && gc < W)
                v = x_c[(long)gr * W + gc];
            s_x[r][c] = v;
        }
#endif
        __syncthreads();

        const float* sw = s_w + ic * (KH * KW);

#if (STRIDE_H == 1) && (DIL_H == 1) && (STRIDE_W == 1) && (DIL_W == 1)
        // Sliding-window register blocking: each input row segment is loaded
        // once into registers and reused across kernel taps and output rows.
        constexpr int ROWS = OUT_H + KH - 1;
#pragma unroll
        for (int t = 0; t < ROWS; ++t) {
            float v[RLEN];
            const float* rowp = &s_x[in_r0 + t][in_c0];
#if ((OUT_W % 4) == 0) && ((RLEN % 4) == 0)
#pragma unroll
            for (int q = 0; q < RLEN / 4; ++q) {
                reinterpret_cast<float4*>(v)[q] =
                    reinterpret_cast<const float4*>(rowp)[q];
            }
#else
#pragma unroll
            for (int j = 0; j < RLEN; ++j) v[j] = rowp[j];
#endif
#pragma unroll
            for (int i = 0; i < OUT_H; ++i) {
                const int kh = t - i;
                if (kh >= 0 && kh < KH) {
                    const float* wrow = sw + kh * KW;
#pragma unroll
                    for (int kw = 0; kw < KW; ++kw) {
                        const float wv = wrow[kw];
#pragma unroll
                        for (int j = 0; j < OUT_W; ++j) {
                            acc[i][j] = fmaf(wv, v[j + kw], acc[i][j]);
                        }
                    }
                }
            }
        }
#else
#pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
#pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                const float wv = sw[kh * KW + kw];
                const int rr = in_r0 + kh * DIL_H;
                const int cc = in_c0 + kw * DIL_W;
#pragma unroll
                for (int i = 0; i < OUT_H; ++i) {
#pragma unroll
                    for (int j = 0; j < OUT_W; ++j) {
                        acc[i][j] = fmaf(wv, s_x[rr + i * STRIDE_H][cc + j * STRIDE_W], acc[i][j]);
                    }
                }
            }
        }
#endif
        __syncthreads();
    }

    // ---- epilogue ----
    const float bv = HAS_BIAS ? bias[oc] : 0.f;
    const long out_base = ((long)blockIdx.z * (long)H_out) * (long)W_out;

#pragma unroll
    for (int i = 0; i < OUT_H; ++i) {
        const int orow = tile_y + base_r + i;
        if (orow < H_out) {
            const int ocol = tile_x + base_c;
            float vals[OUT_W];
#pragma unroll
            for (int j = 0; j < OUT_W; ++j) vals[j] = acc[i][j] + bv;
            float* op = y + out_base + (long)orow * W_out + ocol;
            if (ocol + OUT_W <= W_out && ((((uintptr_t)op) & 15) == 0)) {
                float4* op4 = reinterpret_cast<float4*>(op);
#pragma unroll
                for (int j4 = 0; j4 < OUT_W / 4; ++j4) {
                    op4[j4] = make_float4(vals[j4 * 4], vals[j4 * 4 + 1],
                                          vals[j4 * 4 + 2], vals[j4 * 4 + 3]);
                }
#pragma unroll
                for (int j = (OUT_W / 4) * 4; j < OUT_W; ++j) op[j] = vals[j];
            } else {
#pragma unroll
                for (int j = 0; j < OUT_W; ++j) {
                    if (ocol + j < W_out) op[j] = vals[j];
                }
            }
        }
    }
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only fp32 supported");
    TORCH_CHECK(w.size(1) == IN_C && w.size(2) == KH && w.size(3) == KW,
                "weight shape mismatch");
    TORCH_CHECK(x.size(1) == IN_C, "input channel mismatch");

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const long N  = xc.size(0);
    const long H  = xc.size(2);
    const long W  = xc.size(3);
    const long OC = wc.size(0);
    const long H_out = (H + 2 * PAD_H - DIL_H * (KH - 1) - 1) / STRIDE_H + 1;
    const long W_out = (W + 2 * PAD_W - DIL_W * (KW - 1) - 1) / STRIDE_W + 1;

    auto y = torch::empty({N, OC, H_out, W_out}, xc.options());
    if (y.numel() == 0) return y;

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

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

    auto stream = at::cuda::getCurrentCUDAStream();
    if (has_bias) {
        direct_conv2d_kernel<true><<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            y.data_ptr<float>(), (int)H, (int)W, (int)OC, (int)H_out, (int)W_out);
    } else {
        direct_conv2d_kernel<false><<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), nullptr,
            y.data_ptr<float>(), (int)H, (int)W, (int)OC, (int)H_out, (int)W_out);
    }
    return y;
}
"""


class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, groups=groups, bias=bias)

        self._groups = groups
        self._stride = _pair(stride)
        self._padding = _pair(padding)
        self._dilation = _pair(dilation)
        kh, kw = _pair(kernel_size)

        self._ext = None
        self._native = (groups != 1)
        if not self._native:
            macros = {
                "IN_C": int(in_channels), "KH": kh, "KW": kw,
                "SH": self._stride[0], "SW": self._stride[1],
                "PH": self._padding[0], "PW": self._padding[1],
                "DH": self._dilation[0], "DW": self._dilation[1],
            }
            src = _TEMPLATE
            for k, v in macros.items():
                src = src.replace("{%s}" % k, str(v))
            name = "direct_conv2d_c{}_k{}x{}_s{}x{}_p{}x{}_d{}x{}_v2".format(*macros.values())
            self._ext = load_inline(
                name=name,
                cpp_sources=_CPP_SRC,
                cuda_sources=src,
                functions=["conv2d_forward"],
                verbose=False,
                extra_cuda_cflags=["-O3", "--use_fast_math"],
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (not self._native) and x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            w = self.conv2d.weight
            b = self.conv2d.bias
            if b is None:
                b = x.new_empty(0)
            return self._ext.conv2d_forward(x, w, b)
        return self.conv2d(x)
```