CUDA 算子优化：Reduce 类

1
__global__ void sum_v1(const float* input, float* output, int N) {
2
    int tid = threadIdx.x;
3
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
4

5
    __shared__ float sdata[256];
6

7
    // Load data to shared memory with boundary check
8
    sdata[tid] = (idx < N) ? input[idx] : 0.0f;
9
    __syncthreads();
10

11
    // Reduction in shared memory
12
    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
13
        if (tid < offset) {
14
            sdata[tid] += sdata[tid + offset];
15
        }
16
        __syncthreads();
17
    }
18

19
    if (tid == 0) output[blockIdx.x] = sdata[0];
20
}

1
__global__ void sum_v4(const float* input, float* output, int N) {
2
    int tid = threadIdx.x;
3
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
4
    int warp_id = tid / 32;
5
    int lane_id = tid % 32;
6
    int num_warps = blockDim.x / 32;
7

8
    __shared__ float sdata[32];
9

10
    // Load data with boundary check
11
    float val = (idx < N) ? input[idx] : 0.0f;
12

13
    // Warp-level reduction
14
    val = warp_reduce(val);
15

16
    // Only first lane of each warp writes to shared memory
17
    if (lane_id == 0) {
18
        sdata[warp_id] = val;
19
    }
20
    __syncthreads();
21

22
    // First warp reduces the warp results
23
    if (warp_id == 0) {
24
        val = (lane_id < num_warps) ? sdata[lane_id] : 0.0f;
25
        val = warp_reduce(val);
26
        if (lane_id == 0) {
27
            atomicAdd(output, val);
28
        }
29
    }
30
}

1
__global__ void sum_v5(const float* input, float* output, int N) {
2
    int tid = threadIdx.x;
3
    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
4
    int warp_id = tid / 32;
5
    int lane_id = tid % 32;
6
    int num_warps = blockDim.x / 32;
7

8
    __shared__ float sdata[32];
9

10
    // Load data with boundary check
11
    float val = 0.0f;
12
    if (idx < N) {
13
        float4 tmp = *reinterpret_cast<const float4 *>(&input[idx]);
14
        val += tmp.x;
15
        val += tmp.y;
16
        val += tmp.z;
17
        val += tmp.w;
18
    }
19

20
    // Warp-level reduction
21
    val = warp_reduce(val);
22

23
    // Only first lane of each warp writes to shared memory
24
    if (lane_id == 0) {
25
        sdata[warp_id] = val;
26
    }
27
    __syncthreads();
28

29
    // First warp reduces the warp results
30
    if (warp_id == 0) {
31
        val = (lane_id < num_warps) ? sdata[lane_id] : 0.0f;
32
        val = warp_reduce(val);
33
        if (lane_id == 0) {
34
            atomicAdd(output, val);
35
        }
36
    }
37
}

1
void softmax_cpu(const float* input, float* output, int N) {
2
    // 找最大值（防止指数溢出）
3
    float max_val = input[0];
4
    for (int i = 1; i < N; i++) {
5
        max_val = std::max(max_val, input[i]);
6
    }
7

8
    // 计算指数和
9
    float sum = 0.0f;
10
    for (int i = 0; i < N; i++) {
11
        output[i] = expf(input[i] - max_val);
12
        sum += output[i];
13
    }
14

15
    // 归一化
16
    for (int i = 0; i < N; i++) {
17
        output[i] /= sum;
18
    }
19
}

1
__inline__ __device__ float warp_reduce_max(float val) {
2
    for (int offset = 16; offset > 0; offset >>= 1) {
3
        val = fmax(val, __shfl_down_sync(0xffffffff, val, offset));
4
    }
5
    return val;
6
}
7

8
__device__ static float atomicMax(float* address, float val) {
9
    int* address_as_i = (int*)address;
10
    int old = *address_as_i, assumed;
11
    do {
12
        assumed = old;
13
        old = atomicCAS(address_as_i, assumed,
14
            __float_as_int(fmax(val, __int_as_float(assumed))));
15
    } while (assumed != old);
16
    return __int_as_float(old);
17
}
18

19
// Kernel 1: 找最大值
20
__global__ void reduce_max(const float* input, float* max_val, int N) {
21
    int tid = threadIdx.x;
22
    int idx = blockIdx.x * blockDim.x + tid;
23
    int warp_id = tid / 32;
24
    int lane_id = tid % 32;
25
    int num_warps = blockDim.x / 32;
26

27
    __shared__ float sdata[32];
28

29
    float val = (idx < N) ? input[idx] : (-FLT_MAX);
30

31
    val = warp_reduce_max(val);
32

33
    if (lane_id == 0) {
34
        sdata[warp_id] = val;
35
    }
36
    __syncthreads();
37

38
    if (warp_id == 0) {
39
        val = (lane_id < num_warps) ? sdata[lane_id] : (-FLT_MAX);
40
        val = warp_reduce_max(val);
41
        if (lane_id == 0) {
42
            atomicMax(max_val, val);
43
        }
44
    }
45
}
46

47
__inline__ __device__ float warp_reduce_sum(float val) {
48
    for (int offset = 16; offset > 0; offset >>= 1) {
49
        val += __shfl_down_sync(0xffffffff, val, offset);
50
    }
51
    return val;
52
}
53

54
// Kernel 2: 计算 exp(x - max) 并求和
55
__global__ void softmax_exp_sum(
56
    const float* input, const float* max_val, float* sum_val, int N
57
) {
58
    int tid = threadIdx.x;
59
    int idx = blockIdx.x * blockDim.x + tid;
60
    int warp_id = tid / 32;
61
    int lane_id = tid % 32;
62
    int num_warps = blockDim.x / 32;
63

64
    __shared__ float sdata[32];
65

66
    float val = (idx < N) ? expf(input[idx] - *max_val) : 0.0f;
67

68
    val = warp_reduce_sum(val);
69

70
    if (lane_id == 0) {
71
        sdata[warp_id] = val;
72
    }
73
    __syncthreads();
74

75
    if (warp_id == 0) {
76
        val = (lane_id < num_warps) ? sdata[lane_id] : 0.0f;
77
        val = warp_reduce_sum(val);
78
        if (lane_id == 0) {
79
            atomicAdd(sum_val, val);
80
        }
81
    }
82
}
83

84
// Kernel 3: 归一化
85
__global__ void softmax_normalize(
86
    const float* input, const float* max_val,
87
    const float* sum_val, float* output, int N
88
) {
89
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
90
    if (idx < N) output[idx] = expf(input[idx] - *max_val) / (*sum_val);
91
}