ReLU

原始题目：LeetGPU - ReLU

题目描述

编写一个 GPU 程序，对 32 位浮点数向量执行 ReLU（Rectified Linear Unit）激活函数。ReLU 将所有负值设为零，正值保持不变：

$$\text{ReLU}(x) = \max(0, x)$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 中。

示例

示例 1

Input:  [-2.0, -1.0, 0.0, 1.0, 2.0]
Output: [0.0, 0.0, 0.0, 1.0, 2.0]

示例 2

Input:  [-3.5, 0.0, 4.2]
Output: [0.0, 0.0, 4.2]

约束条件

• $1 \le N \le 100{,}000{,}000$。
• 性能测试在 $N = 25{,}000{,}000$ 的规模下进行。

解题思路

ReLU 是最简单的逐元素激活函数之一，每个元素独立计算 $\max(0, x)$，完全数据并行。计算强度极低（1 FLOP + 1 比较 + 2 次 4 字节内存访问），属于典型的内存带宽受限内核。

代码实现

CUDA

#include <cuda_runtime.h>

__global__ void relu_kernel(const float* input, float* output, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        output[i] = fmaxf(0.0f, input[i]);
    }
}

// Grid-Stride Loop
__global__ void relu_grid_stride(const float* input, float* output, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < N; i += stride) {
        output[i] = fmaxf(0.0f, input[i]);
    }
}

extern "C" void solve(const float* input, float* output, int N) {
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    relu_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, N);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl

@triton.jit
def relu_kernel(input_ptr, output_ptr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N
    x = tl.load(input_ptr + idx, mask=mask)
    tl.store(output_ptr + idx, tl.maximum(x, 0.0), mask=mask)