matrix_cl.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_MATRIX_CL_HPP
#define STAN_MATH_OPENCL_MATRIX_CL_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/prim/size.hpp>
#include <stan/math/opencl/err/check_opencl.hpp>
#include <stan/math/prim/err/check_size_match.hpp>
#include <stan/math/opencl/opencl_context.hpp>
#include <stan/math/opencl/ref_type_for_opencl.hpp>
#include <stan/math/opencl/matrix_cl_view.hpp>
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/prim/fun/vec_concat.hpp>
#include <CL/opencl.hpp>
#include <tbb/concurrent_vector.h>
#include <algorithm>
#include <iostream>
#include <string>
#include <type_traits>
#include <vector>
 
namespace stan {
namespace math {
 
// forward declare
template <typename T>
class arena_matrix_cl;
 
template <typename>
class matrix_cl;
 
template <typename T>
class matrix_cl : public matrix_cl_base {
 private:
  cl::Buffer buffer_cl_;  // Holds the allocated memory on the device
  int rows_{0};           // Number of rows.
  int cols_{0};           // Number of columns.
  // Holds info on if matrix is a special type
  matrix_cl_view view_{matrix_cl_view::Entire};
  mutable tbb::concurrent_vector<cl::Event> write_events_;  // Tracks write jobs
  mutable tbb::concurrent_vector<cl::Event> read_events_;   // Tracks reads
 
 public:
  using Scalar = T;  // Underlying type of the matrix
  using type = T;    // Underlying type of the matrix
  // Forward declare the methods that work in place on the matrix
  template <matrix_cl_view matrix_view = matrix_cl_view::Entire>
  inline void zeros_strict_tri();
 
  int rows() const { return rows_; }
 
  int cols() const { return cols_; }
 
  int size() const { return rows_ * cols_; }
 
  const matrix_cl_view& view() const { return view_; }
 
  void view(const matrix_cl_view& view) { view_ = view; }
 
  inline void clear_write_events() const {
    write_events_.clear();
    return;
  }
 
  inline void clear_read_events() const {
    read_events_.clear();
    return;
  }
 
  inline void clear_read_write_events() const {
    read_events_.clear();
    write_events_.clear();
    return;
  }
 
  inline const tbb::concurrent_vector<cl::Event>& write_events() const {
    return write_events_;
  }
 
  inline const tbb::concurrent_vector<cl::Event>& read_events() const {
    return read_events_;
  }
 
  inline const tbb::concurrent_vector<cl::Event> read_write_events() const {
    return vec_concat(this->read_events(), this->write_events());
  }
 
  inline void add_read_event(cl::Event new_event) const {
    this->read_events_.push_back(new_event);
  }
 
  inline void add_write_event(cl::Event new_event) const {
    this->write_events_.push_back(new_event);
  }
 
  inline void add_read_write_event(cl::Event new_event) const {
    this->read_events_.push_back(new_event);
    this->write_events_.push_back(new_event);
  }
 
  inline void wait_for_write_events() const {
    for (cl::Event e : write_events_) {
      e.wait();
    }
    write_events_.clear();
  }
 
  inline void wait_for_read_events() const {
    for (cl::Event e : read_events_) {
      e.wait();
    }
    read_events_.clear();
  }
 
  inline void wait_for_read_write_events() const {
    wait_for_read_events();
    wait_for_write_events();
  }
 
  const cl::Buffer& buffer() const { return buffer_cl_; }
  cl::Buffer& buffer() { return buffer_cl_; }
 
  matrix_cl() {}
  matrix_cl(const cl::Buffer& A, const int R, const int C,
            matrix_cl_view partial_view = matrix_cl_view::Entire)
      : buffer_cl_(A), rows_(R), cols_(C), view_(partial_view) {}
 
  matrix_cl(const matrix_cl<T>& A)
      : rows_(A.rows()), cols_(A.cols()), view_(A.view()) {
    if (A.size() == 0) {
      return;
    }
    buffer_cl_ = cl::Buffer(opencl_context.context(), CL_MEM_READ_WRITE,
                            sizeof(T) * this->size());
    initialize_buffer_cl(A);
  }
 
