Cublas - 列/行明智的操作



我正在寻找一种对列执行操作的方法。 我有 MxN 矩阵,我想在每列上激活 cublas 函数(例如 nrm2(。

我期望得到的结果是:M x 1

我该怎么做?

CUBLAS 没有批处理的 1 级例程,因此没有直接的方法可以在单个调用中计算列或行规范。您可以通过在矩阵的所有行或列上循环多次调用 nrm2 来实现,例如:

#include <cublas_v2.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/transform.h>
#include <thrust/random.h>
#include <thrust/device_vector.h>
#include <iostream>
struct prg
{
float a, b;
__host__ __device__
prg(float _a=0.f, float _b=1.f) : a(_a), b(_b) {};
__host__ __device__
float operator()(const unsigned int n) const
{
thrust::default_random_engine rng;
thrust::uniform_real_distribution<float> dist(a, b);
rng.discard(n);
return dist(rng);
}
};

int main(void)
{
const int M = 1024, N = M;
const int num = N * M;
thrust::device_vector<float> matrix(num);
thrust::device_vector<float> vector(N, -1.0f);
thrust::counting_iterator<unsigned int> index_sequence_begin(0);
thrust::transform(index_sequence_begin,
index_sequence_begin + num,
matrix.begin(),
prg(1.f,2.f));
float* m_d = thrust::raw_pointer_cast(matrix.data());
float* v_d = thrust::raw_pointer_cast(vector.data());
cudaStream_t stream; 
cudaStreamCreate(&stream);
cublasHandle_t handle;
cublasCreate(&handle);
cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_DEVICE);
cublasSetStream(handle, stream);
for(int col=0; col < N; col++) {
cublasSnrm2(handle, M, m_d + col*M, 1, v_d + col);
}
cudaDeviceSynchronize();
for(auto x : vector) {
float normval = x;
std::cout << normval << std::endl;
}
return 0;
}

除非您有非常大的行或列,否则利用流来运行并发内核并减少整体运行时的空间很小,因为每个 nrm2 调用都太短了。因此,运行大量单个内核时会有很多延迟,这会对性能产生负面影响。

一个更好的选择是编写自己的内核来执行此操作。

最新更新