kernel_cl.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_CL_HPP
#define STAN_MATH_OPENCL_KERNEL_CL_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/buffer_types.hpp>
#include <stan/math/opencl/matrix_cl_view.hpp>
#include <stan/math/opencl/opencl_context.hpp>
#include <stan/math/opencl/matrix_cl.hpp>
#include <stan/math/opencl/stringify.hpp>
#include <stan/math/opencl/err/check_opencl.hpp>
#include <stan/math/opencl/kernels/helpers.hpp>
#include <stan/math/prim/fun/vec_concat.hpp>
#include <CL/opencl.hpp>
#include <algorithm>
#include <map>
#include <string>
#include <vector>
#include <utility>
 
namespace stan {
namespace math {
namespace opencl_kernels {
namespace internal {
 
template <typename T, require_not_matrix_cl_t<T>* = nullptr>
inline const T& get_kernel_args(const T& t) {
  return t;
}
 
template <typename K, require_matrix_cl_t<K>* = nullptr>
inline const cl::Buffer& get_kernel_args(const K& m) {
  return m.buffer();
}
 
template <typename T, require_not_matrix_cl_t<T>* = nullptr>
inline void assign_event(const cl::Event& e, const T&) {}
 
template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
          require_same_t<T, in_buffer>* = nullptr>
inline void assign_event(const cl::Event& e, const K& m) {
  m.add_read_event(e);
}
template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
          require_same_t<T, out_buffer>* = nullptr>
inline void assign_event(const cl::Event& e, K& m) {
  m.add_write_event(e);
}
template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
          require_same_t<T, in_out_buffer>* = nullptr>
inline void assign_event(const cl::Event& e, K& m) {
  m.add_read_write_event(e);
}
 
template <typename T, require_same_t<T, cl::Event>* = nullptr>
inline void assign_events(const T&) {}
 
template <typename Arg, typename... Args, typename CallArg,
          typename... CallArgs>
inline void assign_events(const cl::Event& new_event, CallArg& m,
                          CallArgs&... args) {
  assign_event<Arg>(new_event, m);
  assign_events<Args...>(new_event, args...);
}
 
template <typename T, require_not_matrix_cl_t<T>* = nullptr>
inline tbb::concurrent_vector<cl::Event> select_events(const T& m) {
  return tbb::concurrent_vector<cl::Event>{};
}
template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
          require_same_t<T, in_buffer>* = nullptr>
inline const tbb::concurrent_vector<cl::Event>& select_events(const K& m) {
  return m.write_events();
}
template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
          require_any_same_t<T, out_buffer, in_out_buffer>* = nullptr>
inline tbb::concurrent_vector<cl::Event> select_events(K& m) {
  static_assert(!std::is_const<K>::value, "Can not write to const matrix_cl!");
  return m.read_write_events();
}
 
}  // namespace internal
 
inline auto compile_kernel(
    const char* name, const std::vector<std::string>& sources,
    const std::unordered_map<std::string, int>& options) {
  auto base_opts = opencl_context.base_opts();
  for (auto& it : options) {
    if (base_opts[it.first] > it.second) {
      base_opts[it.first] = it.second;
    }
  }
  std::string kernel_opts = "";
  for (auto&& comp_opts : base_opts) {
    kernel_opts += std::string(" -D") + comp_opts.first + "="
                   + std::to_string(comp_opts.second);
  }
  cl::Program program(opencl_context.context(), sources);
  try {
    program.build({opencl_context.device()}, kernel_opts.c_str());
 
    return cl::Kernel(program, name);
  } catch (const cl::Error& e) {
    // in case of CL_BUILD_PROGRAM_FAILURE, print the build error
    if (e.err() == -11) {
      std::string buildlog = program.getBuildInfo<CL_PROGRAM_BUILD_LOG>(
          opencl_context.device()[0]);
      system_error("compile_kernel", name, e.err(), buildlog.c_str());
    } else {
      check_opencl_error(name, e);
    }
  }
  return cl::Kernel();  // never reached because check_opencl_error throws
}
 
template <typename... Args>
struct kernel_cl {
 private:
  const char* name_;
  std::vector<std::string> sources_;
  std::unordered_map<std::string, int> opts_;
  mutable cl::Kernel kernel_;
 
 public:
  kernel_cl(const char* name, std::vector<std::string> sources,
            std::unordered_map<std::string, int> options = {})
      : name_(name), sources_(std::move(sources)), opts_(std::move(options)) {}
 
  template <typename... CallArgs>
  auto operator()(cl::NDRange global_thread_size, CallArgs&&... args) const {
    if (kernel_() == NULL) {
      kernel_ = compile_kernel(name_, sources_, opts_);
      opencl_context.register_kernel_cache(&kernel_);
    }
    cl::EnqueueArgs eargs(opencl_context.queue(),
                          vec_concat(std::vector<cl::Event>{},
                                     internal::select_events<Args>(args)...),
                          global_thread_size);
    cl::KernelFunctor<internal::to_const_buffer_t<Args>&...> kernel_functor(
        kernel_);
    cl::Event kern_event
        = kernel_functor(eargs, internal::get_kernel_args(args)...);
    internal::assign_events<Args...>(kern_event, args...);
    return kern_event;
  }
 
  template <typename... CallArgs>
  auto operator()(cl::NDRange global_thread_size, cl::NDRange thread_block_size,
                  CallArgs&&... args) const {
    if (kernel_() == NULL) {
      kernel_ = compile_kernel(name_, sources_, opts_);
      opencl_context.register_kernel_cache(&kernel_);
    }
    cl::EnqueueArgs eargs(opencl_context.queue(),
                          vec_concat(std::vector<cl::Event>{},
                                     internal::select_events<Args>(args)...),
                          global_thread_size, thread_block_size);
    cl::KernelFunctor<internal::to_const_buffer_t<Args>&...> kernel_functor(
        kernel_);
    cl::Event kern_event
        = kernel_functor(eargs, internal::get_kernel_args(args)...);
    internal::assign_events<Args...>(kern_event, args...);
    return kern_event;
  }
 
  int get_option(const std::string option_name) const {
    return std::min(opts_.at(option_name),
                    opencl_context.base_opts().at(option_name));
  }
};
 
}  // namespace opencl_kernels
}  // namespace math
}  // namespace stan
 
#endif
#endif