You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.

Key Technologies Used:

Inline CUDA Extension in PyTorch: Uses torch.utils.cpp_extension.load_inline() to compile and load CUDA code directly within Python, providing seamless integration without external compilation steps.

Simplified 1D Kernel Design: Implements a streamlined kernel optimized for 1D padding operations, focusing on width dimension processing only.

Direct Global Memory Access: Unlike the 2D/3D versions, this implementation accesses global memory directly without shared memory caching, suitable for the simpler 1D case.

1D Thread Blocking: Employs 1D thread blocks (BLOCK_SIZE = 256) for efficient parallelization across the width dimension.

Grid-Strided Loop Pattern: Uses a strided loop (for (int j = tid; j < W_out; j += BLOCK_SIZE)) to distribute work across threads and handle arbitrary output sizes.

Inline Reflection Logic: Implements reflection indexing directly within the kernel using conditional statements, avoiding separate device function calls.

Batched Channel Processing: Processes multiple batches and channels concurrently through 2D grid dimensions (grid_dim(N, C)).

Memory Layout Optimization: Leverages the natural memory layout of 3D tensors (N, C, W) with straightforward stride calculations.

Comprehensive Error Checking: Includes validation for tensor dimensions, CUDA requirements, contiguity, and padding bounds.

Performance Optimizations:

Minimal kernel design with no synchronization overhead

Coalesced memory access patterns for 1D data

Grid-strided loops for optimal load balancing

Restricted pointers for compiler optimization

Direct indexing calculations

Architecture Features:

Separate C++ interface declaration and CUDA implementation

Template-style parameter passing for padding values

Automatic output tensor allocation with correct dimensions

Efficient handling of 3D tensor layout (N, C, W)

Key Differences from 2D/3D Versions:

No shared memory usage (simpler access pattern)

Direct reflection calculation in kernel

1D thread blocks instead of 2D/3D

Simpler memory addressing

Reduced computational complexity


Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
import torch
import torch.nn as nn
import torch.nn.functional as F

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
BATCH_SIZE = 32
CHANNELS = 64
WIDTH = 128  # W_in
PADDING = (3, 1)  # (padding_left, padding_right)


# -------------------------------------------------------------

class Model(nn.Module):
    """
    nn.ReflectionPad1d 的纯 PyTorch 基准实现
    (使用 F.pad)
    """

    def __init__(self, padding):
        super().__init__()

        if isinstance(padding, int):
            # F.pad 需要 (left, right) 格式
            self.padding_tuple = (padding, padding)
        else:
            self.padding_tuple = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, pad_dim_1_left, ...)
        # 因为我们只 pad 最后一个维度 (dim -1)，所以元组是 (pad_L, pad_R)
        return F.pad(x, self.padding_tuple, mode='reflect')


def get_inputs():
    """
    生成一个 (N, C, W) 形状的输入
    """
    x = torch.randn(BATCH_SIZE, CHANNELS, WIDTH, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [PADDING]
