Add ps optimizer info.

mindspore/ccsrc/parallel/CMakeLists.txt

 file(GLOB_RECURSE _PARALLEL_SRC_FILES RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.cc")
-list(REMOVE_ITEM _PARALLEL_SRC_FILES  "ps/util.cc" "ps/scheduler.cc")
+list(REMOVE_ITEM _PARALLEL_SRC_FILES  "ps/util.cc" "ps/scheduler.cc" "ps/optimizer_info.cc")
 if (ENABLE_DUMP_PROTO)
     list(REMOVE_ITEM _PARALLEL_SRC_FILES "parallel/strategy_checkpoint/parallel_strategy_checkpoint.cc")
 endif ()

mindspore/ccsrc/parallel/ps/optimizer_info.cc

+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+
+#include "parallel/ps/optimizer_info.h"
+#include <memory>
+
+namespace mindspore {
+namespace parallel {
+namespace ps {
+void OptimizerInfo::AddWorkspace(const AddressPtr &workspace) { workspaces_.push_back(workspace); }
+
+const std::vector<AddressPtr> &OptimizerInfo::inputs() { return inputs_; }
+
+const std::vector<AddressPtr> &OptimizerInfo::workspaces() { return workspaces_; }
+
+const std::vector<AddressPtr> &OptimizerInfo::outputs() { return outputs_; }
+
+bool OptimizerInfo::IsSparse() const { return false; }
+
+size_t OptimizerInfo::grad_index() { return 0; }
+
+size_t OptimizerInfo::indices_index() { return 0; }
+
+void OptimizerInfo::UpdateWeight(const WeightPtr &weight) {
+  AddressPtr weight_addr = std::make_shared<kernel::Address>();
+  weight_addr->addr = weight->data();
+  weight_addr->size = weight->size();
+  inputs_[0] = weight_addr;
+}
+
+void DenseOptimInfo::Accumulate(const Values &values, const Lengths &lengths) {
+  float *accum_grad_data = reinterpret_cast<float *>(gradient()->addr);
+  size_t size = gradient()->size / sizeof(float);
+  size_t grad_index = this->grad_index();
+  size_t grad_offset = 0;
+  for (size_t i = 0; i < grad_index; i++) {
+    grad_offset += lengths[i];
+  }
+  float *grad_data = values.data() + grad_offset;
+  CHECK_EQ(size, static_cast<size_t>(lengths[grad_index]));
+
+  for (size_t i = 0; i < size; i++) {
+    accum_grad_data[i] += grad_data[i];
+  }
+}
+
+void SparseOptimInfo::Accumulate(const Values &values, const Lengths &lengths) {
+  // Append grad data to the end
+  float *accum_grad_data = reinterpret_cast<float *>(gradient()->addr);
+
+  size_t grad_index = this->grad_index();
+  size_t grad_offset = 0;
+  for (size_t i = 0; i < grad_index; i++) {
+    grad_offset += lengths[i];
+  }
+  float *incr_grad_data = values.data() + grad_offset;
+  size_t incr_grad_size = lengths[grad_index] * sizeof(float);
+
+  auto ret = memcpy_s(accum_grad_data + grads_offset_, incr_grad_size, incr_grad_data, incr_grad_size);
+  if (ret != 0) {
+    MS_LOG(EXCEPTION) << "memcpy_s error, errorno(" << ret << ")";
+  }
+  grads_offset_ += incr_grad_size;
+  gradient()->size += incr_grad_size;
+
+  // Append indice data to the end
+  int *accum_indices_data = reinterpret_cast<int *>(indices()->addr);
+
+  size_t indices_index = this->indices_index();
+  size_t indice_offset = 0;
+  for (size_t i = 0; i < indices_index; i++) {
+    indice_offset += lengths[i];
+  }
+  int *incr_indice_data = reinterpret_cast<int *>(values.data() + indice_offset);
+  size_t incr_indice_size = lengths[indices_index] * sizeof(float);
+
+  auto ret2 = memcpy_s(accum_indices_data + indices_offset_, incr_indice_size, incr_indice_data, incr_indice_size);
+  if (ret2 != 0) {
+    MS_LOG(EXCEPTION) << "memcpy_s error, errorno(" << ret2 << ")";
+  }
+  indices_offset_ += incr_indice_size;
+  indices()->size += incr_indice_size;
+}
+
+void SparseOptimInfo::Reset() {
+  auto &gradient = this->gradient();
+  gradient->size = 0;
+  auto &indices = this->indices();
+  indices->size = 0;
+  grads_offset_ = 0;
+  indices_offset_ = 0;
+}
+
+MomentumOptimInfo::MomentumOptimInfo(const AddressPtr &weight, const AddressPtr &accumulate,
+                                     const AddressPtr &learning_rate, const AddressPtr &gradient,
+                                     const AddressPtr &momentum) {
+  inputs_.push_back(weight);
+  inputs_.push_back(accumulate);
+  inputs_.push_back(learning_rate);
+  inputs_.push_back(gradient);
+  inputs_.push_back(momentum);
+}
+
+const AddressPtr &MomentumOptimInfo::gradient() { return inputs_[3]; }
+
+const AddressPtr &MomentumOptimInfo::indices() { return inputs_[3]; }
+
+SparseAdamOptimInfo::SparseAdamOptimInfo(const AddressPtr &weight, const AddressPtr &m, const AddressPtr &v,
+                                         const AddressPtr &beta1_power, const AddressPtr &beta2_power,
+                                         const AddressPtr &learning_rate, const AddressPtr &beta1,
+                                         const AddressPtr &beta2, const AddressPtr &epsilon, const AddressPtr &grad,
+                                         const AddressPtr &indices, size_t grads_offset, size_t indices_offset) {
+  inputs_.push_back(weight);
+  inputs_.push_back(m);
+  inputs_.push_back(v);
+  inputs_.push_back(beta1_power);
+  inputs_.push_back(beta2_power);
+  inputs_.push_back(learning_rate);
+  inputs_.push_back(beta1);
+  inputs_.push_back(beta2);
+  inputs_.push_back(epsilon);
+  inputs_.push_back(grad);
+  inputs_.push_back(indices);
+  grads_offset_ = grads_offset;
+  indices_offset_ = indices_offset;
+}
+
+void SparseAdamOptimInfo::Update(const Values &values, const Lengths &lens) {
+  void *data_ptr = values.data();
+  AddressPtr beta1_power = inputs_[3];
+  size_t size = values.size() * sizeof(float);
+  auto ret = memcpy_s(beta1_power->addr, size, data_ptr, size);
+  if (ret != 0) {
+    MS_LOG(EXCEPTION) << "memcpy_s error, errorno(" << ret << ")";
+  }
+}
+
+const AddressPtr &SparseAdamOptimInfo::gradient() { return inputs_[9]; }
+
+const AddressPtr &SparseAdamOptimInfo::indices() { return inputs_[10]; }
+
+bool SparseAdamOptimInfo::IsSparse() const { return true; }
+
+size_t SparseAdamOptimInfo::grad_index() { return 6; }
+
+size_t SparseAdamOptimInfo::indices_index() { return 7; }
+
+SparseFtrlOptimInfo::SparseFtrlOptimInfo(const AddressPtr &weight, const AddressPtr &accum, const AddressPtr &linear,
+                                         const AddressPtr &grad, const AddressPtr &indices, size_t grads_offset,
+                                         size_t indices_offset) {
+  inputs_.push_back(weight);
+  inputs_.push_back(accum);
+  inputs_.push_back(linear);
+  inputs_.push_back(grad);
+  inputs_.push_back(indices);
+  grads_offset_ = grads_offset;
+  indices_offset_ = indices_offset;
+}
+
+const AddressPtr &SparseFtrlOptimInfo::gradient() { return inputs_[3]; }
+
+const AddressPtr &SparseFtrlOptimInfo::indices() { return inputs_[4]; }
+
+bool SparseFtrlOptimInfo::IsSparse() const { return true; }
+
+size_t SparseFtrlOptimInfo::grad_index() { return 0; }
+
+size_t SparseFtrlOptimInfo::indices_index() { return 1; }
+}  // namespace ps
+}  // namespace parallel
+}  // namespace mindspore

mindspore/ccsrc/parallel/ps/optimizer_info.h

+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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 MINDSPORE_MINDSPORE_CCSRC_PARALLEL_PS_OPTIMIZER_INFO_H_
+#define MINDSPORE_MINDSPORE_CCSRC_PARALLEL_PS_OPTIMIZER_INFO_H_
+
+#include <vector>
+#include "kernel/kernel.h"
+#include "parallel/ps/common.h"
+
+namespace mindspore {
+namespace parallel {
+namespace ps {
+using mindspore::kernel::AddressPtr;
+class OptimizerInfo {
+ public:
+  OptimizerInfo() = default;
+  virtual ~OptimizerInfo() = default;
+
+  virtual void Update(const Values &values, const Lengths &lengths) {}
+  virtual void UpdateWeight(const WeightPtr &weight);
+  virtual void Accumulate(const Values &values, const Lengths &lengths) = 0;
+  virtual void Reset() {}
+  void AddWorkspace(const AddressPtr &workspace);
+
+  virtual const AddressPtr &gradient() = 0;
+  virtual const AddressPtr &indices() = 0;
+  const std::vector<AddressPtr> &inputs();
+  const std::vector<AddressPtr> &workspaces();
+  const std::vector<AddressPtr> &outputs();
+
+  virtual bool IsSparse() const;
+  virtual size_t grad_index();
+  virtual size_t indices_index();
+
+ protected:
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspaces_;
+  std::vector<AddressPtr> outputs_;
+};
+
+class DenseOptimInfo : public OptimizerInfo {
+ public:
+  DenseOptimInfo() = default;
+  ~DenseOptimInfo() override = default;
+
+  void Accumulate(const Values &values, const Lengths &lens) override;
+};
+
+class SparseOptimInfo : public OptimizerInfo {
+ public:
+  SparseOptimInfo() = default;
+  ~SparseOptimInfo() override = default;
+
+  void Accumulate(const Values &values, const Lengths &lens) override;
+  void Reset() override;
+
+ protected:
+  size_t grads_offset_{0};
+  size_t indices_offset_{0};
+};
+
+class MomentumOptimInfo : public DenseOptimInfo {
+ public:
+  MomentumOptimInfo(const AddressPtr &weight, const AddressPtr &accumulate, const AddressPtr &learning_rate,
+                    const AddressPtr &gradient, const AddressPtr &momentum);
+  ~MomentumOptimInfo() override = default;
+
+  const AddressPtr &gradient();
+  const AddressPtr &indices();
+};
+
+class SparseAdamOptimInfo : public SparseOptimInfo {
+ public:
+  SparseAdamOptimInfo(const AddressPtr &weight, const AddressPtr &m, const AddressPtr &v, const AddressPtr &beta1_power,
+                      const AddressPtr &beta2_power, const AddressPtr &learning_rate, const AddressPtr &beta1,
+                      const AddressPtr &beta2, const AddressPtr &epsilon, const AddressPtr &grad,
+                      const AddressPtr &indices, size_t grads_offset, size_t indices_offset);
+  ~SparseAdamOptimInfo() override = default;
+
+  void Update(const Values &values, const Lengths &lens) override;
+  const AddressPtr &gradient();
+  const AddressPtr &indices();
+  bool IsSparse() const override;
+  size_t grad_index() override;
+  size_t indices_index() override;
+};
+
+class SparseFtrlOptimInfo : public SparseOptimInfo {
+ public:
+  SparseFtrlOptimInfo(const AddressPtr &weight, const AddressPtr &accum, const AddressPtr &linear,
+                      const AddressPtr &grad, const AddressPtr &indices, size_t grads_offset, size_t indices_offset);
+  ~SparseFtrlOptimInfo() override = default;
+
+  const AddressPtr &gradient();
+  const AddressPtr &indices();
+  bool IsSparse() const override;
+  size_t grad_index() override;
+  size_t indices_index() override;
+};
+}  // namespace ps
+}  // namespace parallel
+}  // namespace mindspore
+#endif  // MINDSPORE_MINDSPORE_CCSRC_PARALLEL_PS_OPTIMIZER_INFO_H_

tests/ut/cpp/CMakeLists.txt

@@ -117,6 +117,7 @@ list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/ir/lite/tensor.cc"
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/parallel/strategy_checkpoint/parallel_strategy_checkpoint.cc")
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/parallel/ps/util.cc")
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/parallel/ps/scheduler.cc")
+list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/parallel/ps/optimizer_info.cc")
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/utils/anf_ir.pb.cc")
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/utils/node_strategy.pb.cc")
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/utils/load_onnx/anf_model_parser.cc")