Sigmoid Activation

原始题目：LeetGPU - Sigmoid Activation

题目描述

编写一个 GPU 程序，对 32 位浮点数向量逐元素应用 sigmoid 激活函数。对于输入向量 $X$ 中的每个元素 $x$，计算：

$$\text{sigmoid}(x) = \frac{1}{1 + e^{-x}}$$

并将结果存储在输出向量 $Y$ 中。sigmoid 函数将任意实数映射到 $(0, 1)$ 区间。

实现要求

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

示例

示例 1

Input:  X = [0.0, 1.0, -1.0, 2.0]
Output: Y = [0.5, 0.7311, 0.2689, 0.8808]

示例 2

Input:  X = [0.5, -0.5, 3.0, -3.0]
Output: Y = [0.6225, 0.3775, 0.9526, 0.0474]

约束条件

• $1 \le N \le 100{,}000{,}000$。
• 输入值为有限的 32 位浮点数。
• 性能测试在 $N = 50{,}000{,}000$ 的规模下进行。

解题思路

Sigmoid 的核心计算瓶颈在 $e^{-x}$。对于大的正输入，$e^{-x} \to 0$，sigmoid $\to 1$；对于大的负输入，$e^{-x}$ 可能很大但分母也会相应变大。数值上 sigmoid 天然有界（$(0,1)$），不需要额外的 max-trick。CUDA 的 expf() 提供 IEEE 标准精度。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>

// 解法一：基础逐元素 sigmoid
__global__ void sigmoid_kernel(const float* X, float* Y, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        Y[i] = 1.0f / (1.0f + expf(-X[i]));
    }
}

// 解法二：Grid-Stride Loop
__global__ void sigmoid_grid_stride(const float* X, float* Y, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < N; i += stride) {
        Y[i] = 1.0f / (1.0f + expf(-X[i]));
    }
}

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

Triton

import triton
import triton.language as tl
import torch


@triton.jit
def sigmoid_kernel(
    X_ptr, Y_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(X_ptr + idx, mask=mask)
    # Triton 编译器自动处理数值稳定性
    tl.store(Y_ptr + idx, tl.sigmoid(x), mask=mask)