Cublas - 列/行明智的操作

Posted 2021-04-21

我正在寻找一种在列上执行操作的方法。我有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调用都会太短。因此，运行大量单个内核会有很多延迟，这会对性能产生负面影响。

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