dnn_lib.hpp

Go to the documentation of this file.

#ifndef LBANN_UTILS_DNN_LIB_DNN_LIB_HPP
#define LBANN_UTILS_DNN_LIB_DNN_LIB_HPP

#include "lbann/base.hpp"
#include "lbann/layers/data_type_layer.hpp"
#include "lbann/layers/layer.hpp"
#include "lbann/utils/exception.hpp"
#include "lbann/utils/gpu/helpers.hpp"
#include <vector>

#include "lbann/proto/layers.pb.h"

#ifdef LBANN_HAS_DNN_LIB

namespace lbann {
namespace dnn_lib {

#if defined LBANN_HAS_CUDNN
using namespace cudnn;
#elif defined LBANN_HAS_MIOPEN
using namespace miopen;
#endif // LBANN_HAS_CUDNN

template <typename T>
struct ScalingParameterT
{
  using type = T;
};

template <typename T>
using ScalingParamType = typename ScalingParameterT<T>::type;

#ifdef LBANN_HAS_GPU_FP16
template <>
struct ScalingParameterT<fp16>
{
  using type = float;
};
#endif // LBANN_USE_GPU_FP16

// Global DNN library objects

void initialize();
void destroy();
dnnHandle_t& get_handle();

// Helper functions for DNN library types

template <typename TensorDataType>
dnnDataType_t get_data_type();

// Wrapper classes for DNN library types

template <typename T>
using BackendHandleType = typename T::handle_type;

class TensorDescriptor
{
public:
  using handle_type = dnnTensorDescriptor_t;

public:
  explicit TensorDescriptor(dnnTensorDescriptor_t desc = nullptr);

  ~TensorDescriptor();

  // Copy-and-swap idiom
  TensorDescriptor(const TensorDescriptor&);
  TensorDescriptor(TensorDescriptor&&);
  TensorDescriptor& operator=(TensorDescriptor);
  friend void swap(TensorDescriptor& first, TensorDescriptor& second);

  void reset(dnnTensorDescriptor_t desc = nullptr);
  dnnTensorDescriptor_t release() noexcept;
  dnnTensorDescriptor_t get() const noexcept;
  operator dnnTensorDescriptor_t() const noexcept;

  void create();
  void set(dnnDataType_t data_type,
           std::vector<int> dims,
           std::vector<int> strides = {});
  template <typename... IntTs>
  void set(dnnDataType_t data_type, IntTs... dims)
  {
    set(data_type, {static_cast<int>(dims)...});
  }
#if !(defined LBANN_HAS_CUDNN)
  void set(dnnDataType_t data_type,
           dnnTensorFormat_t /*format*/,
           const std::vector<int>& dims)
  {
    this->set(data_type, dims);
  }
#endif // !LBANN_HAS_CUDNN

private:
  dnnTensorDescriptor_t desc_ = nullptr;
};

#ifdef LBANN_HAS_CUDNN

class FilterDescriptor
{
public:
  using handle_type = dnnFilterDescriptor_t;

public:
  explicit FilterDescriptor(dnnFilterDescriptor_t desc = nullptr);

  ~FilterDescriptor();

  // Copy-and-swap idiom
  FilterDescriptor(const FilterDescriptor&);
  FilterDescriptor(FilterDescriptor&&);
  FilterDescriptor& operator=(FilterDescriptor);
  friend void swap(FilterDescriptor& first, FilterDescriptor& second);

  void reset(dnnFilterDescriptor_t desc = nullptr);
  dnnFilterDescriptor_t release() noexcept;
  dnnFilterDescriptor_t get() const noexcept;
  operator dnnFilterDescriptor_t() const noexcept;