  matrix_cl(matrix_cl<T>&& A)
      : buffer_cl_(std::move(A.buffer_cl_)),
        rows_(A.rows_),
        cols_(A.cols_),
        view_(A.view_),
        write_events_(std::move(A.write_events_)),
        read_events_(std::move(A.read_events_)) {}
 
  // defined in rev/arena_matrix_cl.hpp
  matrix_cl(const arena_matrix_cl<T>& A);  // NOLINT(runtime/explicit)
 
  template <typename Vec, require_std_vector_vt<is_eigen, Vec>* = nullptr,
            require_st_same<Vec, T>* = nullptr>
  explicit matrix_cl(Vec&& A) try : rows_(A.empty() ? 0 : A[0].size()),
                                    cols_(A.size()) {
    if (this->size() == 0) {
      return;
    }
    cl::Context& ctx = opencl_context.context();
    cl::CommandQueue& queue = opencl_context.queue();
    buffer_cl_ = cl::Buffer(ctx, CL_MEM_READ_WRITE, sizeof(T) * size());
    for (int i = 0, offset_size = 0; i < cols_; i++, offset_size += rows_) {
      check_size_match("matrix constructor", "input rows", A[i].size(),
                       "matrix_cl rows", rows_);
      cl::Event write_event;
      queue.enqueueWriteBuffer(
          buffer_cl_,
          opencl_context.in_order() || std::is_rvalue_reference<Vec&&>::value,
          sizeof(T) * offset_size, sizeof(T) * rows_, A[i].data(), nullptr,
          &write_event);
      this->add_write_event(write_event);
    }
  } catch (const cl::Error& e) {
    check_opencl_error("matrix constructor", e);
  }
 
  matrix_cl(const int rows, const int cols,
            matrix_cl_view partial_view = matrix_cl_view::Entire)
      : rows_(rows), cols_(cols), view_(partial_view) {
    if (this->size() == 0) {
      return;
    }
    cl::Context& ctx = opencl_context.context();
    try {
      int flags = CL_MEM_READ_WRITE;
      if (opencl_context.device()[0].getInfo<CL_DEVICE_HOST_UNIFIED_MEMORY>()) {
        flags |= CL_MEM_ALLOC_HOST_PTR;
      }
      buffer_cl_ = cl::Buffer(ctx, flags, sizeof(T) * rows_ * cols_);
    } catch (const cl::Error& e) {
      check_opencl_error("matrix constructor", e);
    }
  }
 
  template <typename Mat, require_eigen_t<Mat>* = nullptr,
            require_vt_same<Mat, T>* = nullptr>
  explicit matrix_cl(Mat&& A,
                     matrix_cl_view partial_view = matrix_cl_view::Entire)
      : rows_(A.rows()), cols_(A.cols()), view_(partial_view) {
    using Mat_type = std::decay_t<ref_type_for_opencl_t<Mat>>;
    if (this->size() == 0) {
      return;
    }
    initialize_buffer_no_heap_if<
        std::is_same<std::decay_t<Mat>, Mat_type>::value
        && (std::is_lvalue_reference<Mat>::value
            || is_eigen_contiguous_map<Mat>::value)>(A);
  }
 
  template <typename Scal,
            typename = require_same_t<T, std::remove_reference_t<Scal>>>
  explicit matrix_cl(Scal&& A,
                     matrix_cl_view partial_view = matrix_cl_view::Diagonal)
      : rows_(1), cols_(1), view_(partial_view) {
    initialize_buffer<std::is_rvalue_reference<Scal&&>::value>(
        const_cast<const std::decay_t<Scal>*>(&A));
  }
 
  template <typename Vec, require_std_vector_t<Vec>* = nullptr,
            require_vt_same<Vec, T>* = nullptr>
  explicit matrix_cl(Vec&& A,
                     matrix_cl_view partial_view = matrix_cl_view::Entire)
      : matrix_cl(std::forward<Vec>(A), A.size(), 1) {}
 
  template <typename Vec, require_std_vector_t<Vec>* = nullptr,
            require_vt_same<Vec, T>* = nullptr>
  explicit matrix_cl(Vec&& A, const int& R, const int& C,
                     matrix_cl_view partial_view = matrix_cl_view::Entire)
      : rows_(R), cols_(C), view_(partial_view) {
    initialize_buffer_no_heap_if<std::is_lvalue_reference<Vec>::value>(A);
  }
 
