Leaky ReLU

原始题目：LeetGPU - Leaky ReLU

题目描述

编写一个 GPU 程序，对浮点数向量执行 Leaky ReLU 激活函数。Leaky ReLU 定义为：

$$f(x) = \begin{cases} x, & \text{if } x > 0 \\ \alpha x, & \text{if } x \le 0 \end{cases}$$

其中 $\alpha$ 是一个小的正常数（本题中 $\alpha = 0.01$）。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在向量 output 中。
• 使用 $\alpha = 0.01$ 作为泄露系数。

示例

示例 1

Input:  x = [1.0, -2.0, 3.0, -4.0]
Output: y = [1.0, -0.02, 3.0, -0.04]

示例 2

Input:  x = [-1.5, 0.0, 2.5, -3.0]
Output: y = [-0.015, 0.0, 2.5, -0.03]

约束条件

• $1 \le N \le 100{,}000{,}000$。
• $-1000.0 \le input[i] \le 1000.0$。
• 性能测试在 $N = 50{,}000{,}000$ 的规模下进行。

解题思路

Leaky ReLU 与 ReLU 几乎相同，只是负半轴有一个小的斜率 $\alpha$ 而非完全归零。每个元素独立计算条件分支，但 GPU 上的 warp 内线程如果走向不同分支，会导致warp divergence。好在 Leaky ReLU 的分支对应于输入的正负号，如果输入分布均匀，会有约一半的线程走不同分支。可以通过使用 fma 指令或条件移动（cmov 等效）来避免分支。

代码实现

CUDA

#include <cuda_runtime.h>

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

// 无分支版本：用 fma + 条件移动避免 warp divergence
__global__ void leaky_relu_branchless(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) {
        float x = input[i];
        output[i] = fmaxf(x, 0.01f * x);
    }
}

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

Triton

import triton
import triton.language as tl

@triton.jit
def leaky_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.where(x > 0, x, 0.01 * x), mask=mask)