  void create();
  void set(dnnDataType_t data_type,
           dnnTensorFormat_t format,
           const std::vector<int>& dims);
  template <typename... IntTs>
  void set(dnnDataType_t data_type, dnnTensorFormat_t format, IntTs... dims)
  {
    set(data_type, format, {static_cast<int>(dims)...});
  }

private:
  dnnFilterDescriptor_t desc_ = nullptr;
};
#else  // MIOpen and OneDNN
using FilterDescriptor = TensorDescriptor;
#endif // LBANN_HAS_CUDNN

class DropoutDescriptor
{

public:
  explicit DropoutDescriptor(dnnDropoutDescriptor_t desc = nullptr);
  DropoutDescriptor(float dropout,
                    void* states,
                    size_t states_size,
                    unsigned long long seed,
                    bool use_mask,
                    bool state_evo,
                    dnnRNGType_t rng_mode)
  {
    this
      ->set(dropout, states, states_size, seed, use_mask, state_evo, rng_mode);
  }

  ~DropoutDescriptor();

  // Copy-and-swap idiom
  DropoutDescriptor(const DropoutDescriptor&);
  DropoutDescriptor(DropoutDescriptor&&);
  DropoutDescriptor& operator=(DropoutDescriptor);
  friend void swap(DropoutDescriptor& first, DropoutDescriptor& second);

  void reset(dnnDropoutDescriptor_t desc = nullptr);
  dnnDropoutDescriptor_t release() noexcept;
  dnnDropoutDescriptor_t get() const noexcept;
  operator dnnDropoutDescriptor_t() const noexcept;

  void create();
  void set(float dropout,
           void* states,
           size_t states_size,
           unsigned long long seed,
           bool use_mask = false,
           bool state_evo = false,
           dnnRNGType_t rng_mode = DNN_RNG_PSEUDO_XORWOW);

private:
  dnnDropoutDescriptor_t desc_ = nullptr;
};

class RNNDescriptor
{

public:
  explicit RNNDescriptor(dnnRNNDescriptor_t desc = nullptr);

  RNNDescriptor(const RNNDescriptor&) = delete;
  ~RNNDescriptor();

  // Copy-and-swap idiom
  RNNDescriptor(RNNDescriptor&&);
  RNNDescriptor& operator=(RNNDescriptor);
  friend void swap(RNNDescriptor& first, RNNDescriptor& second);

  void reset(dnnRNNDescriptor_t desc = nullptr);
  dnnRNNDescriptor_t release() noexcept;
  dnnRNNDescriptor_t get() const noexcept;
  operator dnnRNNDescriptor_t() const noexcept;

  void create();
  void set(dnnRNNAlgo_t algorithm,
           dnnRNNMode_t cell_mode,
           dnnRNNBiasMode_t bias_mode,
           dnnDirectionMode_t direction_mode,
           dnnRNNInputMode_t input_mode,
           dnnDataType_t data_type,
           dnnDataType_t math_precision,
           dnnMathType_t math_type,
           size_t input_size,
           size_t hidden_size,
           size_t proj_size,
           size_t num_layers,
           dnnDropoutDescriptor_t dropout_desc,
           uint32_t aux_flags);

private:
  dnnRNNDescriptor_t desc_ = nullptr;
};

class ConvolutionDescriptor
{
public:
  using DescriptorHandle_t = dnnConvolutionDescriptor_t;

public:

  explicit ConvolutionDescriptor(DescriptorHandle_t desc = nullptr);

  ~ConvolutionDescriptor();

  ConvolutionDescriptor(const ConvolutionDescriptor&);
  ConvolutionDescriptor(ConvolutionDescriptor&&);

  ConvolutionDescriptor& operator=(ConvolutionDescriptor);



  DescriptorHandle_t release() noexcept;
  DescriptorHandle_t get() const noexcept;
  operator DescriptorHandle_t() const noexcept;



  void swap(ConvolutionDescriptor& other);

  void reset(DescriptorHandle_t desc = nullptr);

  void create();

  void set(std::vector<int> const& pad,
           std::vector<int> const& stride,
           std::vector<int> const& dilation,
           dnnDataType_t data_type,
           dnnConvolutionMode_t mode = DNN_CROSS_CORRELATION);
  void set(size_t array_dim,
           int const pad[],
           int const stride[],
           int const dilation[],
           dnnDataType_t data_type,
           dnnConvolutionMode_t mode = DNN_CROSS_CORRELATION);