  template <typename U, require_same_t<T, U>* = nullptr>
  explicit matrix_cl(const U* A, const int& R, const int& C,
                     matrix_cl_view partial_view = matrix_cl_view::Entire)
      : rows_(R), cols_(C), view_(partial_view) {
    initialize_buffer(A);
  }
 
  // defined in kernel_generator/matrix_cl_conversion.hpp
  template <typename Expr,
            require_all_kernel_expressions_and_none_scalar_t<Expr>* = nullptr,
            require_not_matrix_cl_t<Expr>* = nullptr>
  matrix_cl(const Expr& expression);  // NOLINT(runtime/explicit)
 
  matrix_cl<T>& operator=(matrix_cl<T>&& a) {
    view_ = a.view();
    rows_ = a.rows();
    cols_ = a.cols();
    this->wait_for_read_write_events();
    buffer_cl_ = std::move(a.buffer_cl_);
    write_events_ = std::move(a.write_events_);
    read_events_ = std::move(a.read_events_);
    return *this;
  }
 
  matrix_cl<T>& operator=(const matrix_cl<T>& a) {
    this->view_ = a.view();
    if (a.size() == 0) {
      this->rows_ = a.rows();
      this->cols_ = a.cols();
      return *this;
    }
    this->wait_for_read_write_events();
    if (this->size() != a.size()) {
      buffer_cl_ = cl::Buffer(opencl_context.context(), CL_MEM_READ_WRITE,
                              sizeof(T) * a.size());
    }
    this->rows_ = a.rows();
    this->cols_ = a.cols();
    initialize_buffer_cl(a);
    return *this;
  }
 
  // defined in kernel_generator/matrix_cl_conversion.hpp
  template <typename Expr,
            require_all_kernel_expressions_and_none_scalar_t<Expr>* = nullptr,
            require_not_matrix_cl_t<Expr>* = nullptr>
  matrix_cl<T>& operator=(const Expr& expression);
 
  // defined in rev/arena_matrix_cl.hpp
  matrix_cl<T>& operator=(const arena_matrix_cl<T>& other);
 
  const matrix_cl<T>& eval() const& { return *this; }
  matrix_cl<T> eval() && { return std::move(*this); }
 
  ~matrix_cl() { wait_for_read_write_events(); }
 
  void setZero() {
    if (this->size() == 0) {
      return;
    }
    cl::Event zero_event;
    const std::size_t write_events_size = this->write_events().size();
    const std::size_t read_events_size = this->read_events().size();
    const std::size_t read_write_size = write_events_size + read_events_size;
    std::vector<cl::Event> read_write_events(read_write_size, cl::Event{});
    auto&& read_events_vec = this->read_events();
    auto&& write_events_vec = this->write_events();
    for (std::size_t i = 0; i < read_events_size; ++i) {
      read_write_events[i] = read_events_vec[i];
    }
    for (std::size_t i = read_events_size, j = 0; j < write_events_size;
         ++i, ++j) {
      read_write_events[i] = write_events_vec[j];
    }
    try {
      opencl_context.queue().enqueueFillBuffer(buffer_cl_, static_cast<T>(0), 0,
                                               sizeof(T) * this->size(),
                                               &read_write_events, &zero_event);
    } catch (const cl::Error& e) {
      check_opencl_error("setZero", e);
    }
    this->add_write_event(zero_event);
  }
 
 private:
  template <bool in_order = false>
  cl::Event initialize_buffer(const T* A) {
    cl::Event transfer_event;
    if (this->size() == 0) {
      return transfer_event;
    }
    cl::Context& ctx = opencl_context.context();
    cl::CommandQueue& queue = opencl_context.queue();
    try {
      buffer_cl_ = cl::Buffer(ctx, CL_MEM_READ_WRITE, sizeof(T) * size());
      queue.enqueueWriteBuffer(buffer_cl_,
                               opencl_context.in_order() || in_order, 0,
                               sizeof(T) * size(), A, nullptr, &transfer_event);
      this->add_write_event(transfer_event);
    } catch (const cl::Error& e) {
      check_opencl_error("initialize_buffer", e);
    }
    return transfer_event;
  }
 
