为何thrust::reduce_by_key比基于atomicAdd()的thrust::for_each慢近75倍？

问题描述 #

我对thrust::reduce_by_key的性能不满意，尝试多种改写方式（包括移除permutation iterator）后性能提升甚微。但改用基于atomicAdd()的thrust::for_each实现后，获得了近75倍的性能提升！两种实现结果完全一致。请问造成这种巨大性能差异的最主要原因是什么？

完整对比代码 #

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <ctime>
#include <iostream>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/discard_iterator.h>
#include <thrust/sort.h>

constexpr int NumberOfOscillators = 100;
int SeedRange = 500;

struct GetProduct
{
    template<typename Tuple>
    __host__ __device__
        int operator()(const Tuple & t)
    {
        return  thrust::get<0>(t) * thrust::get<1>(t);
    }
};

int main()
{
    using namespace std;
    using namespace thrust::placeholders;
    
    /* BEGIN INITIALIZATION */

    thrust::device_vector<int> dv_OscillatorsVelocity(NumberOfOscillators);
    thrust::device_vector<int> dv_outputCompare(NumberOfOscillators);
    thrust::device_vector<int> dv_Connections_Strength((NumberOfOscillators - 1) * NumberOfOscillators);
    thrust::device_vector<int> dv_Connections_Active((NumberOfOscillators - 1) * NumberOfOscillators);
    thrust::device_vector<int> dv_Connections_TerminalOscillatorID_Map(0);
    thrust::device_vector<int> dv_Permutation_Connections_To_TerminalOscillators((NumberOfOscillators - 1) * NumberOfOscillators);
    thrust::device_vector<int> dv_Connection_Keys((NumberOfOscillators - 1) * NumberOfOscillators);
    
    srand((unsigned int)time(NULL));

    thrust::fill(dv_OscillatorsVelocity.begin(), dv_OscillatorsVelocity.end(), 0);

    for (int c = 0; c < NumberOfOscillators * (NumberOfOscillators - 1); c++)
    {
        dv_Connections_Strength[c] = (rand() % SeedRange) - (SeedRange / 2);

        dv_Connections_Active[c] = 0;
    }

    int curOscillatorIndx = -1;
    for (int c = 0; c < NumberOfOscillators * NumberOfOscillators; c++)
    {
        if (c % NumberOfOscillators == 0)
        {
            curOscillatorIndx++;
        }

        if (c % NumberOfOscillators != curOscillatorIndx)
        {
            dv_Connections_TerminalOscillatorID_Map.push_back(c % NumberOfOscillators);
        }
    }

    for (int n = 0; n < NumberOfOscillators; n++)
    {
        for (int p = 0; p < NumberOfOscillators - 1; p++)
        {
            thrust::copy_if(
                thrust::device,
                thrust::make_counting_iterator<int>(0),
                thrust::make_counting_iterator<int>(dv_Connections_TerminalOscillatorID_Map.size()), // indices from 0 to N
                dv_Connections_TerminalOscillatorID_Map.begin(), // array data
                dv_Permutation_Connections_To_TerminalOscillators.begin() + (n * (NumberOfOscillators - 1)), // result will be written here
                _1 == n);
        }
    }

    for (int c = 0; c < NumberOfOscillators * (NumberOfOscillators - 1); c++)
    {
        dv_Connection_Keys[c] = c / (NumberOfOscillators - 1);
    }

    /* END INITIALIZATION */


    /* BEGIN COMPARISON */

    auto t = clock();

    for (int x = 0; x < 5000; ++x) //Set x maximum to a reasonable number while testing performance.
    {
        thrust::reduce_by_key(
            thrust::device,
            //dv_Connection_Keys = 0,0,0,...1,1,1,...2,2,2,...3,3,3...
            dv_Connection_Keys.begin(), //keys_first    The beginning of the input key range.
            dv_Connection_Keys.end(), //keys_last   The end of the input key range.
            thrust::make_permutation_iterator(
                thrust::make_transform_iterator(
                    thrust::make_zip_iterator(
                        thrust::make_tuple(
                            dv_Connections_Strength.begin(),
                            dv_Connections_Active.begin()
                        )
                    ),
                    GetProduct()
                ),
                dv_Permutation_Connections_To_TerminalOscillators.begin()
            ), //values_first   The beginning of the input value range.
            thrust::make_discard_iterator(), //keys_output  The beginning of the output key range.
            dv_OscillatorsVelocity.begin() //values_output  The beginning of the output value range.
        );
    }

    std::cout << "iterations    time for original: " << (clock() - t) * (1000.0 / CLOCKS_PER_SEC) << "ms\n" << endl << endl;

    thrust::copy(dv_OscillatorsVelocity.begin(), dv_OscillatorsVelocity.end(), dv_outputCompare.begin());


    t = clock();

    for (int x = 0; x < 5000; ++x) //Set x maximum to a reasonable number while testing performance.
    {
        thrust::for_each(
            thrust::device,
            thrust::make_counting_iterator(0),
            thrust::make_counting_iterator(0) + dv_Connections_Active.size(),
            [
                s = dv_OscillatorsVelocity.size() - 1,
                dv_b = thrust::raw_pointer_cast(dv_OscillatorsVelocity.data()),
                dv_c = thrust::raw_pointer_cast(dv_Permutation_Connections_To_TerminalOscillators.data()), //3,6,9,0,7,10,1,4,11,2,5,8
                dv_ppa = thrust::raw_pointer_cast(dv_Connections_Active.data()),
                dv_pps = thrust::raw_pointer_cast(dv_Connections_Strength.data())
            ] __device__(int i) {
                const int readIndex = i / s;
                atomicAdd(
                    dv_b + readIndex,
                    (dv_ppa[dv_c[i]] * dv_pps[dv_c[i]])
                );
            }
        );
    }

    std::cout << "iterations    time for new: " << (clock() - t) * (1000.0 / CLOCKS_PER_SEC) << "ms\n" << endl << endl;

    std::cout << "***" << (dv_OscillatorsVelocity == dv_outputCompare ? "success" : "fail") << "***\n";

    /* END COMPARISON */

    return 0;
}