  void set_math_mode(dnnMathType_t math_type);

  void set_group_count(int num_groups);


private:
  DescriptorHandle_t desc_ = nullptr;
};

void swap(ConvolutionDescriptor& lhs, ConvolutionDescriptor& rhs);

class PoolingDescriptor
{
public:
  using DescriptorHandle_t = dnnPoolingDescriptor_t;

public:

  explicit PoolingDescriptor(DescriptorHandle_t desc = nullptr);

  ~PoolingDescriptor();

  PoolingDescriptor(const PoolingDescriptor&);
  PoolingDescriptor(PoolingDescriptor&&);

  PoolingDescriptor& operator=(PoolingDescriptor);



  DescriptorHandle_t release() noexcept;
  DescriptorHandle_t get() const noexcept;
  operator DescriptorHandle_t() const noexcept;



  void swap(PoolingDescriptor& other);

  void reset(DescriptorHandle_t desc = nullptr);

  void create();
  void set(pooling_mode mode,
           dnnNanPropagation_t maxpoolingNanOpt,
           std::vector<int> const& window_dims,
           std::vector<int> const& padding,
           std::vector<int> const& stride);
  void set(pooling_mode mode,
           dnnNanPropagation_t nan_prop,
           int num_dims,
           int const window_dims[],
           int const padding[],
           int const stride[]);


private:
  DescriptorHandle_t desc_ = nullptr;
};

void swap(PoolingDescriptor& lhs, PoolingDescriptor& rhs);

class LRNDescriptor
{
public:
  using DescriptorHandle_t = dnnLRNDescriptor_t;

public:

  explicit LRNDescriptor(DescriptorHandle_t desc = nullptr);

  ~LRNDescriptor();

  LRNDescriptor(const LRNDescriptor&);
  LRNDescriptor(LRNDescriptor&&);

  LRNDescriptor& operator=(LRNDescriptor);



  DescriptorHandle_t release() noexcept;
  DescriptorHandle_t get() const noexcept;
  operator DescriptorHandle_t() const noexcept;



  void swap(LRNDescriptor& other);

  void reset(DescriptorHandle_t desc = nullptr);

  void create();
  void set(unsigned n,
           double alpha,
           double beta,
           double k,
           dnnLRNMode_t mode = DNN_LRN_CROSS_CHANNEL);


private:
  DescriptorHandle_t desc_ = nullptr;
};

void swap(LRNDescriptor& lhs, LRNDescriptor& rhs);

// DNN library tensor managers

template <typename TensorDataType>
class layer_tensor_manager
{
public:
  using LayerType = data_type_layer<TensorDataType>;

public:
  layer_tensor_manager(const LayerType* l = nullptr);
  virtual ~layer_tensor_manager() = default;

  const LayerType* get_layer() const { return m_layer; }
  void set_layer(const LayerType* l);

  virtual TensorDescriptor& get_prev_activations(int parent_index = 0) = 0;
  virtual TensorDescriptor& get_activations(int child_index = 0) = 0;
  virtual TensorDescriptor& get_prev_error_signals(int child_index = 0) = 0;
  virtual TensorDescriptor& get_error_signals(int parent_index = 0) = 0;

protected:
  layer_tensor_manager(const layer_tensor_manager&) = default;
  layer_tensor_manager& operator=(const layer_tensor_manager&) = default;
  layer_tensor_manager(layer_tensor_manager&&) = default;
  layer_tensor_manager& operator=(layer_tensor_manager&&) = default;

  void set_num_parents(int num_parents);
  void set_num_children(int num_children);

