layer.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.

#ifndef LBANN_LAYERS_LAYER_HPP_INCLUDED
#define LBANN_LAYERS_LAYER_HPP_INCLUDED

#include "lbann/base.hpp"
#include "lbann/utils/exception.hpp"
#include "lbann/utils/typename.hpp"
#include <string>
#include <vector>
#ifdef LBANN_HAS_ONNX
#include <onnx/onnx_pb.h>
#endif

#define LBANN_DEFINE_LAYER_BUILDER(LAYER_NAME)                                 \
  template <typename TensorDataType, data_layout Layout, El::Device Device>    \
  std::unique_ptr<Layer> build_##LAYER_NAME##_layer_from_pbuf(                 \
    lbann_comm*,                                                               \
    lbann_data::Layer const&)

#define LBANN_LAYER_DEFAULT_BUILDER(LAYER_NAME)                                \
  template <typename TensorDataType, data_layout Layout, El::Device Device>    \
  std::unique_ptr<Layer> build_##LAYER_NAME##_layer_from_pbuf(                 \
    lbann_comm* comm,                                                          \
    lbann_data::Layer const&)                                                  \
  {                                                                            \
    using LayerType = LAYER_NAME##_layer<TensorDataType, Layout, Device>;      \
    return std::make_unique<LayerType>(comm);                                  \
  }

#define LBANN_LAYER_BUILDER_ETI(LAYER_NAME, T, Device)                         \
  template std::unique_ptr<Layer>                                              \
    build_##LAYER_NAME##_layer_from_pbuf<T,                                    \
                                         ::lbann::data_layout::DATA_PARALLEL,  \
                                         Device>(lbann_comm*,                  \
                                                 lbann_data::Layer const&);    \
  template std::unique_ptr<Layer>                                              \
    build_##LAYER_NAME##_layer_from_pbuf<T,                                    \
                                         ::lbann::data_layout::MODEL_PARALLEL, \
                                         Device>(lbann_comm*,                  \
                                                 lbann_data::Layer const&)

// Forward-declare protobuf classes
namespace lbann_data {
class Layer;
}

