Merge pull request !31717 from kisnwang/clean_code
This commit is contained in:
i-robot 2022-03-23 01:20:30 +00:00 committed by Gitee
commit 651b5ebc12
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
4 changed files with 38 additions and 26 deletions

View File

@ -28,17 +28,14 @@ const AnfNodePtr CustomOpConstInputToAttr::Process(const FuncGraphPtr &, const A
if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
return nullptr;
}
auto cnode = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(cnode);
if (!IsPrimitiveCNode(cnode, prim::kPrimCustom)) {
return nullptr;
}
auto primitive = common::AnfAlgo::GetCNodePrimitive(cnode);
MS_EXCEPTION_IF_NULL(primitive);
mindspore::HashSet<size_t> attr_indices;
GetCustomOpAttrIndex(primitive, &attr_indices);
GetCustomOpAttrIndex(common::AnfAlgo::GetCNodePrimitive(cnode), &attr_indices);
if (attr_indices.empty()) {
return nullptr;
}

View File

@ -24,6 +24,10 @@
namespace mindspore {
namespace kernel {
namespace {
constexpr size_t kNcdhwShapeSize = 5;
} // namespace
bool HostCheck::CheckValidDeviceShape(const AnfNodePtr &node) {
size_t real_input_num = common::AnfAlgo::GetInputTensorNum(node);
for (size_t i = 0; i < real_input_num; i++) {
@ -71,7 +75,7 @@ std::vector<int64_t> HostCheck::GetFinalInferShape(const AnfNodePtr &node, const
MS_LOG(DEBUG) << "Get Device Shape using a shape size is less than 4 ,should be Padding shape by Default firstly";
temp_shape = trans::PaddingShapeTo4dDefault(infer_shape);
}
if (infer_shape.size() != trans::kNcdhw && k3DFormatSet.find(format) != k3DFormatSet.end()) {
if (infer_shape.size() != kNcdhwShapeSize && k3DFormatSet.find(format) != k3DFormatSet.end()) {
temp_shape = trans::PaddingShapeTo5dDefault(infer_shape);
}
return temp_shape;

View File

@ -72,8 +72,8 @@ bool EmbeddingLookUpCommGradCpuKernelMod::Launch(const std::vector<kernel::Addre
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', memset failed. Error no: " << ret;
}
const std::vector<int> &rank_group = {0, 1, 2, 3, 4, 5, 6, 7};
size_t input_split_lens = input_size / split_num_ / sizeof(float_t);
size_t output_split_lens = output_size / split_num_ / sizeof(float_t);
size_t input_split_lens = (input_size / split_num_) / sizeof(float_t);
size_t output_split_lens = (output_size / split_num_) / sizeof(float_t);
for (size_t i = 0; i < split_num_; ++i) {
(void)MPIAllGather(input_addr + i * input_split_lens, output_addr + i * output_split_lens, rank_group,
input_split_lens);

View File

@ -24,6 +24,17 @@ namespace {
constexpr size_t kSparseApplyProximalAdagradInputsNum = 7;
constexpr size_t kSparseApplyProximalAdagradWorkspaceSize = 4;
constexpr char kKernelName[] = "SparseApplyProximalAdagrad";
constexpr size_t kVarIndex = 0;
constexpr size_t kAccIndex = 1;
constexpr size_t kLRIndex = 2;
constexpr size_t kL1Index = 3;
constexpr size_t kL2Index = 4;
constexpr size_t kGradIndex = 5;
constexpr size_t kIndicesIndex = 6;
constexpr size_t kWorkSpaceIndex0 = 0;
constexpr size_t kWorkSpaceIndex1 = 1;
constexpr size_t kWorkSpaceIndex2 = 2;
constexpr size_t kWorkSpaceIndex3 = 3;
template <typename T>
void ComputeProximalAdagrad(MultiThreadComputeParams<T> *input_params, size_t start, size_t end) {
@ -84,13 +95,13 @@ void SparseApplyProximalAdagradCpuKernelMod::InitInputOutputSize(const CNodePtr
void SparseApplyProximalAdagradCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
MS_EXCEPTION_IF_NULL(kernel_node);
kernel_name_ = common::AnfAlgo::GetCNodeName(kernel_node);
std::vector<size_t> var_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
std::vector<size_t> accum_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
std::vector<size_t> lr_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
std::vector<size_t> l1_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 3);
std::vector<size_t> l2_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 4);
std::vector<size_t> grad_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 5);
std::vector<size_t> indices_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 6);
std::vector<size_t> var_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kVarIndex);
std::vector<size_t> accum_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kAccIndex);
std::vector<size_t> lr_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kLRIndex);
std::vector<size_t> l1_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kL1Index);
std::vector<size_t> l2_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kL2Index);
std::vector<size_t> grad_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kGradIndex);
std::vector<size_t> indices_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, kIndicesIndex);
if (var_shape.empty()) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_
<< "', the dimension of 'var' should be at least 1-D, but got scalar or None.";
@ -142,23 +153,23 @@ void SparseApplyProximalAdagradCpuKernelMod::InitKernel(const CNodePtr &kernel_n
<< "', 'l2' should be a scalar,and dimension of 'l2' should be 0,but got the dimension of 'l2': "
<< Vector2Str(l2_shape);
}
indices_data_type_ = AnfAlgo::GetInputDeviceDataType(kernel_node, 6);
indices_data_type_ = AnfAlgo::GetInputDeviceDataType(kernel_node, kIndicesIndex);
}
template <typename T>
void SparseApplyProximalAdagradCpuKernelMod::LaunchKernel(const std::vector<kernel::AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &workspace) const {
auto var = reinterpret_cast<float *>(inputs[0]->addr);
auto accum = reinterpret_cast<float *>(inputs[1]->addr);
auto lr = reinterpret_cast<float *>(inputs[2]->addr)[0];
auto l1 = reinterpret_cast<float *>(inputs[3]->addr)[0];
auto l2 = reinterpret_cast<float *>(inputs[4]->addr)[0];
auto grad = reinterpret_cast<float *>(inputs[5]->addr);
auto indices = reinterpret_cast<T *>(inputs[6]->addr);
auto new_grad = reinterpret_cast<float *>(workspace[0]->addr);
auto new_indices = reinterpret_cast<T *>(workspace[1]->addr);
auto workspace_grad = reinterpret_cast<float *>(workspace[2]->addr);
auto workspace_indices = reinterpret_cast<T *>(workspace[3]->addr);
auto var = reinterpret_cast<float *>(inputs[kVarIndex]->addr);
auto accum = reinterpret_cast<float *>(inputs[kAccIndex]->addr);
auto lr = reinterpret_cast<float *>(inputs[kLRIndex]->addr)[0];
auto l1 = reinterpret_cast<float *>(inputs[kL1Index]->addr)[0];
auto l2 = reinterpret_cast<float *>(inputs[kL2Index]->addr)[0];
auto grad = reinterpret_cast<float *>(inputs[kGradIndex]->addr);
auto indices = reinterpret_cast<T *>(inputs[kIndicesIndex]->addr);
auto new_grad = reinterpret_cast<float *>(workspace[kWorkSpaceIndex0]->addr);
auto new_indices = reinterpret_cast<T *>(workspace[kWorkSpaceIndex1]->addr);
auto workspace_grad = reinterpret_cast<float *>(workspace[kWorkSpaceIndex2]->addr);
auto workspace_indices = reinterpret_cast<T *>(workspace[kWorkSpaceIndex3]->addr);
SparseGradient<T> unique_sparse_grad({new_grad, new_indices, indices_size_});
SparseGradient<T> workspace_sparse_grad({workspace_grad, workspace_indices, indices_size_});