c++ - 如何使用跳过部分 device_vector 的自定义仿函数实现 thrust::transform？

Question

我正在从事一个项目(本质上是一个物理模拟)，我需要在多个时间步长上对大量节点执行计算。我目前通过编写一个在 thrust::transform 中调用的自定义仿函数来实现每种类型的计算。

作为一个最小的例子(使用伪代码)，假设我有一些数据都共享一个共同的结构，但可以分解成不同的类型(A、B 和 C)，例如都有一个

double value.

因此，我将此数据存储在单个 device_vector 中，如下所示:

class Data {
    thrust::device_vector<double> values;
    unsigned values_begin_A, values_end_A;
    unsigned values_begin_B, values_end_B;
    unsigned values_begin_C, values_end_C;
}

其中类型 A 占据 vector 的第一部分，然后是类型 B，然后是类型 C。为了保持跟踪，我保存了每种类型的开始/结束索引值。

不同类型的数据需要由不同的仿函数作用(例如，仿函数 1 应用于类型 A 和 B；仿函数 2 应用于 A、B 和 C；仿函数 3 应用于 A 和 C)。每个仿函数都需要访问 vector 中值的索引，由 counting_iterator 提供，并将结果存储在单独的 vector 中。

struct my_functor : public thrust::unary_function< thrust::tuple<unsigned, double> , double > {

    __host__ __device__
    double operator() (const thrust::tuple<unsigned, double> index_value) {

        // Do something with the index and value.

        return result;
    }
}

我的问题是我不知道实现最后一个作用于类型 A 和 C 值同时跳过 B 的仿函数的最佳方法。特别是，我正在寻找一种推力友好的解决方案，它可以像我一样合理地扩展添加更多节点类型和更多仿函数(作用于新旧类型的组合)，同时仍然获得并行化的好处。

我想出了四个选项:

选项 1:

对每种数据类型进行一次转换调用，例如

void Option_One(thrust::device_vector<double>& result) {
    // Multiple transform calls.

    thrust::counting_iterator index(0);

    // Apply functor to 'A' values.
    thrust::transform( 
        thrust::make_zip_iterator(thrust::make_tuple(index, values.begin())),
        thrust::make_zip_iterator(thrust::make_tuple(index, values.begin())) + values_end_A,
        result.begin(),
        my_functor());

    // Apply functor to 'C' values.
    thrust::transform( 
        thrust::make_zip_iterator(thrust::make_tuple(index, values.begin())) + values_begin_C,
        thrust::make_zip_iterator(thrust::make_tuple(index, values.begin())) + values_end_C,
        result.begin() + values_begin_C,
        my_functor());
}

这看起来相当简单，但以牺牲效率为代价，因为我牺牲了并行计算 A 和 C 的能力。

选项 2:

将值复制到临时 vector 中，对临时 vector 调用转换，然后将临时结果复制回结果中。这看起来像很多来回复制，但只允许在 A 和 C 上一起调用一次转换。

void Option_Two(thrust::device_vector<double>& result) {

    // Copy 'A' and 'C' values into temporary vector
    thrust::device_vector<double> temp_values_A_and_C(size_A + size_C);
    thrust::copy(values.begin(), values.begin() + values_end_A, temp_values_A_and_C.begin());
    thrust::copy(values.begin() + values_begin_C, values.begin() + values_end_C, temp_values_A_and_C.begin() + values_end_A);

    // Store results in temporary vector.
    thrust::device_vector<double> temp_results_A_and_C(size_A + size_C);

    thrust::transform( 
        thrust::make_zip_iterator(thrust::make_tuple(index, temp_values_A_and_C.begin())),
        thrust::make_zip_iterator(thrust::make_tuple(index, temp_values_A_and_C.begin())) + size_A + size_C,
        temp_results_A_and_C.begin(),
        my_functor());


    // Copy temp results back into result
    // ....
}

选项 3:

对所有值调用转换，但更改仿函数以检查索引并仅对 A 或 C 范围内的索引起作用。

struct my_functor_with_index_checking : public thrust::unary_function< thrust::tuple<unsigned, double> , double > {

    __host__ __device__
    double operator() (const thrust::tuple<unsigned, double> index_value) {

        if ( (index >= values_begin_A && index <= values_end_A ) ||
            ( index >= values_begin_C && index <= values_end_C ) ) {

                // Do something with the index and value.
                return result;
             }
        else {
            // Do nothing;
            return 0; //Result is 0 by default.
        }
    }
}

void Option_Three(thrust::device_vector<double>& result) {

    // Apply functor to all values, but check index inside functor.
    thrust::transform( 
        thrust::make_zip_iterator(thrust::make_tuple(index, values.begin())),
        thrust::make_zip_iterator(thrust::make_tuple(index, values.begin())) + values.size(),
        result.begin(),
        my_functor_with_index_checking());
}

选项 4:

我想出的最后一个选择是创建一个基于 counting_iterator 的自定义迭代器，该迭代器通常在 A 范围内计数，但一旦到达末尾就会跳到 C 的开头A 的。这似乎是一个优雅的解决方案，但我不知道该怎么做。

void Option_Four(thrust::device_vector<double>& result) {

    // Create my own version of a counting iterator
    // that skips from the end of 'A' to the beginning of 'C'
    // I don't know how to do this!
    FancyCountingIterator fancyIndex(0); 

    thrust::transform( 
        thrust::make_zip_iterator(thrust::make_tuple(fancyIndex, values.begin())),
        thrust::make_zip_iterator(thrust::make_tuple(fancyIndex, values.begin())) + values.size(),
        result.begin(),
        my_functor());
}

Answer 1

将 permutation_iterator 与自定义 transform_iterator(您正在寻找的奇特迭代器)结合使用。

Data d; //assuming this has values.
unsigned A_size = d.values_end_A - d.values_begin_A;
unsigned C_size = d.values_end_C - d.values_begin_C;
auto A_C_index_iter = thrust::make_transform_iterator( thrust::make_counting_iterator(0), 
[&]__device__(int i) {
  if (i<A_size)
    return i+d.values_begin_A; 
  else 
    return (i-A_size)+d.values_begin_C;
});
auto permuted_input_iter = thrust::make_permutation_iterator(values.begin(), A_C_index_iter);
auto permuted_output_iter = thrust::make_permutation_iterator(result.begin(), A_C_index_iter);
thrust::transform(permuted_input_iter, permuted_input_iter + A_size + C_size, permuted_output_iter);

这利用了完全并行性 (A_size + C_size)。