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

Call functor for all combinations in Cuda/Thrust

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

本质上,我想为这些索引的每个组合创建一个 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;
}

我相当确定这很容易实现,我想我只是缺少正确的搜索词。无论如何,感谢所有帮助:)

  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. 向量上的 [] 等操作在设备代码中是不可能的。因此,我们必须将您的仿函数转换为处理原始数据指针,而不是直接进行推力向量运算。

考虑到这些注意事项,以及我们处理您的 ij 索引的方式略有修改,我认为您的要求并不难。

基本策略是按照您的建议创建长度为 N*M 的结果向量,然后在仿函数运算符中创建索引 ij。这样一来,我们只需要将一个索引传递给仿函数,例如使用thrust::transformthrust::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 所建议的。