# Optimized Model with Custom CUDA Convolution Kernel

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

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

// ---------------------------------------------------------------------------
// Fast direct 2D convolution kernel (fp32), specialized for
//   KH x KW kernel, dilation (DH, DW), stride 1.
//
// Strategy (register-blocked, smem-tiled direct convolution):
//   * Each thread block computes a tile of TILE_H x TILE_W output pixels
//     for OC_TILE output channels of one batch element.
//   * The required input halo tile (IN_ROWS x IN_COLS) for one input channel
//     is staged in shared memory (zero-filled outside the image -> padding).
//   * Each thread produces 4 spatial outputs (columns tx, tx+16, tx+32,
//     tx+48 -> conflict-avoiding strided layout) x OC_TILE output channels,
//     giving 32 independent FMA chains per tap (high FMA-pipe utilization).
// ---------------------------------------------------------------------------

#define TILE_H   16
#define TILE_W   64
#define OC_TILE  8
#define NTHREADS 256

template <int KH, int KW, int DH, int DW>
__global__ void __launch_bounds__(NTHREADS)
conv2d_direct_kernel(const float* __restrict__ input,
                     const float* __restrict__ weight,
                     const float* __restrict__ bias,
                     float* __restrict__ output,
                     int N, int IC, int H, int W,
                     int OC, int OH, int OW,
                     int pad_h, int pad_w) {
    constexpr int IN_ROWS = TILE_H + (KH - 1) * DH;   // 16 + 8  = 24
    constexpr int IN_COLS = TILE_W + (KW - 1) * DW;   // 64 + 24 = 88
    constexpr int NTAP    = KH * KW;

    __shared__ float  s_in[IN_ROWS * IN_COLS];
    __shared__ float4 s_w[NTAP * 2];                  // [tap][8 oc] as 2 x float4

    const int tx  = threadIdx.x & 15;                 // 0..15
    const int ty  = threadIdx.x >> 4;                 // 0..15
    const int tid = threadIdx.x;

    const int ow0 = blockIdx.x * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int n   = blockIdx.z / (OC / OC_TILE);
    const int oc0 = (blockIdx.z % (OC / OC_TILE)) * OC_TILE;

    float acc[OC_TILE][4];
    #pragma unroll
    for (int o = 0; o < OC_TILE; ++o) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) acc[o][j] = 0.f;
    }

    const int in_row0 = oh0 - pad_h;
    const int in_col0 = ow0 - pad_w;
    const size_t in_batch_stride = (size_t)IC * H * W;
    const float* in_base = input + (size_t)n * in_batch_stride;

    const bool interior = (in_row0 >= 0) && (in_row0 + IN_ROWS <= H) &&
                          (in_col0 >= 0) && (in_col0 + IN_COLS <= W);

    for (int ic = 0; ic < IC; ++ic) {
        // ---- stage weight tile: s_w[tap][oc_local] -------------------------
        {
            const float* wsrc = weight + ((size_t)oc0 * IC + ic) * NTAP;
            #pragma unroll
            for (int k = 0; k < (NTAP * OC_TILE + NTHREADS - 1) / NTHREADS; ++k) {
                int idx = tid + k * NTHREADS;
                if (idx < NTAP * OC_TILE) {
                    int ocl = idx / NTAP;             // 0..7
                    int tap = idx - ocl * NTAP;       // 0..NTAP-1 (coalesced)
                    reinterpret_cast<float*>(s_w)[tap * OC_TILE + ocl] =
                        wsrc[(size_t)ocl * IC * NTAP + tap];
                }
            }
        }

        // ---- stage input tile ----------------------------------------------
        {
            const float* src_base = in_base + (size_t)ic * H * W;
            #pragma unroll
            for (int k = 0; k < (IN_ROWS * (IN_COLS / 4) + NTHREADS - 1) / NTHREADS; ++k) {
                int idx = tid + k * NTHREADS;
                if (idx < IN_ROWS * (IN_COLS / 4)) {
                    int r  = idx / (IN_COLS / 4);
                    int cq = idx - r * (IN_COLS / 4);
                    float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
                    if (interior) {
                        const float* p = src_base + (size_t)(in_row0 + r) * W +
                                         (in_col0 + cq * 4);
                        v = *reinterpret_cast<const float4*>(p);
                    } else {
                        int gr = in_row0 + r;
                        if ((unsigned)gr < (unsigned)H) {
                            int gc = in_col0 + cq * 4;
                            const float* p = src_base + (size_t)gr * W;
                            float t0 = ((unsigned)(gc    ) < (unsigned)W) ? p[gc    ] : 0.f;
                            float t1 = ((unsigned)(gc + 1) < (unsigned)W) ? p[gc + 1] : 0.f;
                            float t2 = ((unsigned)(gc + 2) < (unsigned)W) ? p[gc + 2] : 0.f;
                            float t3 = ((unsigned)(gc + 3) < (unsigned)W) ? p[gc + 3] : 0.f;
                            v = make_float4(t0, t1, t2, t3);
                        }
                    }
                    reinterpret_cast<float4*>(s_in)[r * (IN_COLS / 4) + cq] = v;
                }
            }
        }
        __syncthreads();

        // ---- compute: 45 taps x (4 loads + 8 weights + 32 FMA) --------------
        const int sbase = ty * IN_COLS + tx;
        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                const int off = sbase + kh * DH * IN_COLS + kw * DW;
                float a0 = s_in[off];
                float a1 = s_in[off + 16];
                float a2 = s_in[off + 32];
                float a3 = s_in[off + 48];
                float4 wA = s_w[(kh * KW + kw) * 2];
                float4 wB = s_w[(kh * KW + kw) * 2 + 1];

                acc[0][0] += a0 * wA.x; acc[0][1] += a1 * wA.x;
                acc[0][2] += a2 * wA.x; acc[0][3] += a3 * wA.x;
                acc[1][0] += a0 * wA.y; acc[1][1] += a1 * wA.y;
                acc[1][2] += a2 * wA.y; acc[1][3] += a3 * wA.y;
                acc[2][0] += a0 * wA.z; acc[2][1] += a1 * wA.z;
                acc[2][2] += a2 * wA.z; acc[2][3] += a3 * wA.z;
                acc[3][0] += a0 * wA.w; acc[3][1] += a1 * wA.w;
                acc[3][2] += a2 * wA.w; acc[3][3] += a3 * wA.w;
                acc[4][0] += a0 * wB.x; acc[4][1] += a1 * wB.x;
                acc[4][2] += a2 * wB.x; acc[4][3] += a3 * wB.x;
                acc[5][0] += a0 * wB.y; acc[5][1] += a1 * wB.y;
                acc[5][2] += a2 * wB.y; acc[5][3] += a3 * wB.y;
                acc[6][0] += a0 * wB.z; acc[6][1] += a1 * wB.z;
                acc[6][2] += a2 * wB.z; acc[6][3] += a3 * wB.z;
                acc[7][0] += a0 * wB.w; acc[7][1] += a1 * wB.w;
                acc[7][2] += a2 * wB.w; acc[7][3] += a3 * wB.w;
            }
        }
        __syncthreads();
    }

    // ---- epilogue -----------------------------------------------------------
    int oh = oh0 + ty;
    if (oh < OH) {
        const size_t out_batch_stride = (size_t)OC * OH * OW;
        float* out_base = output + (size_t)n * out_batch_stride +
                          (size_t)oh * OW + ow0 + tx;
        #pragma unroll
        for (int o = 0; o < OC_TILE; ++o) {
            float bv = (bias != nullptr) ? bias[oc0 + o] : 0.f;
            float* dst = out_base + (size_t)(oc0 + o) * OH * OW;
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                int ow = ow0 + tx + 16 * j;
                if (ow < OW) dst[16 * j] = acc[o][j] + bv;
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Generic naive fallback kernel: correct for any configuration.
// ---------------------------------------------------------------------------
__global__ void conv2d_naive_kernel(const float* __restrict__ input,
                                    const float* __restrict__ weight,
                                    const float* __restrict__ bias,
                                    float* __restrict__ output,
                                    int N, int IC, int H, int W,
                                    int OC, int OH, int OW,
                                    int KH, int KW,
                                    int sh, int sw,
                                    int ph, int pw,
                                    int dh, int dw) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * OC * OH * OW;
    if (idx >= total) return;

    int ow = idx % OW;
    int oh = (idx / OW) % OH;
    int oc = (idx / ((long long)OW * OH)) % OC;
    int n  = idx / ((long long)OW * OH * OC);

    float sum = (bias != nullptr) ? bias[oc] : 0.f;
    for (int ic = 0; ic < IC; ++ic) {
        for (int kh = 0; kh < KH; ++kh) {
            int ih = oh * sh - ph + kh * dh;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < KW; ++kw) {
                int iw = ow * sw - pw + kw * dw;
                if ((unsigned)iw >= (unsigned)W) continue;
                sum += input[((size_t)(n * IC + ic) * H + ih) * W + iw] *
                       weight[((size_t)(oc * IC + ic) * KH + kh) * KW + kw];
            }
        }
    }
    output[((size_t)(n * OC + oc) * OH + oh) * OW + ow] = sum;
}