namespace lbann {

// Forward declarations
class lbann_comm;
class description;
class Layer;
class model;
class lbann_summary;
class weights;
using ViewingWeightsPtr = std::weak_ptr<weights>;
#ifdef LBANN_HAS_DISTCONV
class distconv_adapter;
#endif // LBANN_HAS_DISTCONV
namespace callback {
class sync_layers;
} // namespace callback
class KFAC;
template <hydrogen::Device Device>
class kfac_block_fc_conv;
template <hydrogen::Device Device>
class kfac_block_channelwise_fc;
template <hydrogen::Device Device>
class kfac_block_bn;
template <hydrogen::Device Device>
class kfac_block_gru;

using OwningLayerPtr = std::shared_ptr<Layer>;
using ViewingLayerPtr = std::weak_ptr<Layer>;

struct ParallelStrategy
{
  int sample_groups = 0;
  int sample_splits = 0;
  int depth_groups = 0;
  int depth_splits = 0;
  int height_groups = 0;
  int height_splits = 0;
  int width_groups = 0;
  int width_splits = 0;
  int channel_groups = 0;
  int channel_splits = 0;
  int filter_groups = 0;
  int filter_splits = 0;
  int replications = 0;
  bool enable_subgraph = false;
  int sub_branch_tag = 0;
  int sub_branch_resource_percentage = 0;
  bool operator==(const ParallelStrategy& ps) const
  {
    return sample_groups == ps.sample_groups &&
           sample_splits == ps.sample_splits &&
           depth_groups == ps.depth_groups && depth_splits == ps.depth_splits &&
           height_groups == ps.height_groups &&
           height_splits == ps.height_splits &&
           width_groups == ps.width_groups && width_splits == ps.width_splits &&
           channel_groups == ps.channel_groups &&
           channel_splits == ps.channel_splits &&
           filter_groups == ps.filter_groups &&
           filter_splits == ps.filter_splits &&
           replications == ps.replications &&
           sub_branch_tag == ps.sub_branch_tag &&
           sub_branch_resource_percentage ==
             ps.sub_branch_resource_percentage &&
           enable_subgraph == ps.enable_subgraph;
  }
  bool operator!=(const ParallelStrategy& ps) const { return !(*this == ps); }
};

inline std::ostream& operator<<(std::ostream& os, const ParallelStrategy& ps)
{
  os << "{"
     << "N: " << ps.sample_groups << "/" << ps.sample_splits << ", "
     << "C: " << ps.channel_groups << "/" << ps.channel_splits << ", "
     << "D: " << ps.depth_groups << "/" << ps.depth_splits << ", "
     << "H: " << ps.height_groups << "/" << ps.height_splits << ", "
     << "W: " << ps.width_groups << "/" << ps.width_splits << ", "
     << "F: " << ps.filter_groups << "/" << ps.filter_splits << ", "
     << "R: " << ps.replications << ", "
     << "T: " << ps.sub_branch_tag << ", "
     << "%: " << ps.sub_branch_resource_percentage << ", "
     << "e: " << ps.enable_subgraph << "}";
  return os;
}

inline std::ostream& print_parallel_strategy_header(std::ostream& os)
{
  os << "Axis over which DistConv can parallelize:\n"
     << "\tSamples in the mini-batch (N)\n"
     << "\tDepth, Height, and Width (D x H x W)\n"
     << "\tChannel (C)\n"
     << "\tFilters (F)\n"
     << "\tReplications (R): Number of times the layer is replicated (for FC "
        "layers right now)\n"
     << "\tBranch number in the subgraph (T)\n"
     << "\tPercentage of parent resources to be allocated to this branch (%)\n"
     << "\tEnable subgraph for the layer (e)\n"
     << "\nFor each of the above dimensions there are two fields:\n"
     << "\t# Groups (G): refers to how many reduced-order tensors exist with "
        "respect to that dimension"
     << std::endl
     << "\t             e.g. For a kD tensor you would have a stack of G "
        "(k-1)D tensors"
     << std::endl
     << "\t\t[N, C, D, H, W]" << std::endl
     << "\t\t[2, 1, 4, 1, 1] ---" << std::endl
     << "\t\t                  |" << std::endl
     << "\t\t                  V" << std::endl
     << "\t\t  4 Depth groups: [N, C, H, W]" << std::endl
     << "\t\t                  [2, 1, 1, 1]" << std::endl
     << "\t\t                  [2, 1, 1, 1]" << std::endl
     << "\t\t                  [2, 1, 1, 1]" << std::endl
     << "\t\t                  [2, 1, 1, 1]" << std::endl
     << "\t\t[1, 1, 4, 1, 2] ---" << std::endl
     << "\t\t                  |" << std::endl
     << "\t\t                  V" << std::endl
     << "\t\t 2 Sample groups: [C, D, H, W]" << std::endl
     << "\t\t                  [1, 4, 1, 1]" << std::endl
     << "\t\t                  [1, 4, 1, 1]" << std::endl
     << "\n\tSplit per Dimension (S): Number of groups the dimension is split "
        "over (i.e. split K times) (aka H2 split shape) (must divide groups "
        "evenly).\n"
     << std::endl;

  os << "Reporting order for the parallel strategy" << std::endl;
  os << "{N: G/S"
     << ", C: G/S"
     << ", D: G/S"
     << ", H: G/S"
     << ", W: G/S"
     << ", F: G/S"
     << ", R:"
     << ", T:"
     << ", %:"
     << ", e:"
     << "}";
  return os;
}

enum SubGraphCommunication
{
  PT2PT = 0,
  COLL = 10,
  COLL_OPT = 2
};

class Layer
{
  friend class callback::sync_layers;
  friend class KFAC;
  template <hydrogen::Device Device>
  friend class kfac_block_fc_conv;
  template <hydrogen::Device Device>
  friend class kfac_block_channelwise_fc;
  template <hydrogen::Device Device>
  friend class kfac_block_bn;
  template <hydrogen::Device Device>
  friend class kfac_block_gru;

public:
  Layer();
  virtual ~Layer() = default;

  virtual Layer* copy() const = 0;



  void set_name(const std::string name) { m_name = name; }
  void set_model(model* m) { m_model = m; }



  std::string get_name() const { return m_name; }

