CUDA – cuBLAS でベクトル演算 SAXPY

CUDA には、線形代数演算を行うための cuBLAS ライブラリが付属しています。
これを利用してベルトル演算を行う例を掲載します。

目次

cuBLAS で saxpy 演算

saxpy とは “Single precision A X Plus Y” の略です。
単精度における、ax + y の演算のことを表します。
ここで、a はスカラー型、x と y はベクトルを表します。(行列を表す場合もあります)

cuBLAS には、cublasSaxpy という関数が用意されており、この演算を行うことができます。

以下、配列長5のベクトル x, y を生成して、2 * x + y を計算するサンプルです。

#include <iostream>
#include <iterator>
#include <numeric>
#include <string>
#include <vector>
#include <cublas_v2.h>

void print_vector(const std::string &label, const std::vector<float> &x)
{
    std::cout << label << ": [ ";
    std::copy(x.begin(), x.end(), std::ostream_iterator<float>(std::cout, " "));
    std::cout << "]" << std::endl;
}

int main()
{
    size_t n = 5;

    // x の初期化 [ 0 1 2 3 4 ]
    std::vector<float> x(n);
    std::iota(x.begin(), x.end(), 0);
    print_vector("x", x);

    // y の初期化 [ 1 1 1 1 1 ]
    std::vector<float> y(n);
    std::fill(y.begin(), y.end(), 1);
    print_vector("y", y);

    // x を GPU メモリに転送
    float *dev_x;
    cudaMalloc((void **)&dev_x, n * sizeof(*dev_x));
    cublasSetVector(n, sizeof(float), x.data(), 1, dev_x, 1);

    // y を GPU メモリに転送
    float *dev_y;
    cudaMalloc((void **)&dev_y, n * sizeof(*dev_y));
    cublasSetVector(n, sizeof(float), y.data(), 1, dev_y, 1);

    // y = 2 * x + y を計算
    float a = 2.0f;
    cublasHandle_t handle;
    cublasCreate(&handle);
    cublasSaxpy(handle, n, &a, dev_x, 1, dev_y, 1);
    cublasDestroy(handle);

    // 計算結果 y を GPU メモリから取得して表示
    // [ 1 3 5 7 9 ]
    cublasGetVector(n, sizeof(float), dev_y, 1, y.data(), 1);
    print_vector("2x+y", y);

    cudaFree(dev_y);
    cudaFree(dev_x);

    return 0;
}

コンパイルするときは cublas のライブラリをリンクする必要があります。

$ nvcc -lcublas main.cu && ./a.out
x: [ 0 1 2 3 4 ]
y: [ 1 1 1 1 1 ]
2x+y: [ 1 3 5 7 9 ]

cuBLAS で saxpy 演算(thrust 利用版)

thrust ライブラリを利用することで、前後の処理をより簡潔に記述することができます。

以下は thrust を利用して同じ計算を行うサンプルです。

#include <cublas_v2.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/sequence.h>

void print_device_vector(const std::string &label, const thrust::device_vector<float> &x)
{
    std::cout << label << ": [ ";
    thrust::copy(x.begin(), x.end(), std::ostream_iterator<float>(std::cout, " "));
    std::cout << "]" << std::endl;
}

int main()
{
    size_t n = 5;

    // x の初期化 [ 0 1 2 3 4 ]
    thrust::device_vector<float> x(n);
    thrust::sequence(x.begin(), x.end());
    print_device_vector("x", x);

    // y の初期化 [ 1 1 1 1 1 ]
    thrust::device_vector<float> y(n);
    thrust::fill(y.begin(), y.end(), 1.0f);
    print_device_vector("y", y);

    // y = 2 * x + y を計算
    float a = 2.0f;
    cublasHandle_t handle;
    cublasCreate(&handle);
    cublasSaxpy(
        handle, n, &a,
        thrust::raw_pointer_cast(x.data()), 1,
        thrust::raw_pointer_cast(y.data()), 1);
    cublasDestroy(handle);

    // 計算結果 y を表示 [ 1 3 5 7 9 ]
    print_device_vector("2x+y", y);

    return 0;
}
$ nvcc -lcublas main_thrust.cu && ./a.out
x: [ 0 1 2 3 4 ]
y: [ 1 1 1 1 1 ]
2x+y: [ 1 3 5 7 9 ]
  • URLをコピーしました!
  • URLをコピーしました!

コメント

コメントする

目次