weights_proxy.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_WEIGHTS_WEIGHTS_PROXY_HPP_INCLUDED
#define LBANN_WEIGHTS_WEIGHTS_PROXY_HPP_INCLUDED

#include "lbann_config.hpp"

#include "lbann/base.hpp"
#include "lbann/weights/data_type_weights.hpp"
#include "lbann/weights/weights.hpp"

#if defined LBANN_DEBUG
#include "lbann/utils/exception.hpp"
#define LBANN_DEBUG_ASSERT_POINTER(ptr)                                        \
  do {                                                                         \
    if (!ptr)                                                                  \
      LBANN_ERROR("Pointer \"" #ptr "\" is null.");                            \
  } while (0)
#define LBANN_IN_DEBUG_MODE true
#else
#define LBANN_DEBUG_ASSERT_POINTER(ptr)
#define LBANN_IN_DEBUG_MODE false
#endif

namespace lbann {

template <typename TensorDataType>
class WeightsProxy
{
  using ValuesType = El::AbstractDistMatrix<TensorDataType>;
  using ValuesPtrType = std::unique_ptr<ValuesType>;

public:

  WeightsProxy() = default;

  WeightsProxy(ViewingWeightsPtr const& w)
  {
    if (!w.expired()) {
      this->setup(w);
    }
  }

  WeightsProxy(WeightsProxy const& other)
  {
    if (!other.master_weights_.expired()) {
      this->setup(other.master_weights_);
    }
  }

  template <typename T>
  WeightsProxy(WeightsProxy<T> const& other) : WeightsProxy()
  {
    auto ptr = other.master_weights_pointer();
    if (!ptr.expired()) {
      this->setup(ptr);
    }
  }

  WeightsProxy(WeightsProxy&& other) noexcept
    : master_weights_{std::move(other.master_weights_)},
      values_{std::move(other.values_)}
  {
    other.clear();
  }

  ~WeightsProxy() noexcept { this->clear(); }



  WeightsProxy& operator=(WeightsProxy const& other)
  {
    WeightsProxy<TensorDataType>(other).swap(*this);
    return *this;
  }

  template <typename T>
  WeightsProxy& operator=(WeightsProxy<T> const& other)
  {
    WeightsProxy<TensorDataType>(other).swap(*this);
    return *this;
  }

  WeightsProxy& operator=(WeightsProxy&& other) noexcept
  {
    // "Move-and-swap" idiom
    WeightsProxy<TensorDataType>(std::move(other)).swap(*this);
    return *this;
  }



  void clear() noexcept
  {
    master_weights_.reset();
    values_.reset();
  }

  void setup(ViewingWeightsPtr const& w)
  {
    master_weights_ = w;
    if (master_weights_.expired()) {
      values_.reset();
    }
    else {
      values_ = setup_values_(*master_weights_.lock());
    }
  }

  void synchronize_with_master()
  {
    if (!empty()) {
      const auto& master_values = master_weights_.lock()->get_values();
      if (values_->Viewing()) {
        El::LockedView(*values_,
                       dynamic_cast<const ValuesType&>(master_values));
      }
      else {
        El::Copy(master_values, *values_);
      }
    }
  }



  bool empty() const noexcept { return values_ == nullptr; }

  ValuesType const& values() const noexcept(!LBANN_IN_DEBUG_MODE)
  {
    LBANN_DEBUG_ASSERT_POINTER(values_);
    return *values_;
  }

  weights const& master_weights() const
  {
    LBANN_DEBUG_ASSERT_POINTER(master_weights_.lock());
    return *master_weights_.lock();
  }

  ViewingWeightsPtr master_weights_pointer() const
    noexcept(!LBANN_IN_DEBUG_MODE)
  {
    LBANN_DEBUG_ASSERT_POINTER(master_weights_.lock());
    return master_weights_;
  }



  void swap(WeightsProxy<TensorDataType>& other)
  {
    std::swap(master_weights_, other.master_weights_);
    std::swap(values_, other.values_);
  }


private:

  ValuesPtrType
  setup_values_(data_type_weights<TensorDataType> const& dtw) const
  {
    auto const& vals = dtw.get_values();
    ValuesPtrType ret(vals.Construct(vals.Grid(), vals.Root()));
    El::LockedView(*ret, vals);
    return ret;
  }

  template <typename OtherT>
  ValuesPtrType setup_values_(data_type_weights<OtherT> const& w) const
  {
    return setup_values_as_copy_(w);
  }

  ValuesPtrType setup_values_(weights const& w) const
  {
    if (auto dtw = dynamic_cast<data_type_weights<TensorDataType> const*>(&w))
      return setup_values_(*dtw);
    return setup_values_as_copy_(w);
  }

  ValuesPtrType setup_values_as_copy_(weights const& w) const
  {
    // In this case, w has some other dynamic type. So we need to
    // deep-copy every time. Thus, we allocate a target for this deep
    // copy here.
    ValuesPtrType ret{ValuesType::Instantiate(w.get_matrix_distribution())};
    ret->Resize(w.get_matrix_height(), w.get_matrix_width());
    return ret;
  }
private:
  // These members should never observably differ in nullity.

  ViewingWeightsPtr master_weights_;

  ValuesPtrType values_;

};

// Conform to LBANN's scheme
template <typename TensorDataType>
using weights_proxy = WeightsProxy<TensorDataType>;

} // namespace lbann
#undef LBANN_IN_DEBUG_MODE
#endif // LBANN_WEIGHTS_WEIGHTS_PROXY_HPP_INCLUDED