  model* get_model() const noexcept { return m_model; }

  virtual std::string get_type() const = 0;

  virtual std::string get_datatype_name() const = 0;

  virtual description get_description() const;

  virtual data_layout get_data_layout() const = 0;
  virtual El::Device get_device_allocation() const = 0;

  int get_expected_num_parent_layers() const noexcept
  {
    return m_expected_num_parent_layers;
  }

  int get_expected_num_child_layers() const noexcept
  {
    return m_expected_num_child_layers;
  }

  virtual int get_backprop_requirements() const
  {
    return ERROR_SIGNALS | PREV_ACTIVATIONS | ACTIVATIONS | WEIGHTS;
  }

  ParallelStrategy& get_parallel_strategy() noexcept
  {
    return m_parallel_strategy;
  }
  ParallelStrategy const& get_parallel_strategy() const noexcept
  {
    return m_parallel_strategy;
  }


  virtual bool can_run_inplace() const { return false; }

#ifdef LBANN_HAS_GPU
  bool using_gpus() const { return get_device_allocation() == El::Device::GPU; }
#else
  bool using_gpus() const noexcept { return false; }
#endif // LBANN_HAS_GPU



  virtual void forward_prop() = 0;
  void back_prop();

  bool update();

  virtual void setup(size_t max_mini_batch_size,
                     const std::vector<El::Grid*>& grids);

  virtual void check_setup();


  void write_proto(lbann_data::Layer& proto) const;


  // FIXME (trb 10/03/2023): The lbann_summary class should be
  // reevaluated. This strikes me as a feature that should be moved to
  // a callback and/or replaced by proper performance
  // profilers/counters (e.g., Caliper). More directly: I'm not sure
  // anyone has used the summarizer in at least 2 years and it might
  // be time to trim that out.
  void summarize_stats(lbann_summary& summarizer, int step);
  virtual void summarize_matrices(lbann_summary& summarizer, int step) = 0;

  void reset_counters();



  // (trb 10/03/2023): unused, but accessor is used; kept for symmetry
  void set_communication_flag(SubGraphCommunication type)
  {
    subgraph_communication_method = type;
  }

  // (trb 10/03/2023): used, keeping setter, which is unused.
  SubGraphCommunication get_communication_flag()
  {
    return subgraph_communication_method;
  }

  // (trb 10/03/2023): used
  void set_num_spliting_groups(El::Int spliting_groups)
  {
    m_num_spliting_groups = spliting_groups;
  }

  // (trb 10/03/2023): used
  El::Int get_num_spliting_groups() const { return m_num_spliting_groups; }

  // (trb 10/03/2023): USED BUT THIS IS BAD BECAUSE m_mygrid IS NEVER SET!
  std::shared_ptr<El::Grid> get_mygrid() const
  {
    LBANN_ERROR("This function should not be used.");
    return nullptr;
  }

  // (trb 10/03/2023): used, model.cpp
  void reset_inter_subgrid_vc_comm(std::shared_ptr<El::mpi::Comm> mpi_comm)
  {
    m_interSubGridVCComm = std::move(mpi_comm);
  }

  // (trb 10/03/2023): used
  void set_subgraph_parallelism_execution()
  {
    m_subgraph_parallelism_execution = true;
  }

  // layer-level sub-graph parallelism execution
  // (trb 10/03/2023): used
  bool subgraph_parallelism_execution() const noexcept
  {
    return m_subgraph_parallelism_execution;
  }

  // (trb 10/03/2023): used
  void set_run_layer_in_subgraph() { run_layer_in_subgraph = true; }

  // (trb 10/03/2023): used
  bool get_run_layer_in_subgraph() const noexcept
  {
    return run_layer_in_subgraph;
  }

private:
  virtual void write_specific_proto(lbann_data::Layer& proto) const = 0;

public:
#ifdef LBANN_HAS_ONNX

  virtual void fill_onnx_node(onnx::GraphProto& graph) const;

private:
  virtual std::string get_onnx_op_type() const;
#endif // LBANN_HAS_ONNX

public:

  const Layer& get_parent_layer(size_t index = 0) const;
  const Layer& get_child_layer(size_t index = 0) const;

