param.h

/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); 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 SINGA_UTILS_PARAM_H_
#define SINGA_UTILS_PARAM_H_

#include <memory>
#include <string>
#include <vector>
#include "communication/msg.h"
#include "proto/job.pb.h"
#include "utils/blob.h"

namespace singa {

class ParamGenerator {
 public:
  static ParamGenerator* Create(const ParamGenProto& proto);

  virtual ~ParamGenerator() {}

  virtual void Init(const ParamGenProto& proto) { proto_ = proto; }
  virtual void Fill(Blob<float>* data);

 protected:
  ParamGenProto proto_;
};

class GaussianGen : public ParamGenerator {
 public:
  void  Fill(Blob<float>* data) override;
};

class GaussianSqrtFanInGen : public GaussianGen {
 public:
  void  Fill(Blob<float>* data) override;
};

class UniformGen : public ParamGenerator {
 public:
  void  Fill(Blob<float>* data) override;
};

class UniformSqrtFanInGen : public UniformGen {
 public:
  void Fill(Blob<float>* data) override;
};

class UniformSqrtFanInOutGen : public UniformGen {
 public:
  void Fill(Blob<float>* data) override;
};

class Param {
 public:
  static Param* Create(const ParamProto& proto);

  Param() {}
  virtual ~Param() {}
  void Init(const ParamProto& proto) { proto_ = proto; }
  virtual void Setup(const std::vector<int>& shape);
  /*
   * Fill the values according to init method, e.g., gaussian distribution.
   *
   * @param version initial version
   */
  virtual void InitValues();
  virtual void InitValues(int version);
  void ShareFrom(const Param& other);
  void FromProto(const BlobProto& blob);
  void ToProto(BlobProto* blob);
  void AddSlice(int slice_id, int size);
  inline float lr_scale() const { return proto_.lr_scale(); }
  inline float wd_scale() const { return proto_.wd_scale(); }
  inline const std::string& name() const { return proto_.name(); }
  inline void set_name(const std::string& name) { proto_.set_name(name); }
  inline int owner() const { return proto_.owner(); }
  inline int id() const { return proto_.id(); }
  inline void set_id(int id) {
    proto_.set_id(id);
    proto_.set_owner(id);
  }
  inline int version() const { return data_->version(); }
  inline void set_version(int v) { data_->set_version(v); }
  inline int local_version() const { return local_version_; }
  inline void set_local_version(int v) { local_version_ = v; }
  inline const std::string& share_from() const { return proto_.share_from(); }
  inline int size() const { return data_->count(); }
  inline const Blob<float>& data() const { return *data_; }
  inline Blob<float>* mutable_data() { return data_.get(); }
  inline const Blob<float> &grad() const { return grad_; }
  inline Blob<float> *mutable_grad() { return &grad_; }
  inline float* mutable_cpu_data() { return data_->mutable_cpu_data(); }
  inline float* mutable_cpu_grad() { return grad_.mutable_cpu_data(); }
  inline float* mutable_cpu_history() { return history_.mutable_cpu_data(); }
  inline int slice_start() const { return slice_start_; }
  inline int num_slices() const { return num_slices_; }

  virtual Msg* GenPutMsg(bool copy, int slice_idx);
  virtual Msg* GenGetMsg(bool copy, int slice_idx);
  virtual Msg* GenUpdateMsg(bool copy, int slice_idx);
  virtual Msg* GenSyncMsg(int offset, int size);
  virtual Msg* HandlePutMsg(Msg** msg, bool reserve);
  virtual Msg* HandleGetMsg(Msg** msg, bool reserve);
  virtual void ParseUpdateMsgs(const std::vector<Msg*>& msgs);
  virtual const std::vector<Msg*>
    GenUpdateResponseMsgs(std::vector<Msg*>* msgs, bool reserve);
  virtual Msg* HandleSyncMsg(Msg** msg, bool reserve);
  virtual int ParseGetResponseMsg(Msg* msg, int slice_idx);
  virtual int ParseUpdateResponseMsg(Msg* msg, int slice_idx);
  virtual int ParseSyncResponseMsg(Msg* msg, int slice_idx);

 protected:
  void ParseResponseMsg(Msg* msg, int slice_idx);

 protected:
  int local_version_ = -1;
  // the ID of the first slice
  int slice_start_ = 0;
  int num_slices_ = 0;
  // offset and size of each slice
  std::vector<int> slice_offset_;
  std::vector<int> slice_size_;
  // for debug checking
  // since put request has no feedback, we do not track its pending status
  std::vector<bool> pending_get_;
  std::vector<bool> pending_update_;
  int num_pending_requests_ = 0;
  // data field
  std::shared_ptr<Blob<float>> data_ = nullptr;
  // gradient, history gradient of this parameter
  Blob<float> grad_, history_;
  ParamProto proto_;
};

class ParamEntry {
 public:
  ParamEntry() {}
  ParamEntry(int total, Param* p);
  void AddParam(bool local, Param* p);
  int next_version = -1;  // next_version & num_update are directly used by stub
  int num_update = 0;
  int num_local = 0;  
  int num_total = 0;  
  std::vector<Param*> shares;
};

inline int ParamTrgt(int param_id, int slice_id) {
  return (param_id << 16) | slice_id;
}

inline int ParamID(int param_trgt) {
  return param_trgt >> 16;
}

inline int SliceID(int param_trgt) {
  static const int mask = (1 << 16) -1;
  return param_trgt & mask;
}

}  // namespace singa

#endif  // SINGA_UTILS_PARAM_H_