CUDA入门系列之矩阵表示

一、代码

#ifndef CUDA_MAT_MAT_HPP
#define CUDA_MAT_MAT_HPP

#include <iostream>
#include <exception>
#include <sstream>

#include <cuda.h>
#include <cuda_runtime.h>

using namespace std;

#define CHECK(exp) \
if (exp != cudaSuccess) { \
    auto err = cudaGetLastError(); \
    printf("[ERROR] [%s:%d] %s(%d).\n", __FILE__, __LINE__, cudaGetErrorString(err), err); \
    exit(-1); \
}

template<class T>
class Mat {
  template<class K>
  friend ostream &operator<<(ostream &stream, const Mat<K> &mat);

public:
  Mat() : row_(0), col_(0), size_(0), cpu_data_(nullptr), gpu_data_(nullptr) {};

  Mat(size_t r, size_t c, bool use_cuda = false) : row_(r), col_(c), size_(r * c * sizeof(T)) {
    if (!use_cuda) {
      CHECK(cudaMallocHost(&cpu_data_, size_));
      gpu_data_ = nullptr;
    } else {
      cpu_data_ = nullptr;
      CHECK(cudaMalloc(&gpu_data_, size_));
    }
  }

  Mat(const initializer_list<T> &init_data) : Mat(1, init_data.size()) {
    T *q = cpu_data_;
    for (auto p = init_data.begin(); p != init_data.end(); p++, q++) {
      *q = *p;
    }
  }

  Mat(const initializer_list<initializer_list<T>> &init_data) : Mat(init_data.size(), init_data.begin()->size()) {
    T *q = cpu_data_;
    for (auto p1 = init_data.begin(); p1 != init_data.end(); p1++) {
      if (p1->size() != col_) {
        stringstream ss;
        ss << "The expected col is " << col_ << ", but the actual value is " << p1->size();
        throw invalid_argument(ss.str().c_str());
      }

      for (auto p2 = p1->begin(); p2 != p1->end(); p2++, q++) {
        *q = *p2;
      }
    }
  }

  Mat(const Mat &rhs) noexcept: row_(rhs.row_), col_(rhs.col_), size_(rhs.size_) {
    if (rhs.cpu_data_) {
      CHECK(cudaMallocHost(&cpu_data_, size_));
      CHECK(cudaMemcpy(cpu_data_, rhs.cpu_data_, size_, cudaMemcpyHostToHost));
    } else cpu_data_ = nullptr;

    if (rhs.gpu_data_) {
      CHECK(cudaMalloc(&gpu_data_, size_));
      CHECK(cudaMemcpy(gpu_data_, rhs.gpu_data_, size_, cudaMemcpyDeviceToDevice));
    } else gpu_data_ = nullptr;
  }

  Mat(Mat &&rhs) noexcept: row_(rhs.row_), col_(rhs.col_), size_(rhs.size_), cpu_data_(rhs.cpu_data_),
                           gpu_data_(rhs.gpu_data_) {
    rhs.col_ = rhs.row_ = rhs.row_ = 0;
    rhs.gpu_data_ = nullptr;
    rhs.cpu_data_ = nullptr;
  }

  Mat &operator=(const Mat &rhs) {
    Mat tmp(rhs);
    swap(tmp);
    return *this;
  }

  ~Mat() {
    if (cpu_data_)cudaFreeHost(cpu_data_);
    if (gpu_data_)cudaFree(gpu_data_);
  }

  Mat &cuda() {
    if (cpu_data_) {
      if (!gpu_data_)CHECK(cudaMalloc(&gpu_data_, size_));
      CHECK(cudaMemcpy(gpu_data_, cpu_data_, size_, cudaMemcpyHostToHost));
    }
    return *this;
  }

  Mat &cpu() {
    if (gpu_data_) {
      if (!cpu_data_)CHECK(cudaMallocHost(&cpu_data_, size_));
      CHECK(cudaMemcpy(cpu_data_, gpu_data_, size_, cudaMemcpyDeviceToHost));
    }
    return *this;
  }

  void swap(Mat<T> &rhs) {
    std::swap(row_, rhs.row_);
    std::swap(col_, rhs.col_);
    std::swap(size_, rhs.size_);
    std::swap(cpu_data_, rhs.cpu_data_);
    std::swap(gpu_data_, rhs.gpu_data_);
  }

  inline const T *cpu_data() const { return cpu_data_; }

  inline T *cpu_data_mutable() { return cpu_data_; }

  inline const T *gpu_data() const { return gpu_data_; };

  inline T *gpu_data_mutable() { return gpu_data_; }

  inline size_t row() const { return row_; }

  inline size_t col() const { return col_; }

private:
  T *cpu_data_, *gpu_data_;
  size_t row_, col_, size_;
};

template<class T>
ostream &operator<<(ostream &stream, const Mat<T> &mat) {
  if (mat.cpu_data_ == nullptr) {
    return stream;
  }

  T *p = mat.cpu_data_;
  for (size_t r = 0; r < mat.row_; r++) {
    for (size_t c = 0; c < mat.col_; c++) {
      stream << *p++ << " ";
    }
    stream << "\n";
  }

  return stream;
}
#endif //CUDA_MAT_MAT_HPP

二、main.cpp

#include <iostream>

#include "mat.hpp"

using namespace std;

int main() {
  Mat<float> a(3, 4);
  Mat<float> b = {1, 2, 3, 4};
  Mat<float> c = {{1, 2, 3}, {4, 5, 6}};

  cout << c << endl;
  return 0;
}

三、CMakeLists.txt

cmake_minimum_required(VERSION 3.17)
project(cuda_mat)

set(CMAKE_CXX_STANDARD 14)

find_package(CUDA)
find_package(CUDAToolkit)

cuda_add_executable(cuda_mat main.cpp mat.hpp)
target_link_libraries(cuda_mat CUDA::cudart)

赏

非常感谢,我会继续努力的！

支付宝

微信

• 本文作者： killf
• 本文链接： http://www.killf.info/编程语言/CUDA/CUDA入门系列之矩阵表示/
• 版权声明： 本博客所有文章除特别声明外，均采用 Apache License 2.0 许可协议。转载请注明出处！