调用 Cuda/Thrust 中所有组合的函子

Question

我有两个索引集，一个在 [0, N] 范围内，一个在 [0, M] 范围内，其中 N != M。索引用于引用不同 thrust::device_vectors 中的值。

基本上，我想为这些索引的每个组合创建一个 GPU 线程，因此 N*M 个线程。每个线程都应根据索引组合计算一个值，并将结果存储在另一个 thrust::device_vector 中，该索引也基于输入组合。

这似乎是一个相当标准的问题，但我无法找到一种方法来解决这个问题。文档只提到了问题，其中向量的元素 i 需要用另一个向量的元素 i 计算一些东西。有thrust::permutation_iterator，但据我了解，它只给了我重新排序数据的选项，我还必须指定顺序。

一些代码：

#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <iostream>
int main()
{
    // Initialize some data
    const int N = 2;
    const int M = 3;
    thrust::host_vector<int> vec1_host(N);
    thrust::host_vector<int> vec2_host(M);
    vec1_host[0] = 1;
    vec1_host[1] = 5;
    vec2_host[0] = -3;
    vec2_host[1] = 42;
    vec2_host[2] = 9;

    // Copy to device
    thrust::device_vector<int> vec1_dev = vec1_host;
    thrust::device_vector<int> vec2_dev = vec2_host;

    // Allocate device memory to copy results to
    thrust::device_vector<int> result_dev(vec1_host.size() * vec2_host.size());

    // Create functor I want to call on every combination
    struct myFunctor
    {
        thrust::device_vector<int> const& m_vec1;
        thrust::device_vector<int> const& m_vec2;
        thrust::device_vector<int>& m_result;

        myFunctor(thrust::device_vector<int> const& vec1, thrust::device_vector<int> const& vec2, thrust::device_vector<int>& result)
        : m_vec1(vec1), m_vec2(vec2), m_result(result)
        {
        }

        __host__ __device__
        void operator()(size_t i, size_t j) const
        {
            m_result[i + j * m_vec1.size()] = m_vec1[i] + m_vec1[j];
        }
    } func(vec1_dev, vec2_dev, result_dev);

    // How do I create N*M threads, each of which calls func(i, j) ?

    // Copy results back
    thrust::host_vector<int> result_host = result_dev;
    for(int i : result_host)
        std::cout << i << ", ";
    std::cout << std::endl;
    // Expected output:
    // -2, 2, 43, 47, 10, 14
    return 0;
}

我很确定这很容易实现，我想我只是缺少正确的搜索词。无论如何，所有帮助表示赞赏:)

Answer 1

1. 大概在你的仿函数运算符中，而不是这个：

       m_result[i + j * m_vec1.size()] = m_vec1[i] + m_vec1[j];
                                              ^           ^
   

   你的意思是：

       m_result[i + j * m_vec1.size()] = m_vec1[i] + m_vec2[j];
                                              ^           ^
   

2. 我认为可能有很多方法可以解决这个问题，但为了不争论与问题无关的事情，我会尽量接近你给定的代码。

3. 像[] 这样的矢量操作在设备代码中是不可能的。因此，我们必须将您的仿函数转换为处理原始数据指针，而不是直接进行推力向量操作。

有了这些注意事项，以及我们处理您的 i 和 j 索引的方式稍作修改，我认为您的要求并不难。

基本策略是按照您的建议创建一个长度为N*M 的结果向量，然后在仿函数运算符中创建索引i 和j。这样做时，我们只需将一个索引传递给函子，例如使用thrust::transform 或 thrust::for_each 创建我们的输出：

$ cat t79.cu
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/for_each.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/execution_policy.h>
#include <iostream>

    struct myFunctor
    {
        const int *m_vec1;
        const int *m_vec2;
        int *m_result;
        size_t v1size;
        myFunctor(thrust::device_vector<int> const& vec1, thrust::device_vector<int> const& vec2, thrust::device_vector<int>& result)
        {
          m_vec1 = thrust::raw_pointer_cast(vec1.data());
          m_vec2 = thrust::raw_pointer_cast(vec2.data());
          m_result = thrust::raw_pointer_cast(result.data());
          v1size = vec1.size();
        }

        __host__ __device__
        void operator()(const size_t  x) const
        {
            size_t i = x%v1size;
            size_t j = x/v1size;
            m_result[i + j * v1size] = m_vec1[i] + m_vec2[j];
        }
    };

int main()
{
    // Initialize some data
    const int N = 2;
    const int M = 3;
    thrust::host_vector<int> vec1_host(N);
    thrust::host_vector<int> vec2_host(M);
    vec1_host[0] = 1;
    vec1_host[1] = 5;
    vec2_host[0] = -3;
    vec2_host[1] = 42;
    vec2_host[2] = 9;

    // Copy to device
    thrust::device_vector<int> vec1_dev = vec1_host;
    thrust::device_vector<int> vec2_dev = vec2_host;

    // Allocate device memory to copy results to
    thrust::device_vector<int> result_dev(vec1_host.size() * vec2_host.size());

    // How do I create N*M threads, each of which calls func(i, j) ?
    thrust::for_each_n(thrust::device, thrust::counting_iterator<size_t>(0), (N*M), myFunctor(vec1_dev, vec2_dev, result_dev));
    // Copy results back
    thrust::host_vector<int> result_host = result_dev;
    for(int i : result_host)
        std::cout << i << ", ";
    std::cout << std::endl;
    // Expected output:
    // -2, 2, 43, 47, 10, 14
    return 0;
}
$ nvcc -std=c++11 -arch=sm_61 -o t79 t79.cu
$ ./t79
-2, 2, 43, 47, 10, 14,
$

回想起来，我认为这或多或少正是@eg0x20 的建议。