  std::vector<const Layer*> get_parent_layers() const;
  std::vector<const Layer*> get_child_layers() const;

  size_t find_parent_layer_index(const Layer& l) const;
  size_t find_child_layer_index(const Layer& l) const;

  int get_num_parents() const noexcept { return m_parent_layers.size(); }
  int get_num_children() const noexcept { return m_child_layers.size(); }



  void add_parent_layer(ViewingLayerPtr parent);
  void add_child_layer(ViewingLayerPtr child);

  void replace_parent_layer(ViewingLayerPtr l, size_t index);
  void replace_child_layer(ViewingLayerPtr l, size_t index);

  void clear_parent_layers() { m_parent_layers.clear(); }
  void clear_child_layers() { m_child_layers.clear(); }

  ViewingLayerPtr get_parent_layer_pointer(size_t index) const;
  ViewingLayerPtr get_child_layer_pointer(size_t index) const;

  virtual std::vector<ViewingLayerPtr> get_layer_pointers();
  virtual void set_layer_pointers(std::vector<ViewingLayerPtr> layers);



  std::vector<ViewingWeightsPtr> get_weights_pointers() const;
  void set_weights_pointers(std::vector<ViewingWeightsPtr> ptrs);

  void replace_weights(Layer const& other_layer);



  virtual const BaseDistMat& get_activations(const Layer& child) const = 0;
  virtual const BaseDistMat& get_error_signals(const Layer& parent) const = 0;



  std::vector<int> get_input_dims(size_t input_index = 0) const;
  int get_input_size(size_t input_index = 0) const;
  std::vector<int> get_output_dims(size_t output_index = 0) const;
  int get_output_size(size_t output_index = 0) const;

  void set_output_dims(std::vector<int> dims, size_t output_index = 0);

  El::Int infer_mini_batch_size_from_parents() const;
  virtual El::Int current_output_mini_batch_size() const = 0;
  virtual El::Int
  infer_mini_batch_size_from_parents_or_default_to_current() const = 0;


  lbann_comm* get_comm() const;

  int get_grid_tag() const noexcept;
  void set_grid_tag(int tag);


  void set_hint_layer(ViewingLayerPtr l);

  const Layer* get_hint_layer() const;



  void freeze();
  void unfreeze();
  bool is_frozen() const;


  virtual void set_keep_error_signals(bool) = 0;

  bool runs_inplace() const { return m_runs_inplace; }


  template <typename ArchiveT>
  void serialize(ArchiveT& ar);


protected:
  Layer(Layer&& other) = default;
  Layer(Layer const& other);
  Layer& operator=(Layer&& other) = default;
  Layer& operator=(Layer const& other);

  void add_weights(ViewingWeightsPtr w)
  {
    m_weights.emplace_back(std::move(w));
  }
  size_t num_weights() const noexcept { return m_weights.size(); }
  bool has_weights() const noexcept { return num_weights() > 0; }
  bool has_weights(size_t idx) const noexcept
  {
    return ((idx < m_weights.size()) && (!m_weights[idx].expired()));
  }
  void set_num_weights(size_t n) { m_weights.resize(n); }
  void set_weights(size_t idx, ViewingWeightsPtr w)
  {
    m_weights.at(idx) = std::move(w);
  }
  weights const& get_weights(size_t idx) const;

  weights& get_weights(size_t idx);

  void add_as_gradient_source();

  void remove_as_gradient_source();

  // ===========================================================
  // Setup helper functions
  // ===========================================================

  void setup_grid();
  virtual void setup_pointers();
  virtual void setup_dims();
  virtual void setup_matrices(const std::vector<El::Grid*>& grids) = 0;
  virtual void setup_data(size_t max_mini_batch_size){};
  virtual void setup_gpu() {}

  // ===========================================================
  // Forward prop step helper functions
  // ===========================================================

  virtual void fp_setup_inputs() = 0;
  virtual void fp_setup_outputs() = 0;
  virtual void fp_compute() = 0;

  // ===========================================================
  // Back prop step helper functions
  // ===========================================================

  virtual void bp_setup_gradient_wrt_inputs() = 0;
  virtual void bp_compute(){};

  // ===========================================================
  // Update step helper functions
  // ===========================================================

