如何加速 CUDA 中的 saxpy 列表

Question

我有一个矩阵和一个长度为 BATCH_SIZE 的三元组列表：

float matrix[3000][3000]

(float vector[3000], uint32_t index, float alpha)

对于每个三连音，我用以下方式进行 saxpy：

matrix[index] += vector * alpha.

我做了很多批次。

我想知道在 CUDA 中加速代码的最佳方法是什么。我也愿意使用 Python 库。

Answer 1

这是使用thrust 的可能方法。它大致遵循 cmets 中概述的顺序。基本思想是将所有内容展平，并将添加到同一向量元素的元素减少在一起：

1. 将每个vector 元素乘以其对应的alpha。
2. 创建索引以允许对 vector 元素进行排序，从而将类似索引的元素组合在一起
3. sort_by_key，将相似索引元素组合在一起
4. reduce_by_key，将相似索引元素组合在一起
5. 将减少的值添加到matrix 中的适当元素

以下示例显示了上述方法，并从时序角度将其与 loop-saxpy 方法进行了比较。在我的 Tesla V100 CUDA 10.1 上，“并行”方法比 loop-saxpy 方法快大约 3 倍，BATCH_SIZE 为 1000，index 被限制在相对较小的占用空间：

$ cat t1714.cu
#include <iostream>
#include <thrust/device_vector.h>
#include <thrust/transform.h>
#include <thrust/sort.h>
#include <thrust/reduce.h>
#include <thrust/fill.h>
#include <thrust/generate.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/copy.h>

#include <cstdlib>
#include <time.h>
#include <sys/time.h>
#define USECPSEC 1000000ULL

unsigned long long dtime_usec(unsigned long long start){

  timeval tv;
  gettimeofday(&tv, 0);
  return ((tv.tv_sec*USECPSEC)+tv.tv_usec)-start;
}


typedef float dt;
const int BATCH_SIZE = 1000;
const int ds = 3000;
const int RANGE = 10;

struct my_gen
{
  int range;
  my_gen(int r) : range(r) {};
  __host__
  int operator ()(){
    return rand()%range;
    }
};

//struct my_seq


using namespace thrust::placeholders;

int main(){

// data setup

  thrust::device_vector<dt> matrix(ds*ds);
  thrust::device_vector<dt> result1 = matrix;
  thrust::device_vector<dt> vector(ds*BATCH_SIZE);
  thrust::host_vector<int> h_index(BATCH_SIZE);
  thrust::device_vector<dt>  alpha(BATCH_SIZE);

  thrust::generate(h_index.begin(), h_index.end(), my_gen(RANGE));
  thrust::fill(vector.begin(), vector.end(), 1.0);
  thrust::fill(alpha.begin(), alpha.end(), 1.0);

// first, time serial op

  thrust::host_vector<dt> h_alpha = alpha;

  unsigned long long my_dt = dtime_usec(0);
  for (int i = 0; i < BATCH_SIZE; i++)
    thrust::transform(vector.begin()+(i*ds), vector.begin()+((i+1)*ds), result1.begin()+(h_index[i]*ds), result1.begin()+(h_index[i]*ds), _1*h_alpha[i]+_2);
  cudaDeviceSynchronize();
  my_dt = dtime_usec(my_dt);
  std::cout << "serial time: " << my_dt/(float)USECPSEC << " seconds" << std::endl;


// now time parallel op
  thrust::device_vector<int> vi(ds*BATCH_SIZE);
  thrust::device_vector<int> index = h_index;
  thrust::device_vector<int> keys(ds*BATCH_SIZE);
  thrust::device_vector<dt> vr(ds*BATCH_SIZE);
  auto my_ti = thrust::make_transform_iterator(thrust::counting_iterator<int>(0), _1%ds);
  my_dt = dtime_usec(0);
  thrust::transform(vector.begin(), vector.end(), thrust::make_permutation_iterator(alpha.begin(), thrust::make_transform_iterator(thrust::counting_iterator<int>(0), _1/ds)), vector.begin(), _1*_2); // multiply by alpha
  thrust::transform(my_ti, my_ti+(ds*BATCH_SIZE), thrust::make_transform_iterator(thrust::make_permutation_iterator(index.begin(), thrust::make_transform_iterator(thrust::counting_iterator<int>(0), _1/ds)), _1*ds), vi.begin(), _1 + _2);
  thrust::sort_by_key(vi.begin(), vi.end(), vector.begin()); // collect like indices
  int key_size = (thrust::reduce_by_key(vi.begin(), vi.end(), vector.begin(), keys.begin(), vr.begin())).first - keys.begin(); // add like indices together
  auto my_pi = thrust::make_permutation_iterator(matrix.begin(), keys.begin());
  thrust::transform(vr.begin(), vr.begin()+key_size, my_pi, my_pi, _1 + _2); // perform final add
  cudaDeviceSynchronize();
  my_dt = dtime_usec(my_dt);
  std::cout << "parallel time: " << my_dt/(float)USECPSEC << " seconds" << std::endl;
  thrust::host_vector<dt> h_matrix = matrix;
  thrust::host_vector<dt> h_result1 = result1;
  for (int i = 0; i < ds*ds; i++) if (h_matrix[i] != h_result1[i]) {std::cout << " mismatch at " << i << " was: " << h_matrix[i] << " should be: " << h_result1[i] << std::endl; return 0;}
}
$ nvcc -o t1714 t1714.cu -std=c++11
$ ./t1714
serial time: 0.018015 seconds
parallel time: 0.005337 seconds
$

我不声明此代码或我发布的任何其他代码的正确性。使用我发布的任何代码的任何人都需要自担风险。我只是声称我试图解决原始帖子中的问题，并提供一些解释。我并不是说我的代码没有缺陷，或者它适用于任何特定目的。使用（或不使用）风险自负。