serialize_matrices.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.
#pragma once
#ifndef LBANN_UTILS_SERIALIZATION_SERIALIZE_MATRICES_HPP_
#define LBANN_UTILS_SERIALIZATION_SERIALIZE_MATRICES_HPP_

#include "cereal_utils.hpp"
#include "rooted_archive_adaptor.hpp"

#include <lbann/utils/exception.hpp>

#include <El.hpp>
#include <stdexcept>

// These really belong in Elemental; let's just extend that.
namespace El {

template <typename ArchiveT, typename T>
void save(ArchiveT& ar, ::El::AbstractMatrix<T> const& mat);

template <typename ArchiveT,
          typename T,
          ::El::Device D,
          lbann::utils::WhenTextArchive<ArchiveT> = 1>
void save(ArchiveT& ar, ::El::Matrix<T, D> const& mat);

namespace details {
template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void do_save(ArchiveT& ar, ::El::Matrix<T, ::El::Device::CPU> const& mat);

#ifdef LBANN_HAS_GPU
template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void do_save(ArchiveT& ar, ::El::Matrix<T, ::El::Device::GPU> const& mat);
#endif // LBANN_HAS_GPU
} // namespace details

template <typename ArchiveT,
          typename T,
          ::El::Device D,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void save(ArchiveT& ar, ::El::Matrix<T, D> const& mat);

// Special treatment for the rooted archive.
template <typename ArchiveT, typename T, ::El::Device D>
void save(lbann::RootedOutputArchiveAdaptor<ArchiveT>& ar,
          ::El::Matrix<T, D> const& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void load(ArchiveT& archive, ::El::AbstractMatrix<T>& mat);

template <typename ArchiveT,
          typename T,
          ::El::Device D,
          lbann::utils::WhenTextArchive<ArchiveT> = 1>
void load(ArchiveT& archive, ::El::Matrix<T, D>& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void load(ArchiveT& archive, ::El::Matrix<T, ::El::Device::CPU>& mat);

#if defined LBANN_HAS_GPU

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void load(ArchiveT& archive, ::El::Matrix<T, ::El::Device::GPU>& mat);
#endif // defined LBANN_HAS_GPU

template <typename ArchiveT, typename T, ::El::Device D>
void load(lbann::RootedInputArchiveAdaptor<ArchiveT>& ar,
          ::El::Matrix<T, D>& mat);

// DistMatrix

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenTextArchive<ArchiveT> = 1>
void save(ArchiveT& ar, ::El::AbstractDistMatrix<T> const& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenTextArchive<ArchiveT> = 1>
void load(ArchiveT& ar, ::El::AbstractDistMatrix<T>& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void save(ArchiveT& ar, ::El::AbstractDistMatrix<T> const& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void load(ArchiveT& ar, ::El::AbstractDistMatrix<T>& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenTextArchive<ArchiveT> = 1>
void save(lbann::RootedOutputArchiveAdaptor<ArchiveT>& ar,
          ::El::AbstractDistMatrix<T> const& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenTextArchive<ArchiveT> = 1>
void load(lbann::RootedInputArchiveAdaptor<ArchiveT>& ar,
          ::El::AbstractDistMatrix<T>& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void save(lbann::RootedOutputArchiveAdaptor<ArchiveT>& ar,
          ::El::AbstractDistMatrix<T> const& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void load(lbann::RootedInputArchiveAdaptor<ArchiveT>& ar,
          ::El::AbstractDistMatrix<T>& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void save(lbann::RootedOutputArchiveAdaptor<ArchiveT>& ar,
          ::El::DistMatrix<T, ::El::CIRC, ::El::CIRC> const& mat);

template <typename ArchiveT,
          typename T,
          lbann::utils::WhenNotTextArchive<ArchiveT> = 1>
void load(lbann::RootedInputArchiveAdaptor<ArchiveT>& ar,
          ::El::DistMatrix<T, ::El::CIRC, ::El::CIRC>& mat);

} // namespace El

// Dealing with smart pointers and object construction

namespace lbann {
namespace utils {

struct grid_manager
{
  grid_manager(Grid const& g);
  ~grid_manager();
};

Grid const& get_current_grid() noexcept;

lbann_comm& get_current_comm() noexcept;
} // namespace utils
} // namespace lbann

namespace cereal {

template <typename DataT,
          ::El::Dist CDist,
          ::El::Dist RDist,
          ::El::DistWrap Wrap,
          ::El::Device D>
struct LoadAndConstruct<::El::DistMatrix<DataT, CDist, RDist, Wrap, D>>
{
  using DistMatrixType = ::El::DistMatrix<DataT, CDist, RDist, Wrap, D>;
  using CircMatrixType =
    ::El::DistMatrix<DataT, ::El::CIRC, ::El::CIRC, Wrap, ::El::Device::CPU>;

  template <
    typename ArchiveT,
    ::h2::meta::EnableWhen<::lbann::utils::IsBuiltinArchive<ArchiveT>, int> = 0>
  static void load_and_construct(ArchiveT& ar,
                                 cereal::construct<DistMatrixType>& construct);

  template <typename ArchiveT>
  static void load_and_construct(lbann::RootedInputArchiveAdaptor<ArchiveT>& ar,
                                 cereal::construct<DistMatrixType>& construct);
}; // struct LoadAndConstruct<::El::DistMatrix<DataT, CDist, RDist, Wrap, D>>
} // namespace cereal

#endif // LBANN_UTILS_SERIALIZATION_SERIALIZE_MATRICES_HPP_