咨询CUDA Thrust中基于索引数组的元素计算实现方法

嘿，这个问题我熟！用CUDA Thrust实现你说的CPU逻辑其实很直接，核心就是用permutation迭代器来根据索引数组访问大数组的元素，再配合thrust::transform完成逐元素计算。我给你一步步拆解，再附上完整代码示例：

核心思路 #

你的CPU逻辑是：遍历索引数组，用每个索引从data1和data2中取对应元素，执行do_something后存入结果数组。在Thrust里，thrust::permutation_iterator正好用来做「通过索引数组间接访问另一个数组」的操作，再结合thrust::transform就能完成批量的元素计算。

而且你提到索引数组是已排序的，这反而能帮Thrust优化内存访问——因为排序后的索引会让GPU对data1/data2的访问是连续的，触发内存合并（coalescing），性能会比随机索引好很多。

完整实现代码 #

首先，你需要把do_something改成GPU可执行的函数（用__device__修饰），然后用Thrust的设备向量和迭代器完成计算：

#include <thrust/device_vector.h>
#include <thrust/transform.h>
#include <thrust/iterator/permutation_iterator.h>
#include <vector>

// 定义GPU端可执行的do_something函数
__device__ float do_something(float a, float b) {
    // 这里替换成你的实际计算逻辑，比如a + b、a*b等
    return a * b + 0.5f;
}

int main() {
    // CPU端原始数据
    std::vector<int> index(20);
    std::vector<float> data1(100), data2(100), result(20);
    
    // 假设这里已经给index、data1、data2赋过值了
    // ...

    // 1. 将CPU数据拷贝到GPU设备向量
    thrust::device_vector<int> d_index(index.begin(), index.end());
    thrust::device_vector<float> d_data1(data1.begin(), data1.end());
    thrust::device_vector<float> d_data2(data2.begin(), data2.end());
    thrust::device_vector<float> d_result(index.size()); // 结果向量大小和索引数组一致

    // 2. 用transform+permutation_iterator完成核心计算
    thrust::transform(
        // 第一个输入：data1[index[i]]的迭代器
        thrust::make_permutation_iterator(d_data1.begin(), d_index.begin()),
        thrust::make_permutation_iterator(d_data1.begin(), d_index.end()),
        // 第二个输入：data2[index[i]]的迭代器
        thrust::make_permutation_iterator(d_data2.begin(), d_index.begin()),
        // 输出结果的起始迭代器
        d_result.begin(),
        // 自定义计算lambda（也可以直接传do_something，如果它符合函数对象要求）
        [] __device__ (float val1, float val2) {
            return do_something(val1, val2);
        }
    );

    // 3. 将GPU结果拷贝回CPU
    thrust::copy(d_result.begin(), d_result.end(), result.begin());

    return 0;
}

关键部分解释 #

• thrust::make_permutation_iterator(data_begin, index_begin)：这个迭代器会返回data_begin[index[i]]的值，完美对应你CPU代码里的data1[index[i]]和data2[index[i]]。
• thrust::transform：它会遍历两个输入迭代器的元素，对每一组元素执行你指定的计算，然后把结果存入输出迭代器。
• 排序索引的优势：因为索引是排序好的，GPU访问data1/data2时是按内存地址递增的顺序读取，这会触发GPU的内存合并访问，大幅提升内存读取效率，这是随机索引做不到的。

额外注意点 #

• 确保你的do_something函数是__device__修饰的，或者是能在GPU上执行的函数对象（比如继承thrust::binary_function的仿函数）。
• 如果你的CUDA版本比较旧（低于7.0），可能不支持lambda表达式，这时候可以把计算逻辑写成一个仿函数：

  struct DoSomething {
      __device__ float operator()(float a, float b) {
          return do_something(a, b);
      }
  };
  

  然后在transform里传DoSomething()即可。

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