  virtual bool update_compute() { return true; }

  // ===========================================================
  // Protected class members
  // ===========================================================

  int m_expected_num_parent_layers = 1;
  int m_expected_num_child_layers = 1;

  model* m_model = nullptr;

  bool m_frozen;

  EvalType m_fp_time;
  EvalType m_fp_compute_time;
  EvalType m_bp_time;
  EvalType m_bp_compute_time;
  EvalType m_update_time;

  std::string m_name;

  bool m_runs_inplace = false;

  // -------------------------------------------------------
  // Objects for sub-grid parallelism
  // -------------------------------------------------------

  int m_grid_tag = -1;

  // -------------------------------------------------------
  // Objects from old sub-grid parallelism implementation
  // -------------------------------------------------------

  SubGraphCommunication subgraph_communication_method = PT2PT;

  // Model-level: Is subgraph parallelism enabled for this Model?
  // Layer-level: Does this layer need subgraph execution (like split and sum
  // layers) Process-level: Does this layer exist in the given process (For e.g.
  // some layers will only run on a subset of processes defined by grid tag)
  // Layer-level sub-graph execution
  bool m_subgraph_parallelism_execution = false;
  // Process-level sub-graph execution
  bool run_layer_in_subgraph = false;

  std::unique_ptr<std::set<int>> m_subgrid_ranks;

  El::Int m_num_spliting_groups = 1;
  std::shared_ptr<El::mpi::Comm> m_interSubGridVCComm;

private:

  friend void attempt_move_error_signal(Layer& parent,
                                        Layer const& child,
                                        std::unique_ptr<BaseDistMat> signal);
  friend void attempt_view_error_signal(Layer& parent,
                                        Layer const& child,
                                        const BaseDistMat& signals);
  friend void deep_copy_error_signal(Layer& parent,
                                     Layer const& child,
                                     const BaseDistMat& signals);

  virtual void back_prop_impl_() = 0;

  virtual void allocate_new_gradients_() = 0;

  virtual void propagate_error_signals_to_parents_() = 0;

  virtual void clear_prev_error_signals_() = 0;

  virtual void move_or_copy_prev_error_signal_(
    const Layer& child,
    std::unique_ptr<El::BaseDistMatrix> signal) = 0;

  virtual void
  view_or_copy_prev_error_signal_(const Layer& child,
                                  const El::BaseDistMatrix& signal) = 0;

  virtual void
  deep_copy_prev_error_signal_(const Layer& child,
                               const El::BaseDistMatrix& signal) = 0;


  // ===========================================================
  // Private class members
  // ===========================================================

  std::vector<ViewingLayerPtr> m_parent_layers;
  std::vector<ViewingLayerPtr> m_child_layers;

  std::vector<ViewingWeightsPtr> m_weights;

  std::vector<std::vector<int>> m_output_dims_list;

  ViewingLayerPtr m_hint_layer;

  ParallelStrategy m_parallel_strategy;

#ifdef LBANN_HAS_DISTCONV
private:
  friend class distconv_adapter;

public:
  bool distconv_enabled() const;
  virtual bool keep_original_inputs(int index) const;
  virtual bool keep_original_outputs(int index) const;
  virtual bool keep_original_gradient_wrt_inputs(int index) const;
  virtual bool keep_original_gradient_wrt_outputs(int index) const;
  virtual const distconv_adapter& get_distconv_adapter() const;
  virtual distconv_adapter& get_distconv_adapter();

protected:
  virtual bool is_distconv_supported() const { return false; }
  void prepare_distconv();
  virtual void setup_distconv_adapter() = 0;
  std::unique_ptr<distconv_adapter>& get_distconv_adapter_ptr()
  {
    return m_dc;
  };
  const std::unique_ptr<distconv_adapter>& get_distconv_adapter_ptr() const
  {
    return m_dc;
  };

private:
  mutable bool m_distconv_enabled = false;
  mutable bool m_distconv_enabled_set = false;
  std::unique_ptr<distconv_adapter> m_dc;
#else
public:
  bool distconv_enabled() const { return false; }

#endif // LBANN_HAS_DISTCONV
};

} // namespace lbann

#endif // LBANN_LAYERS_LAYER_HPP_INCLUDED