forked from huawei/mindspore2022
249 lines
11 KiB
C++
249 lines
11 KiB
C++
/**
|
|
* Copyright 2020-2022 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 "ps/optimizer_info_builder.h"
|
|
#include <vector>
|
|
#include <memory>
|
|
#include <functional>
|
|
#include "plugin/device/cpu/kernel/ps/sparse_apply_ftrl_ps_kernel.h"
|
|
|
|
namespace mindspore {
|
|
namespace ps {
|
|
using mindspore::kernel::ps::SparseApplyFtrlPSKernelMod;
|
|
OptimizerInfo *OptimizerInfoBuilder::Build(const std::shared_ptr<PServerKernel> &pserver_kernel,
|
|
const WeightPtr &weight, const Keys &keys, const Values &values,
|
|
const Lengths &lens, const InputsShapePtr &inputs_shape, size_t worker_num,
|
|
bool sharded) {
|
|
MS_EXCEPTION_IF_NULL(pserver_kernel);
|
|
MS_EXCEPTION_IF_NULL(weight);
|
|
MS_EXCEPTION_IF_NULL(inputs_shape);
|
|
OptimizerInfo *optim_info =
|
|
BuildInputs(weight, keys, values, lens, inputs_shape, worker_num, pserver_kernel, sharded);
|
|
MS_EXCEPTION_IF_NULL(optim_info);
|
|
std::vector<size_t> ws_sizes = pserver_kernel->workspace_sizes();
|
|
BuildWorkspaces(optim_info, ws_sizes, worker_num);
|
|
BuildOutputs(optim_info, worker_num);
|
|
return optim_info;
|
|
}
|
|
|
|
void OptimizerInfoBuilder::BuildWorkspaces(OptimizerInfo *info, const std::vector<size_t> &ws_sizes, size_t) {
|
|
MS_EXCEPTION_IF_NULL(info);
|
|
for (size_t i = 0; i < ws_sizes.size(); i++) {
|
|
size_t size = ws_sizes[i];
|
|
AddressPtr workspace = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(workspace);
|
|
workspace->addr = new float[size];
|
|
MS_EXCEPTION_IF_NULL(workspace->addr);
|
|
workspace->size = size;
|
|
info->AddWorkspace(workspace);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
AddressPtr OptimizerInfoBuilder::GenInputAddrPtr(const std::string &optim_type, const std::string &input_name,
|
|
void *ps_data, const Lengths &ps_lens,
|
|
const InputsShapePtr &inputs_shape) {
|
|
MS_EXCEPTION_IF_NULL(ps_data);
|
|
// Take note of that the data type maybe inconsistent in ps_data.
|
|
MS_LOG(INFO) << "Get input address pointer for optimizer:" << optim_type << ", input name:" << input_name;
|
|
AddressPtr addr_ptr = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(addr_ptr);
|
|
|
|
if (kOptimToOriginIdx.count(optim_type) == 0 || kOptimToPSSendIdx.count(optim_type) == 0) {
|
|
MS_LOG(EXCEPTION) << "Optimizer type " << optim_type << " in not supported.";
|
|
}
|
|
const OptimOriginIdx &origin_input_map = kOptimToOriginIdx.at(optim_type);
|
|
const OptimPSSendIdx &ps_send_index_map = kOptimToPSSendIdx.at(optim_type);
|
|
if (ps_send_index_map.count(input_name) == 0 || origin_input_map.count(input_name) == 0) {
|
|
MS_LOG(EXCEPTION) << "Optimizer " << optim_type << " has no input for " << input_name;
|
|
}
|
|
size_t ps_index = ps_send_index_map.at(input_name);
|
|
if (ps_index == INDEX_NOT_SEND) {
|
|
MS_LOG(EXCEPTION) << "Input " << input_name << " is not supposed to be sent to PS.";
|
|
}
|
|
|
|
size_t addr_data_size, addr_data_offset;
|
|
if (inputs_shape != nullptr) {
|
|
// addr_data_size should be calculated by inputs_shape if it's passed.
|
|
size_t origin_index = origin_input_map.at(input_name);
|
|
EXC_IF_VEC_IDX_OOB((*inputs_shape), origin_index);
|
|
MS_EXCEPTION_IF_NULL((*inputs_shape)[origin_index]);
|
|
auto shape = *((*inputs_shape)[origin_index]);
|
|
addr_data_size = std::accumulate(shape.begin(), shape.end(), worker_num_, std::multiplies<size_t>());
|
|
} else {
|
|
EXC_IF_VEC_IDX_OOB(ps_lens, ps_index);
|
|
addr_data_size = IntToSize(ps_lens[ps_index]);
|
|
}
|
|
addr_data_offset =
|
|
IntToSize(std::accumulate(ps_lens.begin(), ps_lens.begin() + SizeToInt(ps_index), 0, std::plus<int>()));
|
|
|
|
// The size in ps_lens instead of addr_data_size is the size of real data.
|
|
T *buffer = new T[addr_data_size];
|
|
addr_ptr->size = IntToSize(ps_lens[ps_index]) * sizeof(T);
|
|
addr_ptr->addr = buffer;
|
|
|
|
size_t dst_size = addr_ptr->size;
|
|
size_t src_size = addr_ptr->size;
|
|
void *dst_data = addr_ptr->addr;
|
|
void *src_data = reinterpret_cast<T *>(ps_data) + addr_data_offset;
|
|
MS_EXCEPTION_IF_NULL(dst_data);
|
|
MS_EXCEPTION_IF_NULL(src_data);
|
|
int64_t ret = memcpy_s(dst_data, dst_size, src_data, src_size);
|
|
if (ret != 0) {
|
|
MS_LOG(EXCEPTION) << "memcpy_s error, errorno(" << ret << ")";
|
|
delete[] buffer;
|
|
buffer = nullptr;
|
|
return nullptr;
|
|
}
|
|
return addr_ptr;
|
|
}
|
|
|
|
OptimizerInfo *MomentumOptimInfoBuilder::BuildInputs(const WeightPtr &weight, const Keys &, const Values &values,
|
|
const Lengths &lens, const InputsShapePtr &, size_t,
|
|
const std::shared_ptr<PServerKernel> &, bool) {
|
|
MS_EXCEPTION_IF_NULL(weight);
|
|
AddressPtr weight_addr = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(weight_addr);
|
|
weight_addr->addr = weight->data();
|
|
weight_addr->size = weight->size() * sizeof(float);
|
|
|
|
AddressPtr accumulate = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(accumulate);
|
|
|
|
accumulate->addr = new float[weight->size()];
|
|
MS_EXCEPTION_IF_NULL(accumulate->addr);
|
|
accumulate->size = sizeof(float) * weight->size();
|
|
int64_t ret = memset_s(accumulate->addr, accumulate->size, 0x00, accumulate->size);
|
|
if (ret != 0) {
|
|
MS_LOG(EXCEPTION) << "memset_s error, errorno(" << ret << ")";
|
|
delete[] reinterpret_cast<float *>(accumulate->addr);
|
|
accumulate->addr = nullptr;
|
|
return nullptr;
|
|
}
|
|
|
|
AddressPtr learning_rate = GenInputAddrPtr<float>(kApplyMomentum, "lr", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(learning_rate);
|
|
AddressPtr gradient = GenInputAddrPtr<float>(kApplyMomentum, "grad", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(gradient);
|
|
AddressPtr momentum = GenInputAddrPtr<float>(kApplyMomentum, "momentum", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(momentum);
|
|
return new MomentumOptimInfo(weight_addr, accumulate, learning_rate, gradient, momentum);
|
|
}
|
|
|
|
OptimizerInfo *SparseAdamOptimInfoBuilder::BuildInputs(const WeightPtr &weight, const Keys &, const Values &values,
|
|
const Lengths &lens, const InputsShapePtr &inputs_shape, size_t,
|
|
const std::shared_ptr<PServerKernel> &, bool sharded) {
|
|
AddressPtr weight_addr = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(weight_addr);
|
|
weight_addr->addr = weight->data();
|
|
weight_addr->size = weight->size() * sizeof(float);
|
|
|
|
AddressPtr m = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(m);
|
|
|
|
m->addr = new float[weight->size()];
|
|
MS_EXCEPTION_IF_NULL(m->addr);
|
|
m->size = weight->size() * sizeof(float);
|
|
int64_t ret = memset_s(m->addr, m->size, 0x00, m->size);
|
|
if (ret != 0) {
|
|
MS_LOG(EXCEPTION) << "memset_s error, errorno(" << ret << ")";
|
|
delete[] reinterpret_cast<float *>(m->addr);
|
|
m->addr = nullptr;
|
|
return nullptr;
|
|
}
|
|
|
|
AddressPtr v = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(v);
|
|
|
|
v->addr = new float[weight->size()];
|
|
MS_EXCEPTION_IF_NULL(v->addr);
|
|
v->size = weight->size() * sizeof(float);
|
|
ret = memset_s(v->addr, v->size, 0x00, v->size);
|
|
if (ret != 0) {
|
|
MS_LOG(EXCEPTION) << "memset_s error, errorno(" << ret << ")";
|
|
delete[] reinterpret_cast<float *>(v->addr);
|
|
v->addr = nullptr;
|
|
delete[] reinterpret_cast<float *>(m->addr);
|
|
m->addr = nullptr;
|
|
return nullptr;
|
|
}
|
|
|
|
AddressPtr beta1_power = GenInputAddrPtr<float>(kSparseAdam, "beta1_power", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(beta1_power);
|
|
AddressPtr beta2_power = GenInputAddrPtr<float>(kSparseAdam, "beta2_power", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(beta2_power);
|
|
AddressPtr learning_rate = GenInputAddrPtr<float>(kSparseAdam, "lr", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(learning_rate);
|
|
AddressPtr beta1 = GenInputAddrPtr<float>(kSparseAdam, "beta1", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(beta1);
|
|
AddressPtr beta2 = GenInputAddrPtr<float>(kSparseAdam, "beta2", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(beta2);
|
|
AddressPtr epsilon = GenInputAddrPtr<float>(kSparseAdam, "eps", const_cast<float *>(values.data()), lens);
|
|
MS_EXCEPTION_IF_NULL(epsilon);
|
|
AddressPtr grad = GenInputAddrPtr<float>(kSparseAdam, "grad", const_cast<float *>(values.data()), lens, inputs_shape);
|
|
MS_EXCEPTION_IF_NULL(grad);
|
|
AddressPtr indices =
|
|
GenInputAddrPtr<float>(kSparseAdam, "indices", const_cast<float *>(values.data()), lens, inputs_shape);
|
|
MS_EXCEPTION_IF_NULL(indices);
|
|
return new SparseAdamOptimInfo(weight_addr, m, v, beta1_power, beta2_power, learning_rate, beta1, beta2, epsilon,
|
|
grad, indices, sharded);
|
|
}
|
|
|
|
OptimizerInfo *SparseFtrlOptimInfoBuilder::BuildInputs(const WeightPtr &weight, const Keys &, const Values &values,
|
|
const Lengths &lens, const InputsShapePtr &inputs_shape, size_t,
|
|
const std::shared_ptr<PServerKernel> &pserver_kernel,
|
|
bool sharded) {
|
|
MS_EXCEPTION_IF_NULL(inputs_shape);
|
|
AddressPtr weight_addr = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(weight_addr);
|
|
weight_addr->addr = weight->data();
|
|
weight_addr->size = weight->size() * sizeof(float);
|
|
|
|
AddressPtr accum = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(accum);
|
|
|
|
accum->addr = new float[weight->size()];
|
|
MS_EXCEPTION_IF_NULL(accum->addr);
|
|
accum->size = weight->size() * sizeof(float);
|
|
for (size_t i = 0; i < weight->size(); i++) {
|
|
float *tmp = reinterpret_cast<float *>(accum->addr);
|
|
tmp[i] = std::dynamic_pointer_cast<SparseApplyFtrlPSKernelMod>(pserver_kernel)->init_accum();
|
|
}
|
|
|
|
AddressPtr linear = std::make_shared<kernel::Address>();
|
|
MS_EXCEPTION_IF_NULL(linear);
|
|
|
|
linear->addr = new float[weight->size()];
|
|
MS_EXCEPTION_IF_NULL(linear->addr);
|
|
linear->size = weight->size() * sizeof(float);
|
|
int64_t ret = memset_s(linear->addr, weight->size() * sizeof(float), 0x00, weight->size() * sizeof(float));
|
|
if (ret != 0) {
|
|
MS_LOG(EXCEPTION) << "memset_s error, errorno(" << ret << ")";
|
|
delete[] reinterpret_cast<float *>(linear->addr);
|
|
linear->addr = nullptr;
|
|
return nullptr;
|
|
}
|
|
|
|
AddressPtr grad = GenInputAddrPtr<float>(kSparseFtrl, "grad", const_cast<float *>(values.data()), lens, inputs_shape);
|
|
MS_EXCEPTION_IF_NULL(grad);
|
|
AddressPtr indices =
|
|
GenInputAddrPtr<float>(kSparseFtrl, "indices", const_cast<float *>(values.data()), lens, inputs_shape);
|
|
MS_EXCEPTION_IF_NULL(indices);
|
|
return new SparseFtrlOptimInfo(weight_addr, accum, linear, grad, indices, sharded);
|
|
}
|
|
} // namespace ps
|
|
} // namespace mindspore
|