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

Tanimoto Coefficient (Jaccard Index): Similarity measure between sets or vectors

Fused Kernel Design: Combines dot product, squared sums, and division in single kernel

Vectorized Memory Access (float4): Uses 128-bit wide loads (4 floats) for improved memory bandwidth

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

Shared Memory Reduction: Uses __shared__ arrays for efficient block-level parallel reduction

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

Precision Control: Uses volatile keyword to disable FMA (Fused Multiply-Add) for exact numerical matching

Compiler Precision Flags: Uses -ftz=false, -prec-div=true, -prec-sqrt=true for high precision

Dynamic Shared Memory Allocation: Allocates shared memory for three reduction buffers at kernel launch

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

Numerical Stability: Adds epsilon (eps) to prevent division by zero

Memory Coalescing: Ensures contiguous tensor layout for optimal memory access patterns

Automatic Type Conversion: Converts inputs to float32 and CUDA device if needed

Boundary Checking: Validates block indices to prevent out-of-bounds access

Three-Accumulator Design: Maintains separate accumulators for dot product, x squared, and y squared

Exact Numerical Matching: Designed to 100% match PyTorch's step-by-step floating point rounding

Optimized Memory Layout: Uses vectorized pointer arithmetic with float4 type casting

Efficient Resource Utilization: Maximizes memory bandwidth while maintaining numerical precision


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, C, H, W = 32, 64, 56, 56
EPS = 1e-6


class Tanimoto(nn.Module):


    def __init__(self, eps=1e-6):
        super().__init__()
        self.eps = eps
        # 我们将在 C, H, W 维度上进行归约
        self.reduction_dims = (1, 2, 3)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

        x_dot_y = torch.sum(x * y, dim=self.reduction_dims)


        x_norm_sq = torch.sum(x * x, dim=self.reduction_dims)
        y_norm_sq = torch.sum(y * y, dim=self.reduction_dims)


        denominator = x_norm_sq + y_norm_sq - x_dot_y


        similarity = (x_dot_y + self.eps) / (denominator + self.eps)

        return similarity


class Model(nn.Module):


    def __init__(self):
        super().__init__()
        self.op = Tanimoto(EPS)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return self.op(x, y)


def get_inputs():

    x = torch.randn(N, C, H, W, dtype=torch.float32)
    y = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x, y]


def get_init_inputs():

    return []