  template <bool in_order = false>
  cl::Event initialize_buffer(T* A) {
    cl::Event transfer_event;
    if (this->size() == 0) {
      return transfer_event;
    }
    cl::Context& ctx = opencl_context.context();
    cl::CommandQueue& queue = opencl_context.queue();
    try {
      if (opencl_context.device()[0].getInfo<CL_DEVICE_HOST_UNIFIED_MEMORY>()) {
        constexpr auto copy_or_share
            = CL_MEM_COPY_HOST_PTR * INTEGRATED_OPENCL
              | (CL_MEM_USE_HOST_PTR * !INTEGRATED_OPENCL);
        buffer_cl_
            = cl::Buffer(ctx, CL_MEM_READ_WRITE | copy_or_share,
                         sizeof(T) * size(), A);  // this is always synchronous
      } else {
        buffer_cl_ = cl::Buffer(ctx, CL_MEM_READ_WRITE, sizeof(T) * size());
        queue.enqueueWriteBuffer(
            buffer_cl_, opencl_context.in_order() || in_order, 0,
            sizeof(T) * size(), A, nullptr, &transfer_event);
        this->add_write_event(transfer_event);
      }
    } catch (const cl::Error& e) {
      check_opencl_error("initialize_buffer", e);
    }
    return transfer_event;
  }
 
  template <bool No_heap, typename U, std::enable_if_t<No_heap>* = nullptr>
  void initialize_buffer_no_heap_if(U&& obj) {
    if (this->size() == 0) {
      return;
    }
    initialize_buffer(obj.data());
  }
  // we need separate overloads as obj.data() might not be available when second
  // overload is called.
  template <bool No_heap, typename U, std::enable_if_t<!No_heap>* = nullptr>
  void initialize_buffer_no_heap_if(U&& obj) {
    using U_val = std::decay_t<ref_type_for_opencl_t<U>>;
    if (this->size() == 0) {
      return;
    }
    auto* obj_heap = new U_val(std::move(obj));
    try {
      cl::Event e = initialize_buffer(obj_heap->data());
      if (opencl_context.device()[0].getInfo<CL_DEVICE_HOST_UNIFIED_MEMORY>()) {
        buffer_cl_.setDestructorCallback(&delete_it_destructor<U_val>,
                                         obj_heap);
      } else {
        e.setCallback(CL_COMPLETE, &delete_it_event<U_val>, obj_heap);
      }
    } catch (...) {
      delete obj_heap;
      throw;
    }
  }
 
  void initialize_buffer_cl(const matrix_cl<T>& A) {
    cl::Event cstr_event;
    std::vector<cl::Event>* dep_events = new std::vector<cl::Event>(
        A.write_events().begin(), A.write_events().end());
    try {
      opencl_context.queue().enqueueCopyBuffer(A.buffer(), this->buffer(), 0, 0,
                                               A.size() * sizeof(T), dep_events,
                                               &cstr_event);
      if (opencl_context.device()[0].getInfo<CL_DEVICE_HOST_UNIFIED_MEMORY>()) {
        buffer_cl_.setDestructorCallback(
            &delete_it_destructor<std::vector<cl::Event>>, dep_events);
      } else {
        cstr_event.setCallback(
            CL_COMPLETE, &delete_it_event<std::vector<cl::Event>>, dep_events);
      }
      this->add_write_event(cstr_event);
      A.add_read_event(cstr_event);
    } catch (const cl::Error& e) {
      delete dep_events;
      check_opencl_error("copy (OpenCL)->(OpenCL)", e);
    } catch (...) {
      delete dep_events;
      throw;
    }
  }
 
  template <typename U>
  static void delete_it_event(cl_event e, cl_int status, void* container) {
    delete static_cast<U*>(container);
  }
 
  template <typename U>
  static void delete_it_destructor(cl_mem buff, void* container) {
    delete static_cast<U*>(container);
  }
};
}  // namespace math
}  // namespace stan
 
#endif
#endif