operator.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_OPERATORS_OPERATOR_HPP_INCLUDED
#define LBANN_OPERATORS_OPERATOR_HPP_INCLUDED

#include "lbann/base.hpp"
#include "lbann/proto/datatype_helpers.hpp"
#include "lbann/utils/cloneable.hpp"
#include "lbann/utils/describable.hpp"
#include "lbann/utils/exception.hpp"
#include "lbann/utils/protobuf_serializable.hpp"
#include "lbann/utils/tensor.hpp"
#include "lbann/utils/typename.hpp"

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

#include <h2/meta/Core.hpp>
#include <h2/meta/TypeList.hpp>

#include <google/protobuf/message.h>

#include <string>
#include <vector>

namespace lbann {

using supported_operator_data_type = h2::meta::TL<
#ifdef LBANN_HAS_GPU_FP16
  fp16,
#endif
#ifdef LBANN_HAS_HALF
  cpu_fp16,
#endif
  float,
  double,
  El::Complex<float>,
  El::Complex<double>>;

template <typename InputT, typename OutputT, El::Device D>
class Operator : public AbstractCloneableBase<Operator<InputT, OutputT, D>>,
                 public Describable,
                 public ProtobufSerializable
{
public:
  using InputTensorType = utils::DistTensorView<InputT, D>;
  using OutputTensorType = utils::DistTensorView<OutputT, D>;
  using ConstInputTensorType = utils::ConstDistTensorView<InputT, D>;
  using ConstOutputTensorType = utils::ConstDistTensorView<OutputT, D>;

public:
  static_assert(
    h2::meta::tlist::MemberV<InputT, supported_operator_data_type>(),
    "Must use a supported input type.");
  static_assert(
    h2::meta::tlist::MemberV<OutputT, supported_operator_data_type>(),
    "Must use a supported output type.");

  Operator() = default;
  virtual ~Operator() = default;

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

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

  Description get_description() const override;

  void write_proto(google::protobuf::Message& msg) const;


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



  virtual void
  fp_compute(std::vector<ConstInputTensorType> const& inputs,
             std::vector<OutputTensorType> const& outputs) const = 0;

  virtual void
  bp_compute(std::vector<ConstInputTensorType> const& inputs,
             std::vector<ConstOutputTensorType> const& gradient_wrt_outputs,
             std::vector<InputTensorType> const& gradient_wrt_inputs) const;

protected:
  Operator(Operator&& other) noexcept = default;
  Operator& operator=(Operator&& other) noexcept = default;
  Operator(Operator const& other) = default;
  Operator& operator=(Operator const& other) = default;

private:
  virtual void set_proto_params(lbann_data::Operator&) const = 0;
  virtual void do_fill_description(Description&) const = 0;
};

template <typename InputT, typename OutputT, El::Device D>
void Operator<InputT, OutputT, D>::write_proto(
  google::protobuf::Message& msg) const
{
  lbann_data::Operator operator_msg;
  operator_msg.set_input_datatype(proto::ProtoDataType<InputT>);
  operator_msg.set_output_datatype(proto::ProtoDataType<OutputT>);
  operator_msg.set_device_allocation(proto::ProtoDevice<D>);

  this->set_proto_params(operator_msg);

  msg.CopyFrom(operator_msg);
}

template <typename InputT, typename OutputT, El::Device D>
Description Operator<InputT, OutputT, D>::get_description() const
{

  // Construct description object
  Description desc(this->get_type());

  // DataType
  desc.add("Input data type", TypeName<InputT>());
  desc.add("Output data type", TypeName<OutputT>());

  this->do_fill_description(desc);

  return desc;
}

template <typename InputT, typename OutputT, El::Device D>
void Operator<InputT, OutputT, D>::bp_compute(
  std::vector<ConstInputTensorType> const&,
  std::vector<ConstOutputTensorType> const&,
  std::vector<InputTensorType> const&) const
{}

template <typename InputT, typename OutputT, El::Device D>
template <typename ArchiveT>
void Operator<InputT, OutputT, D>::serialize(ArchiveT& ar)
{}

} // namespace lbann
#endif // LBANN_OPERATORS_OPERATOR_HPP_INCLUDED