// ---------------------------------------------------------------------------
// Host entry point
// ---------------------------------------------------------------------------
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias_opt,
                             int64_t sh, int64_t sw,
                             int64_t ph, int64_t pw,
                             int64_t dh, int64_t dw) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only fp32 supported");
    auto x_c = x.contiguous();
    auto w_c = w.contiguous();

    const int N  = x_c.size(0), IC = x_c.size(1), H = x_c.size(2), W = x_c.size(3);
    const int OC = w_c.size(0), KH = w_c.size(2), KW = w_c.size(3);
    const int OH = (H + 2 * (int)ph - (int)dh * (KH - 1) - 1) / (int)sh + 1;
    const int OW = (W + 2 * (int)pw - (int)dw * (KW - 1) - 1) / (int)sw + 1;

    auto out = torch::empty({N, OC, OH, OW}, x_c.options());
    if (out.numel() == 0) return out;

    const float* bias_ptr = nullptr;
    torch::Tensor b_c;
    if (bias_opt.has_value() && bias_opt->defined()) {
        b_c = bias_opt->contiguous();
        bias_ptr = b_c.data_ptr<float>();
    }

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

    const bool fast_ok = ((int)sh == 1) && ((int)sw == 1) &&
                         (KH == 5) && (KW == 9) &&
                         ((int)dh == 2) && ((int)dw == 3) &&
                         (OC % OC_TILE == 0) && (W % 4 == 0);

    if (fast_ok) {
        dim3 block(NTHREADS);
        dim3 grid((OW + TILE_W - 1) / TILE_W,
                  (OH + TILE_H - 1) / TILE_H,
                  N * (OC / OC_TILE));
        conv2d_direct_kernel<5, 9, 2, 3><<<grid, block, 0, stream>>>(
            x_c.data_ptr<float>(), w_c.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N, IC, H, W, OC, OH, OW,
            (int)ph, (int)pw);
    } else {
        long long total = (long long)N * OC * OH * OW;
        int threads = 256;
        long long blocks = (total + threads - 1) / threads;
        conv2d_naive_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            x_c.data_ptr<float>(), w_c.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N, IC, H, W, OC, OH, OW,
            KH, KW, (int)sh, (int)sw, (int)ph, (int)pw, (int)dh, (int)dw);
    }
    return out;
}
"""

cpp_source = (
    "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, "
    "c10::optional<torch::Tensor> bias_opt, "
    "int64_t sh, int64_t sw, int64_t ph, int64_t pw, int64_t dh, int64_t dw);"
)

_custom_conv = load_inline(
    name="custom_conv2d_direct",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv2d accelerated with a custom register-blocked, shared-memory-tiled
    direct convolution CUDA kernel.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: int = 1, padding: tuple = (0, 0),
                 dilation: tuple = (1, 1), bias: bool = False):
        super(ModelNew, self).__init__()
        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size)
        if isinstance(padding, int):
            padding = (padding, padding)
        if isinstance(dilation, int):
            dilation = (dilation, dilation)

        self.stride = stride
        self.padding = padding
        self.dilation = dilation

        # Identical parameter initialization to nn.Conv2d.
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv2d.weight
        b = self.conv2d.bias if self.conv2d.bias is not None else None
        if x.is_cuda and x.dtype == torch.float32:
            return _custom_conv.conv2d_forward(
                x, w, b,
                self.stride, self.stride,
                self.padding[0], self.padding[1],
                self.dilation[0], self.dilation[1])
        return self.conv2d(x)
```