Пишем на Rust + CUDA C / Habr

Всем привет!

В данном руководстве хочу рассказать как подружить CUDA C/С++ и Rust. И в качестве примера напишем небольшую программу на Rust для вычисления скалярного произведения векторов, вычисление скалярного произведения будет производиться на GPU с использованием CUDA C.

Кому интересно под кат!

CUDA C

Первым делом необходимо поставить компилятор CUDA — nvcc. Что такое CUDA и для чего это нужно описывать не буду, об этом можно почитать например тут. Скажу лишь что с её помощью можно писать код, который будет запускаться на NVIDIA видеокартах (в дальнейшем — GPU) и использовать всю их мощь для параллельных вычислений и обработки графики. Еще раз повторюсь, данный туториал не о том как писать код на CUDA, а о том как и из кода на Rust пользоваться её преимуществами и писать параллельные вычисления на GPU.

Итак устанавливаем nvcc и CUDA Toolkit. С этим сложностей не должно возникнуть подробная инструкция: на офф сайте.

RUST + CUDA C

В ходе данного туториала, как уже говорилось ранее, будем писать программу на Rust для нахождения скалярного произведения двух векторов, сам процесс вычисления будет происходить на GPU.

Скалярное произведение двух векторов.

Пусть у нас есть два вектора: $inline$ и $inline$ , скалярное произведение этих векторов:

$a \cdot b = \sum_{i=1}^{n}{a_ib_i}$

Начнем создание нашей программы. Далее я предполагаю, что nvcc успешно установлен, так же стоят rustc и cargo для компиляции rust кода.

Первым делом создадим папку проекта. В папке проекта создадим файл Cargo.toml, в котором находятся инструкции для сборщика cargo. Файл выглядит таким образом:

[package]
name = "rust_cuda" # название программы
version = "0.1.0" # версия программы
authors = ["MoonL1ght <ixav1@icloud.com>"] # информация об авторе
build = "build.rs" # скрипт для сборки rust
links = "cudart" # библиотека cuda, которая линкуется динамически

[dependencies]
libc = "0.2" # библиотека rust для работы С кодом 
rand = "0.5.5" # библиотека rust для работы с случайными величинами

[build-dependencies]
cc = "1.0" # rust пакет для сборки С кода

Так же в корневой папке проекта создаем файл build.rs в котором будут находится инструкции для сборки программы на rust и компиляции CUDA C кода.

В корень проекта добавим папку src в которую поместим файлы с исходным кодом. В папке src создадим четыре файла: main.rs — код основной программы, dot.cpp — С++ binding (обертка для CUDA C), dot_gpu.h, dot_gpu.cu — файл в котором содержится код выполняемый на GPU.

Итого имеем такую структура проекта:

rust-cuda/
    src/
        main.rs
        dot.cpp
        dot_gpu.h
        dot_gpu.cu
    Cargo.toml
    build.rs

В файле build.rs самое главное прописать это:

println!("cargo:rustc-link-search=native=/Developer/NVIDIA/CUDA-10.1/lib");
println!("cargo:rustc-link-search=/Developer/NVIDIA/CUDA-10.1/lib");
println!("cargo:rustc-env=LD_LIBRARY_PATH=/Developer/NVIDIA/CUDA-10.1/lib");
println!("cargo:rustc-link-lib=dylib=cudart");

где /Developer/NVIDIA/CUDA-10.1/lib — путь к исполняемым файлам CUDA, в unix подобной системе этот путь можно узнать например командой:

which nvcc

Помимо этого в файле build.rs нужно указать путь к файлам dot.cpp и dot_gpu.cpp:

.files(&["./src/dot.cpp", "./src/dot_gpu.cu"])

Весь код файла build.rs

extern crate cc;

fn main() {
  cc::Build::new()
    .cuda(true)
    .cpp(true)
    .flag("-cudart=shared")
    .files(&["./src/dot.cpp", "./src/dot_gpu.cu"])
    .compile("dot.a");
  println!("cargo:rustc-link-search=native=/Developer/NVIDIA/CUDA-10.1/lib");
  println!("cargo:rustc-link-search=/Developer/NVIDIA/CUDA-10.1/lib");
  println!("cargo:rustc-env=LD_LIBRARY_PATH=/Developer/NVIDIA/CUDA-10.1/lib");
  println!("cargo:rustc-link-lib=dylib=cudart");
}

Теперь можно приступать к написанию основного кода программы. В файле main.rs нужно создать интерфейс С/С++ функции для вызова непосредственно из кода на Rust. Более подробно об это можно почитать в официальной документации в разделе FFI.

extern "C" {
  // интерфейс C функции для расчета скалярного произведения двух векторов
  fn dot(v1: *mut c_float, v2: *mut c_float, N: size_t) -> c_float;
}

Для её вызова надо использовать unsafe блок кода, в качестве аргументов передаем mutable pointer на тип Vec:

unsafe {
  gpu_res = dot(v1.as_mut_ptr(), v2.as_mut_ptr(), VEC_SIZE);
}

Полный код файла main.rs

extern crate libc;
extern crate rand;

use libc::{c_float, size_t};
use rand::Rng;

const VEC_SIZE: usize = 10;
const MAX: f32 = 10.;
const MIN: f32 = 0.;


extern "C" {
  fn dot(v1: *mut c_float, v2: *mut c_float, N: size_t) -> c_float;
}

fn cpu_dot(v1: Vec<f32>, v2: Vec<f32>) -> f32 {
  let mut res: f32 = 0.;
  for i in 0..v1.len() {
    res += v1[i] * v2[i];
  }
  return res;
}

