CUDA Thrust迭代器：高效实现device_vector批量填充、复制等操作

1. 指定索引起始的批量fill/sequence/transform操作 #

低效for循环实现示例 #

#include <thrust/device_vector.h>

int main() {
    thrust::device_vector<int> d_vec(1000, 0);
    int start_idx = 100;
    int batch_size = 500;

    // 批量fill
    for (int i = start_idx; i < start_idx + batch_size; ++i) {
        d_vec[i] = 42;
    }

    // 批量生成序列
    int seq_start = 10;
    for (int i = start_idx; i < start_idx + batch_size; ++i) {
        d_vec[i] = seq_start + (i - start_idx);
    }

    // 批量transform（元素自增运算）
    for (int i = start_idx; i < start_idx + batch_size; ++i) {
        d_vec[i] = d_vec[i] * 2 + 1;
    }
    return 0;
}

Thrust迭代器优化实现 #

利用thrust::make_slice_iterator限定操作范围，结合Thrust内置函数实现批量操作：

#include <thrust/device_vector.h>
#include <thrust/fill.h>
#include <thrust/sequence.h>
#include <thrust/transform.h>
#include <thrust/functional.h>
#include <thrust/iterator/slice_iterator.h>

int main() {
    thrust::device_vector<int> d_vec(1000, 0);
    int start_idx = 100;
    int batch_size = 500;

    // 构造切片迭代器，锁定[start_idx, start_idx+batch_size)区间
    auto slice_begin = thrust::make_slice_iterator(d_vec.begin(), start_idx);
    auto slice_end = thrust::make_slice_iterator(d_vec.begin(), start_idx + batch_size);

    // 批量fill
    thrust::fill(slice_begin, slice_end, 42);

    // 批量生成连续序列
    thrust::sequence(slice_begin, slice_end, 10);

    // 批量transform：执行x*2+1运算
    thrust::transform(slice_begin, slice_end, slice_begin, 
        [] __device__ (int x) { return x * 2 + 1; });

    return 0;
}

2. 批量复制子向量到目标向量 #

示例：复制源向量中多个不连续子块到目标向量 #

将源向量中3个起始索引为0/200/500、大小为100的子块，复制到目标向量的连续位置：

#include <thrust/device_vector.h>
#include <thrust/copy.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/iterator/counting_iterator.h>

int main() {
    thrust::device_vector<int> d_src(1000, 0);
    thrust::device_vector<int> d_dst(300, 0);

    // 定义子块起始索引和单块大小
    const int block_starts[] = {0, 200, 500};
    const int block_size = 100;

    // 构造排列迭代器，生成所有需要复制的元素索引
    auto idx_iter = thrust::make_permutation_iterator(
        thrust::counting_iterator<int>(0),
        thrust::make_transform_iterator(
            thrust::counting_iterator<int>(0),
            [block_starts, block_size] __device__ (int i) {
                int block_idx = i / block_size;
                int offset = i % block_size;
                return block_starts[block_idx] + offset;
            }
        )
    );

    // 执行批量复制
    thrust::copy(idx_iter, idx_iter + 3 * block_size, d_dst.begin());
    return 0;
}

3. 批量reduce_sum求和 #

示例：将向量分块后独立求和 #

将1000元素的向量分成10个块，每个块独立计算求和结果：

#include <thrust/device_vector.h>
#include <thrust/reduce.h>
#include <thrust/transform.h>
#include <thrust/iterator/counting_iterator.h>

int main() {
    thrust::device_vector<int> d_vec(1000, 1);
    const int num_blocks = 10;
    const int block_size = d_vec.size() / num_blocks;
    thrust::device_vector<int> d_block_sums(num_blocks, 0);

    // 用transform+counting_iterator实现批量归约，避免显式for循环
    thrust::transform(
        thrust::counting_iterator<int>(0),
        thrust::counting_iterator<int>(num_blocks),
        d_block_sums.begin(),
        [&d_vec, block_size] __device__ (int i) {
            auto block_begin = d_vec.begin() + i * block_size;
            return thrust::reduce(block_begin, block_begin + block_size, 0, thrust::plus<int>());
        }
    );
    return 0;
}

4. 嵌套批量归约求和及性能对比 #

方案1：嵌套双层归约 #

先将1000元素分成10个100元素的外层块，每个外层块再分成10个10元素的内层子块，先求内层子块和，再聚合外层块和：

#include <thrust/device_vector.h>
#include <thrust/reduce.h>
#include <thrust/transform.h>
#include <thrust/iterator/counting_iterator.h>

int main() {
    thrust::device_vector<int> d_vec(1000, 1);
    const int outer_block_size = 100;
    const int inner_block_size = 10;
    const int num_outer_blocks = d_vec.size() / outer_block_size;
    const int num_inner_blocks = outer_block_size / inner_block_size;

    // 步骤1：计算所有内层子块的和
    thrust::device_vector<int> d_inner_sums(num_outer_blocks * num_inner_blocks, 0);
    thrust::transform(
        thrust::counting_iterator<int>(0),
        thrust::counting_iterator<int>(num_outer_blocks * num_inner_blocks),
        d_inner_sums.begin(),
        [&d_vec, outer_block_size, inner_block_size] __device__ (int i) {
            int outer_idx = i / num_inner_blocks;
            int inner_idx = i % num_inner_blocks;
            auto begin = d_vec.begin() + outer_idx*outer_block_size + inner_idx*inner_block_size;
            return thrust::reduce(begin, begin + inner_block_size, 0, thrust::plus<int>());
        }
    );

    // 步骤2：聚合外层块的和
    thrust::device_vector<int> d_outer_sums(num_outer_blocks, 0);
    thrust::transform(
        thrust::counting_iterator<int>(0),
        thrust::counting_iterator<int>(num_outer_blocks),
        d_outer_sums.begin(),
        [&d_inner_sums, num_inner_blocks] __device__ (int i) {
            auto begin = d_inner_sums.begin() + i * num_inner_blocks;
            return thrust::reduce(begin, begin + num_inner_blocks, 0, thrust::plus<int>());
        }
    );
    return 0;
}

方案2：单层批量归约 #

直接将1000元素分成10个100元素块求和，代码同第3部分示例。

性能对比（测试环境：RTX 3090，向量大小1e7） #

通用性能对比测试结果 #

简单for循环vsThrust迭代器（操作1e6元素） #

不同迭代器在归约操作中的性能（1e7元素批量归约） #

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