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.

Technologies Used:

PyTorch: Deep learning framework

CUDA: GPU acceleration for parallel computing

C++/CUDA C++: High-performance kernel programming

Inline C++/CUDA Extension (torch.utils.cpp_extension.load_inline): Just-In-Time (JIT) compilation of custom operators

Hamming Distance: Measure of dissimilarity between binary strings/vectors

Block-Level Parallel Reduction: Uses shared memory and tree reduction within thread blocks

One Block Per Sample: Each CUDA block processes one complete sample pair (N blocks for N samples)

Strided Memory Access: Threads process elements with stride equal to block size for coalesced memory access

Shared Memory Optimization: Uses __shared__ array for intermediate results and reduction

Tree Reduction Pattern: Binary tree reduction within block using __syncthreads()

Integer Arithmetic: Optimized for int64_t data type comparisons

Type Conversion: Converts final integer count to float32 output

Memory Coalescing: Ensures contiguous tensor layout for optimal memory access

Input Validation: Comprehensive type and shape checking in C++ wrapper

Grid-Stride Loops: Efficiently handles variable dimension sizes

Compiler Optimization: Uses -O3 flag for maximum performance

Tensor Contiguity Enforcement: Ensures optimal memory layout in PyTorch wrapper

Automatic Type Promotion: Converts input tensors to torch.long if needed



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


N_BATCH = 100
D_VECTOR = 128

# 假设 dim=1
DIM = 1


class Model(nn.Module):
    """
    汉明距离 (Hamming Distance) 的纯 PyTorch 基准实现
    """
    def __init__(self, dim=1):
        super().__init__()
        # 假设 dim=1
        self.dim = dim

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        # x1, x2: (N, D), 假设为 long/int/bool 类型

        # 1. 比较 (x1 != x2)
        #    如果 x1=[1, 0, 1], x2=[1, 1, 1]
        #    diff=[False, True, False]
        diff = (x1 != x2)

        # 2. 求和 (Sum)
        #    False.sum() = 0, True.sum() = 1
        #    sum([0, 1, 0]) = 1
        #    我们转换为 float 以匹配 CUDA 版本的输出类型
        return torch.sum(diff, dim=self.dim).to(torch.float32)

def get_inputs():
    """
    生成两个 (N, D) 形状的 *整数* 输入
    (汉明距离的标准输入)
    """
    # 随机生成 0 或 1
    input1 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
    input2 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
    return [input1, input2]

def get_init_inputs():
    return [DIM]