cristoval
5 years ago
committed by
ZPaC
parent
6ea74a3669
commit
c6262111ef
@ -0,0 +1,104 @@
|
||||
/**
|
||||
* 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 "backend/kernel_compiler/cpu/ps/sparse_apply_lazy_adam_ps_kernel.h"
|
||||
#include <memory>
|
||||
#include "backend/kernel_compiler/common_utils.h"
|
||||
#include "runtime/device/cpu/cpu_device_address.h"
|
||||
#include "frontend/parallel/ps/util.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
namespace ps {
|
||||
void SparseApplyLazyAdamPSKernel::InitKernel(
|
||||
const std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> &shapes) {
|
||||
const std::vector<std::shared_ptr<std::vector<size_t>>> &shape_vec = *shapes;
|
||||
std::vector<size_t> &var_shape = *(shape_vec[0]);
|
||||
std::vector<size_t> &m_shape = *(shape_vec[1]);
|
||||
std::vector<size_t> &v_shape = *(shape_vec[2]);
|
||||
const std::vector<size_t> &grad_shape = *(shape_vec[9]);
|
||||
const std::vector<size_t> &indices_shape = *(shape_vec[10]);
|
||||
|
||||
Shard(&var_shape, 0);
|
||||
Shard(&m_shape, 0);
|
||||
Shard(&v_shape, 0);
|
||||
|
||||
if (!IsSameShape(var_shape, m_shape)) {
|
||||
MS_LOG(EXCEPTION) << "var and m should have the same shape";
|
||||
}
|
||||
if (!IsSameShape(var_shape, v_shape)) {
|
||||
MS_LOG(EXCEPTION) << "var and v should have the same shape";
|
||||
}
|
||||
var_first_dim_size_ = var_shape[0];
|
||||
for (size_t i = 1; i < var_shape.size(); ++i) {
|
||||
if (var_shape[i] != grad_shape[i]) {
|
||||
MS_LOG(EXCEPTION) << "The shape of var and grad must equal in dimension " << i;
|
||||
}
|
||||
var_outer_dim_size_ *= var_shape[i];
|
||||
}
|
||||
if (indices_shape.size() != 1) {
|
||||
MS_LOG(EXCEPTION) << "indices must be 1D";
|
||||
}
|
||||
indices_size_ = indices_shape[0];
|
||||
if (grad_shape[0] != indices_size_) {
|
||||
MS_LOG(ERROR) << "The first dimension of grad shape must be equal to indices";
|
||||
}
|
||||
/*
|
||||
if (AnfAlgo::HasNodeAttr(USE_NESTEROV, kernel_node)) {
|
||||
use_nesterov_ = AnfAlgo::GetNodeAttr<bool>(kernel_node, "use_nesterov");
|
||||
}
|
||||
*/
|
||||
workspace_size_list_.emplace_back(indices_size_ * var_outer_dim_size_ * sizeof(float));
|
||||
workspace_size_list_.emplace_back(indices_size_ * sizeof(int));
|
||||
workspace_size_list_.emplace_back(indices_size_ * var_outer_dim_size_ * sizeof(float));
|
||||
workspace_size_list_.emplace_back(indices_size_ * sizeof(int));
|
||||
workspace_size_list_.emplace_back(var_first_dim_size_ * var_outer_dim_size_ * sizeof(float));
|
||||
}
|
||||
|
||||
void SparseApplyLazyAdamPSKernel::ReInit(
|
||||
const std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> &shapes) {
|
||||
const std::vector<std::shared_ptr<std::vector<size_t>>> &shape_vec = *shapes;
|
||||
const std::vector<size_t> &indices_shape = *(shape_vec[0]);
|
||||
indices_size_ = indices_shape[0];
|
||||
workspace_size_list_[0] = indices_size_ * var_outer_dim_size_ * sizeof(float);
|
||||
workspace_size_list_[1] = indices_size_ * sizeof(int);
|
||||
}
|
||||
|
||||
void SparseApplyLazyAdamPSKernel::ReInit(const std::vector<AddressPtr> &inputs) {
|
||||
const auto &indices_addr = inputs[10];
|
||||
indices_size_ = indices_addr->size / sizeof(int);
|
||||
workspace_size_list_[0] = indices_size_ * var_outer_dim_size_ * sizeof(float);
|
||||
workspace_size_list_[1] = indices_size_ * sizeof(int);
|
||||
}
|
||||
|
||||
bool SparseApplyLazyAdamPSKernel::Execute(const std::vector<AddressPtr> &inputs,
|
||||
const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) {
|
||||
ReInit(inputs);
|
||||
int *indices = reinterpret_cast<int *>(inputs[10]->addr);
|
||||
for (size_t i = 0; i < inputs[10]->size / sizeof(int); i++) {
|
||||
indices[i] -= rank_id_ * var_first_dim_size_;
|
||||
}
|
||||
return Launch(inputs, workspace, outputs);
|
||||
}
|
||||
|
||||
const std::vector<size_t> &SparseApplyLazyAdamPSKernel::input_sizes() const { return GetInputSizeList(); }
|
||||
|
||||
const std::vector<size_t> &SparseApplyLazyAdamPSKernel::output_sizes() const { return GetOutputSizeList(); }
|
||||
|
||||
const std::vector<size_t> &SparseApplyLazyAdamPSKernel::workspace_sizes() const { return GetWorkspaceSizeList(); }
|
||||
} // namespace ps
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,49 @@
|
||||
/**
|
||||
* 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_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_APPLY_LAZY_ADAM_PS_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_APPLY_LAZY_ADAM_PS_KERNEL_H_
|
||||
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
#include "backend/kernel_compiler/cpu/ps/pserver_kernel.h"
|
||||
#include "backend/kernel_compiler/cpu/sparse_apply_lazy_adam_cpu_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
namespace ps {
|
||||
using mindspore::kernel::SparseApplyLazyAdamCPUKernel;
|
||||
class SparseApplyLazyAdamPSKernel : public SparseApplyLazyAdamCPUKernel, public PServerKernel {
|
||||
public:
|
||||
SparseApplyLazyAdamPSKernel(size_t rank_id, size_t pserver_num) : PServerKernel(rank_id, pserver_num) {}
|
||||
~SparseApplyLazyAdamPSKernel() override = default;
|
||||
|
||||
void InitKernel(const std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> &) override;
|
||||
void ReInit(const std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> &) override;
|
||||
bool Execute(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) override;
|
||||
|
||||
const std::vector<size_t> &input_sizes() const override;
|
||||
const std::vector<size_t> &output_sizes() const override;
|
||||
const std::vector<size_t> &workspace_sizes() const override;
|
||||
|
||||
protected:
|
||||
void ReInit(const std::vector<AddressPtr> &) override;
|
||||
};
|
||||
} // namespace ps
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_APPLY_LAZY_ADAM_PS_KERNEL_H_
|
||||
Loading…
Reference in new issue