如何在从较小向量派生的较大向量上使用推力::变换？

Question

输入和启动 arrays：

dv_A = { 5, -3, 2, 6} //4 elements
dv_B = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 }

预期 output：

dv_B = { 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1 }

对于 dv_A{} 中的每个元素，在 dv_B{} 中有 (dv_A.size - 1) 个元素。 这是因为 dv_A 的每个元素都应该在 dv_B 中为每个其他dv_A 元素有一个子元素（即应该排除自身）。 因此，如果 dv_A 中有 4 个元素，则对于每个 dv_A 元素，dv_B 中应该有 3 个元素。

如果对应的 dv_A 元素的值 > 0，我想将每个 dv_B 元素的值转换为 1。对应关系是根据 dv_B 中元素的 position 确定的。 例如：

前 3 个 dv_B 值将由 dv_A[0] 中的值转换，后 3 个 dv_B 值将由 dv_A[1] 中的值转换，等等。

到目前为止，这是我的尝试

thrust::transform(
    dv_B.begin(),
    dv_B.end(),
    thrust::make_transform_iterator(
        dv_A.begin(),
        _1 % dv_A
    ), 
    dv_B.begin(),
    _2 > 0 //When 2nd argument is greater than 0 then set the position in dv_A to 1.
);

Answer 1

串行代码可能如下所示：

for(int i = 0; i < dv_b.size(); i++){
    const int readIndex = i / (dv_a.size() - 1);
    if(dv_a[readIndex] > 0) dv_b[i] = 1;
    else dv_b[i] = 0;
}

可以使用for_each轻松编写。 我认为与使用 transform 和各种花哨的迭代器相比，这使代码更清晰。

thrust::for_each(
    thrust::device,
    thrust::make_counting_iterator(0),
    thrust::make_counting_iterator(0) + dv_b.size(),
    [
     s = dv_a.size() - 1,
     dv_a = thrust::raw_pointer_cast(dv_a.data()),
     dv_b = thrust::raw_pointer_cast(dv_b.data())
    ] __device__ (int i){
        const int readIndex = i / s;
        if(dv_a[readIndex] > 0) dv_b[i] = 1;
        else dv_b[i] = 0;
    }
);

Answer 2

将花哨的迭代器创建打包到一个适当命名的工厂 function 中，也使得这个版本非常可读。 特别是如果你不止一次需要这种模式，这个解决方案可能会更优雅。

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

#include <iostream>

// nvcc doesn't seem to like __device__ or __host__ __device__ lambdas in auto
// return types, so I defined this functor instead
template <typename T>
class Div {
    T div_{};
    public:
    Div(T div) : div_{div} {}
    __host__ __device__ T operator()(T in) const noexcept { return in / div_; }
};

// I stole "repeat" from numpy. Another version using modulo (%) and therefore
// repeating the whole input instead of single elements would be called "tile".
template <class It>
auto make_repeat_it_begin(It input, int repetitions) {
    using diff_t = typename It::difference_type;
    return thrust::make_permutation_iterator(
                input,
                thrust::make_transform_iterator(
                    thrust::make_counting_iterator(diff_t{0}),
                    Div{static_cast<diff_t>(repetitions)}));
}

int main() {
    int A[] = {5, -3, 2, 6};
    constexpr int size_A = sizeof(A) / sizeof(A[0]);

    thrust::host_vector<int> hv_A(A, A + size_A);
    thrust::device_vector<int> dv_A(hv_A);
    thrust::device_vector<int> dv_B(size_A * (size_A - 1));

    auto A_repeat_it = make_repeat_it_begin(dv_A.begin(), size_A - 1);
    
    thrust::transform(A_repeat_it, A_repeat_it + dv_B.size(), 
                      dv_B.begin(),
                      [] __device__ (int a){ return a > 0 ? 1 : 0; });

    thrust::host_vector<int> hv_B(dv_B);
    thrust::copy(hv_B.begin(), hv_B.end(),
                 std::ostream_iterator<int>(std::cout, ","));
}