额外信息 #

• 测试环境：单GTX 980 TI显卡
• 所有“Connection”向量含100*(100-1)=9900个元素
• dv_Connection_Keys中100个唯一key各对应99个元素
• 编译选项：--expt-extended-lambda

核心原因分析 #

造成两者性能巨大差距的最主要原因是**thrust::reduce_by_key的实现特性与当前数据访问模式的不匹配**，具体拆解为以下几点：

1. reduce_by_key的通用实现固有开销
   Thrust的reduce_by_key是通用归约接口，需要兼容任意键序列（包括未排序场景），内部包含键分组、排序（若键无序）、归约等多步逻辑。即便你的键序列是有序的，它仍要执行键遍历、分组判断等额外操作；再加上嵌套迭代器（permutation_iterator+transform_iterator+zip_iterator）带来的多层函数调用和间接访问，GPU指令流水线难以优化这类复杂逻辑，额外开销显著。

2. 非连续内存访问的效率损耗
   reduce_by_key中使用的permutation_iterator导致值的访问是随机索引模式，全局内存访问完全不连续。GPU的全局内存带宽高度依赖连续访问的合并特性，非连续访问会产生大量未合并内存事务，内存利用率骤降。而for_each+atomicAdd实现中，数据读取部分（dv_ppa、dv_pps、dv_c）是连续内存访问，能充分利用GPU的内存合并优化，内存效率远高于前者。

3. 任务粒度与GPU资源利用率差异
   reduce_by_key会将同一key的元素分配到同一线程块处理，而你的每个key仅对应99个元素——这种小粒度任务无法填满GPU的SM（流式多处理器），导致计算资源大量闲置。for_each实现直接将9900个元素映射为9900个线程，每个线程处理一个元素的计算和原子写入，能最大化利用GPU的并发线程资源，让SM的计算单元和内存单元持续忙碌。

4. 原子操作的实际开销被低估
   虽然atomicAdd存在线程竞争，但你的场景中每个key仅对应99次原子写入，GTX 980 Ti的SM配备硬件原子操作单元，且这类低频次的竞争冲突概率极低。相比reduce_by_key的复杂逻辑和内存低效，原子操作的开销几乎可以忽略。

内容的提问来源于stack exchange，提问作者aiwyn