fn main() {
  let mut v1: Vec<f32> = Vec::new();
  let mut v2: Vec<f32> = Vec::new();
  let mut gpu_res: c_float;
  let mut cpu_res: f32 = 0.;

  let mut rng = rand::thread_rng();
  for _ in 0..VEC_SIZE {
    v1.push(rng.gen_range(MIN, MAX));
    v2.push(rng.gen_range(MIN, MAX));
  }

  println!("{:?}", v1);
  println!("{:?}", v2);

  println!("GPU computing started");
  unsafe {
    gpu_res = dot(v1.as_mut_ptr(), v2.as_mut_ptr(), VEC_SIZE);
  }
  println!("GPU computing finished");
  println!("GPU dot product result: {}", gpu_res);
  
  cpu_res = cpu_dot(v1, v2);
  println!("CPU dot product result: {}", cpu_res);
}

Теперь приступаем к написанию обвязки на C++, а так же кода для вычисления скалярного произведения векторов на CUDA C.

В файле dot.cpp напишем функцию обвязку, собственно эту функцию мы и вызываем из Rust кода:

extern "C" {
  float dot(float *v1, float *v2, size_t N) {
    float *gpu_res;
    float res = 0.0;
    gpu_res = gpu_dot(v1, v2, N); // вычисление на GPU
    for (int i = 0; i < blocksPerGrid; i++) {
      res += gpu_res[i];
    }
    free(gpu_res);
    return res;
  }
}

Полный код файла dot.cpp

#include <iostream>
#include "dot_gpu.h"

using namespace std;

void display_vector(float *v, size_t N) {
  cout << "[";
  for (size_t i = 0; i < N; i++) {
    cout << v[i];
    if (i != N - 1) {
      cout << ", ";
    }
  }
  cout << "]" << endl;
}

extern "C" {
  float dot(float *v1, float *v2, size_t N) {
    cout << "Calling gpu dot product" << endl;
    cout << "Got two vectors from rust:" << endl;
    display_vector(v1, N);
    display_vector(v2, N);
    float *gpu_res;
    float res = 0.0;
    gpu_res = gpu_dot(v1, v2, N);
    for (int i = 0; i < blocksPerGrid; i++) {
      res += gpu_res[i];
    }
    free(gpu_res);
    return res;
  }
}

Далее представлен код из файла dot_gpu.cu в котором производится основное вычисление, объяснять сам код в данном туториале не буду, так как он не посвящен программированию на CUDA.

dot_gpu.cu

#include "dot_gpu.h"

__global__ void dot__(float *v1, float *v2, float *res, int N) {
  __shared__ float cache [threadsPerBlock];
  int tid = threadIdx.x + blockIdx.x * blockDim.x;
  int cacheIndex = threadIdx.x;
  float temp = 0.0;
  while (tid < N) {
    temp += v1[tid] * v2[tid];
    tid += blockDim.x * gridDim.x;
  }
  cache[cacheIndex] = temp;

  __syncthreads();

  int i = blockDim.x / 2;
  while (i != 0) {
    if (cacheIndex < i) {
      cache[cacheIndex] += cache[cacheIndex + i];
    }
    __syncthreads();
    i /= 2;   
  }

  if (cacheIndex == 0) {
    res[blockIdx.x] = cache[0];
  }
}

float * gpu_dot (float *v1, float *v2, size_t N) {
	float *dev_v1, *dev_v2, *dev_res, *res;

	res = new float[blocksPerGrid];

	cudaMalloc((void**)&dev_v1, N * sizeof(float));
  cudaMalloc((void**)&dev_v2, N * sizeof(float));
	cudaMalloc((void**)&dev_res, blocksPerGrid * sizeof(float));
	
	cudaMemcpy(dev_v1, v1, N * sizeof(float), cudaMemcpyHostToDevice);
	cudaMemcpy(dev_v2, v2, N * sizeof(float), cudaMemcpyHostToDevice);
	
	dot__<<<blocksPerGrid, threadsPerBlock>>>(dev_v1, dev_v2, dev_res, (int)N);
	cudaMemcpy(res, dev_res, blocksPerGrid * sizeof(float), cudaMemcpyDeviceToHost);

	cudaFree(dev_v1);
  cudaFree(dev_v2);
	cudaFree(dev_res);
	
	return res;
}

Все наша маленькая программа написана и готова к сборке. Для того что бы её собрать в консоли вызовем команду:

cargo build

Для запуска:

cargo run

После сборки программы в основной директории проекта появится папка target. Исполняемый файл нашей программы будет находится в папке: ./target/debug/

При этом если мы просто запустим наш исполняемый файл то получим ошибку: dyld library not loaded. То есть он не может найти путь к динамической библиотеки cuda. Для решения данной проблемы можно перед запуском исполняемого файла в консоли прописать переменную окружения LD_LIBRARY_PATH=path_to_CUDA_lib_directory/ или же создать символьные линки в папке rust toolchain для CUDA:

ln -s /Developer/NVIDIA/CUDA-10.1/lib/* /Users/Alexander/.rustup/toolchains/nightly-x86_64-apple-darwin/lib

где /Users/Alexander/.rustup/toolchains/nightly-x86_64-apple-darwin/lib — мой путь к установленному rust toolchain-у, у вас он может немного отличаться.

При запуске программы через cargo run такой ошибки не возникало, потому что мы прописали переменную окружения LD_LIBRARY_PATH в файле build.rs.

В итоге

Мы имеем возможность запускать код CUDA C прямиком из кода на Rust. Для того, что бы проверить это, мы создали небольшую программу, она работает с векторами и все вычисления производит на GPU. Полный код так же можно посмотреть на github.