  const LayerType* m_layer;
  std::vector<TensorDescriptor> m_prev_activations;
  std::vector<TensorDescriptor> m_activations;
  std::vector<TensorDescriptor> m_prev_error_signals;
  std::vector<TensorDescriptor> m_error_signals;
};

template <typename TensorDataType>
class data_parallel_layer_tensor_manager
  : public layer_tensor_manager<TensorDataType>
{
public:
  using LayerType = data_type_layer<TensorDataType>;

public:
  data_parallel_layer_tensor_manager(const LayerType* l = nullptr);
  data_parallel_layer_tensor_manager(
    const data_parallel_layer_tensor_manager&) = default;
  data_parallel_layer_tensor_manager&
  operator=(const data_parallel_layer_tensor_manager&) = default;
  data_parallel_layer_tensor_manager(data_parallel_layer_tensor_manager&&) =
    default;
  data_parallel_layer_tensor_manager&
  operator=(data_parallel_layer_tensor_manager&&) = default;
  ~data_parallel_layer_tensor_manager() = default;
  TensorDescriptor& get_prev_activations(int parent_index = 0) override;
  TensorDescriptor& get_activations(int child_index = 0) override;
  TensorDescriptor& get_prev_error_signals(int child_index = 0) override;
  TensorDescriptor& get_error_signals(int parent_index = 0) override;
};

template <typename TensorDataType>
class entrywise_layer_tensor_manager
  : public layer_tensor_manager<TensorDataType>
{
public:
  using LayerType = data_type_layer<TensorDataType>;

public:
  entrywise_layer_tensor_manager(const LayerType* l = nullptr);
  entrywise_layer_tensor_manager(const entrywise_layer_tensor_manager&) =
    default;
  entrywise_layer_tensor_manager&
  operator=(const entrywise_layer_tensor_manager&) = default;
  entrywise_layer_tensor_manager(entrywise_layer_tensor_manager&&) = default;
  entrywise_layer_tensor_manager&
  operator=(entrywise_layer_tensor_manager&&) = default;
  ~entrywise_layer_tensor_manager() = default;
  TensorDescriptor& get_prev_activations(int parent_index = 0) override;
  TensorDescriptor& get_activations(int child_index = 0) override;
  TensorDescriptor& get_prev_error_signals(int child_index = 0) override;
  TensorDescriptor& get_error_signals(int parent_index = 0) override;
};

// DNN library algorithm selection

fwd_conv_alg get_fwd_algorithm(bool autotune,
                               bool deterministic,
                               const TensorDescriptor& input_desc,
                               const void* input,
                               const FilterDescriptor& kernel_desc,
                               const void* kernel,
                               const ConvolutionDescriptor& conv_desc,
                               const TensorDescriptor& output_desc,
                               void* output,
                               size_t ws_size,
                               void* ws);

bwd_data_conv_alg
get_bwd_data_algorithm(bool autotune,
                       bool deterministic,
                       const FilterDescriptor& kernel_desc,
                       const void* kernel,
                       const TensorDescriptor& prev_error_signal_desc,
                       const void* prev_error_signal,
                       const ConvolutionDescriptor& conv_desc,
                       const TensorDescriptor& error_signal_desc,
                       void* error_signal,
                       size_t ws_size,
                       void* ws);

bwd_filter_conv_alg
get_bwd_filter_algorithm(bool autotune,
                         bool deterministic,
                         const TensorDescriptor& input_desc,
                         const void* input,
                         const TensorDescriptor& prev_error_signal_desc,
                         const void* prev_error_signal,
                         const ConvolutionDescriptor& conv_desc,
                         const FilterDescriptor& kernel_gradient_desc,
                         void* kernel_gradient,
                         size_t ws_size,
                         void* ws);

void default_to_tensor_ops() noexcept;

dnnMathType_t get_default_convolution_math_type() noexcept;

using ProtoTensorOpEnumType = decltype(lbann_data::DEFAULT_TENSOR_OPS);
dnnMathType_t convert_to_dnn_math_type(ProtoTensorOpEnumType mt);
ProtoTensorOpEnumType convert_to_proto_math_type(dnnMathType_t mt);

} // namespace dnn_lib
} // namespace lbann
#endif // LBANN_HAS_DNN_LIB
#endif // LBANN_UTILS_DNN_LIB_DNN_LIB_HPP