ssk015
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!21180 alltoall and neighbor-exchange op 

Merge pull request !21180 from zhoufeng/all-to-all-op
This commit is contained in:
i-robot 2021-08-02 01:42:42 +00:00 committed by Gitee
parent f1c35128a3 2143241092
commit bb6d1a3019
23 changed files with 945 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/backend/kernel_compiler/hccl/hccl_kernel_metadata.cc
View File

@ -67,7 +67,7 @@ void HcclMetadataInfo(const CNodePtr &kernel_node, std::vector<std::shared_ptr<K
MS_EXCEPTION_IF_NULL(kernel_node);
std::string op_name = AnfAlgo::GetCNodeName(kernel_node);
if (op_name != kAllGather && op_name != kAllReduce && op_name != kBroadcast && op_name != kReduceScatter &&
op_name != kHcomSend && op_name != kReceive) {
op_name != kHcomSend && op_name != kReceive && op_name != kAllToAllv) {
MS_LOG(DEBUG) << "Hccl does not have op [" << op_name << "]";
return;
}

29
mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_all_to_all.cc Normal file
View File

@ -0,0 +1,29 @@
/**
* Copyright 2021 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/hccl/hcom_all_to_all.h"
namespace mindspore::kernel {
HcomAllToAllKernel::HcomAllToAllKernel() {}
HcomAllToAllKernel::~HcomAllToAllKernel() {}
bool HcomAllToAllKernel::Launch(const std::vector<AddressPtr> &, const std::vector<AddressPtr> &,
const std::vector<AddressPtr> &, void *) {
return true;
}
MS_HCCL_REG_KERNEL(AllToAllv, HcomAllToAllKernel);
} // namespace mindspore::kernel

33
mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_all_to_all.h Normal file
View File

@ -0,0 +1,33 @@
/**
* Copyright 2021 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_HCCL_HCOM_ALL_TO_ALL_H_
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HCCL_HCOM_ALL_TO_ALL_H_
#include <memory>
#include <vector>
#include "backend/kernel_compiler/hccl/hccl_kernel.h"
namespace mindspore::kernel {
class HcomAllToAllKernel : public HcclKernel {
public:
HcomAllToAllKernel();
~HcomAllToAllKernel() override;
bool Launch(const std::vector<AddressPtr> &, const std::vector<AddressPtr> &, const std::vector<AddressPtr> &,
void *) override;
};
} // namespace mindspore::kernel
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HCCL_HCOM_ALL_TO_ALL_H_

1
mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_util.h
View File

@ -37,6 +37,7 @@ constexpr auto kBroadcast = "Broadcast";
constexpr auto kHcomSend = "Send";
constexpr auto kReceive = "Receive";
constexpr auto kReduceScatter = "ReduceScatter";
constexpr auto kAllToAllv = "AllToAllv";
/* Correspondence between data_type and hcom data type in Ascend */
static map<int64_t, HcclDataType> CONST_OP_HCOM_DATA_TYPE_MAP = {

178
mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.cc Normal file
View File

@ -0,0 +1,178 @@
/**
* 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/optimizer/ascend/mindir/all_to_all_unify_mindir.h"
#include <vector>
#include <string>
#include "backend/session/anf_runtime_algorithm.h"
#include "runtime/hccl_adapter/hccl_adapter.h"
#include "backend/optimizer/common/helper.h"
namespace mindspore {
namespace opt {
namespace {
constexpr size_t kCNodePrimitiveIdx = 0;
constexpr size_t kAllToAllInputIdx = 1;
void ChangePrimitiveToAllToAllV(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
auto neighbor_exchange = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(neighbor_exchange);
if (neighbor_exchange->size() <= kCNodePrimitiveIdx) {
MS_LOG(EXCEPTION) << "Invalid cnode " << node->DebugString() << " input size " << neighbor_exchange->size();
}
auto prim = GetValueNode<PrimitivePtr>(neighbor_exchange->input(kCNodePrimitiveIdx));
MS_EXCEPTION_IF_NULL(prim);
prim->Named::operator=(Named(kAllToAllVOpName));
}
uint32_t GetRankSize(const std::string &group) {
uint32_t rank_size;
auto hccl_ret = hccl::HcclAdapter::GetInstance().HcclGetRankSize(group, &rank_size);
if (hccl_ret != ::HcclResult::HCCL_SUCCESS) {
MS_LOG(EXCEPTION) << "Get hccl rank size for group " << group << " failed, ret = " << hccl_ret;
}
return rank_size;
}
CNodePtr CreateSplitNode(const FuncGraphPtr &graph, const CNodePtr &all_to_all) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(all_to_all);
int64_t split_count = AnfAlgo::GetNodeAttr<int64_t>(all_to_all, kAttrSplitCount);
int64_t split_dim = AnfAlgo::GetNodeAttr<int64_t>(all_to_all, kAttrSplitDim);
if (all_to_all->size() <= kAllToAllInputIdx) {
MS_LOG(EXCEPTION) << "Invalid cnode " << all_to_all->DebugString() << " input size " << all_to_all->size();
}
auto all_to_all_input = all_to_all->input(kAllToAllInputIdx);
std::vector<AnfNodePtr> split_input = {NewValueNode(std::make_shared<Primitive>(prim::kPrimSplitV->name())),
all_to_all_input};
auto split_v = graph->NewCNode(split_input);
MS_EXCEPTION_IF_NULL(split_v);
auto dtype = AnfAlgo::GetOutputInferDataType(all_to_all_input, 0);
auto shape = AnfAlgo::GetOutputInferShape(all_to_all_input, 0);
if (SizeToLong(shape.size()) <= split_dim) {
MS_LOG(EXCEPTION) << "Invalid split dim " << split_dim << " is over the shape size " << shape.size();
}
if (shape[LongToSize(split_dim)] % split_count != 0) {
MS_LOG(EXCEPTION) << "Invalid split count " << split_count << " cannot be divisible by shape[" << split_dim
<< "] = " << shape[LongToSize(split_dim)];
}
shape[LongToSize(split_dim)] /= split_count;
std::vector<TypeId> dtypes(split_count, dtype);
std::vector<std::vector<size_t>> shapes(split_count, shape);
AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split_v.get());
AnfAlgo::SetNodeAttr(kAttrSplitDim, MakeValue<int64_t>(split_dim), split_v);
AnfAlgo::SetNodeAttr(kAttrNumSplit, MakeValue<int64_t>(split_count), split_v);
AnfAlgo::SetNodeAttr(kAttrSizeSplits, MakeValue(std::vector<int64_t>(split_count, shape[LongToSize(split_dim)])),
split_v);
AnfAlgo::SetNodeAttr("is_backend_insert", MakeValue(true), split_v);
return split_v;
}
CNodePtr CreateAllToAllvNode(const FuncGraphPtr &graph, const CNodePtr &all_to_all, const CNodePtr &split) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(all_to_all);
MS_EXCEPTION_IF_NULL(split);
int64_t split_count = AnfAlgo::GetNodeAttr<int64_t>(all_to_all, kAttrSplitCount);
std::string group = AnfAlgo::GetNodeAttr<std::string>(all_to_all, kAttrGroup);
std::vector<AnfNodePtr> split_outputs;
CreateMultipleOutputsOfAnfNode(graph, split, split_count, &split_outputs);
if (split_outputs.empty()) {
MS_LOG(EXCEPTION) << "Create tuple get item failed.";
}
std::vector<AnfNodePtr> all_to_all_v_input = {NewValueNode(std::make_shared<Primitive>(kAllToAllVOpName))};
all_to_all_v_input.insert(all_to_all_v_input.end(), split_outputs.begin(), split_outputs.end());
auto all_to_all_v = graph->NewCNode(all_to_all_v_input);
MS_EXCEPTION_IF_NULL(all_to_all_v);
auto single_shape = AnfAlgo::GetOutputInferShape(split_outputs[0], 0);
auto single_type = AnfAlgo::GetOutputInferDataType(split_outputs[0], 0);
std::vector<TypeId> dtypes(split_count, single_type);
std::vector<std::vector<size_t>> shapes(split_count, single_shape);
AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, all_to_all_v.get());
uint32_t rank_size = GetRankSize(group);
std::vector<int64_t> rank_ids(rank_size, 0);
for (uint32_t i = 0; i < rank_size; ++i) {
rank_ids[i] = static_cast<int64_t>(i);
}
AnfAlgo::SetNodeAttr(kAttrSendRankIds, MakeValue<std::vector<int64_t>>(rank_ids), all_to_all_v);
AnfAlgo::SetNodeAttr(kAttrRecvRankIds, MakeValue<std::vector<int64_t>>(rank_ids), all_to_all_v);
AnfAlgo::SetNodeAttr(kAttrGroup, MakeValue<std::string>(group), all_to_all_v);
return all_to_all_v;
}
CNodePtr CreateConcatNode(const FuncGraphPtr &graph, const CNodePtr &all_to_all, const CNodePtr &all_to_all_v) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(all_to_all);
MS_EXCEPTION_IF_NULL(all_to_all_v);
int64_t split_count = AnfAlgo::GetNodeAttr<int64_t>(all_to_all, kAttrSplitCount);
int64_t concat_dim = AnfAlgo::GetNodeAttr<int64_t>(all_to_all, kAttrConcatDim);
std::vector<AnfNodePtr> all_to_all_v_outputs;
CreateMultipleOutputsOfAnfNode(graph, all_to_all_v, split_count, &all_to_all_v_outputs);
if (all_to_all_v_outputs.empty()) {
MS_LOG(EXCEPTION) << "Create tuple get item failed.";
}
std::vector<AnfNodePtr> concat_input = {NewValueNode(std::make_shared<Primitive>(kConcatOpName))};
concat_input.insert(concat_input.end(), all_to_all_v_outputs.begin(), all_to_all_v_outputs.end());
auto concat = graph->NewCNode(concat_input);
MS_EXCEPTION_IF_NULL(concat);
auto single_shape = AnfAlgo::GetOutputInferShape(all_to_all_v_outputs[0], 0);
if (LongToSize(concat_dim) >= single_shape.size()) {
MS_LOG(EXCEPTION) << "Invalid concat dim " << concat_dim << " is greater than shape size " << single_shape.size();
}
single_shape[LongToSize(concat_dim)] *= split_count;
AnfAlgo::SetOutputInferTypeAndShape({AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)}, {single_shape},
concat.get());
AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue<int64_t>(concat_dim), concat);
AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(split_count), concat);
std::vector<int64_t> dyn_input_size{split_count};
AnfAlgo::SetNodeAttr(kAttrDynInputSizes, MakeValue(dyn_input_size), concat);
return concat;
}
} // namespace
const BaseRef NeighborExchangeUnifyMindIR::DefinePattern() const {
return VectorRef({prim::kPrimNeighborExchange, std::make_shared<SeqVar>()});
}
const AnfNodePtr NeighborExchangeUnifyMindIR::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(node);
ChangePrimitiveToAllToAllV(node);
return node;
}
const BaseRef AllToAllUnifyMindIR::DefinePattern() const {
return VectorRef({prim::kPrimAllToAll, std::make_shared<SeqVar>()});
}
const AnfNodePtr AllToAllUnifyMindIR::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(node);
auto all_to_all = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(all_to_all);
auto split = CreateSplitNode(graph, all_to_all);
auto all_to_all_v = CreateAllToAllvNode(graph, all_to_all, split);
auto concat = CreateConcatNode(graph, all_to_all, all_to_all_v);
return concat;
}
} // namespace opt
} // namespace mindspore

42
mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.h Normal file
View File

@ -0,0 +1,42 @@
/**
* 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_OPTIMIZER_ASCEND_MINDIR_ALL_TO_ALL_UNIFY_MINDIR_H_
#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_ALL_TO_ALL_UNIFY_MINDIR_H_
#include <memory>
#include "backend/optimizer/common/optimizer.h"
namespace mindspore {
namespace opt {
class NeighborExchangeUnifyMindIR : public PatternProcessPass {
public:
explicit NeighborExchangeUnifyMindIR(bool multigraph = true)
: PatternProcessPass("neighbor_exchange_unify_mindir", multigraph) {}
~NeighborExchangeUnifyMindIR() override = default;
const BaseRef DefinePattern() const override;
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
};
class AllToAllUnifyMindIR : public PatternProcessPass {
public:
explicit AllToAllUnifyMindIR(bool multigraph = true) : PatternProcessPass("all_to_all_unify_mindir", multigraph) {}
~AllToAllUnifyMindIR() override = default;
const BaseRef DefinePattern() const override;
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
};
} // namespace opt
} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_ALL_TO_ALL_UNIFY_MINDIR_H_

3
mindspore/ccsrc/backend/session/ascend_session.cc
View File

@ -44,6 +44,7 @@
#include "backend/optimizer/ascend/mindir/slice_grad_unify_mindir.h"
#include "backend/optimizer/ascend/mindir/avg_pool_grad_unify_mindir.h"
#include "backend/optimizer/ascend/mindir/bn_grad_unify_mindir.h"
#include "backend/optimizer/ascend/mindir/all_to_all_unify_mindir.h"
#include "runtime/device/kernel_adjust.h"
#include "runtime/device/ascend/ascend_stream_assign.h"
#include "backend/session/anf_runtime_algorithm.h"
@ -316,6 +317,8 @@ void AscendSession::UnifyMindIR(const KernelGraphPtr &graph) {
unify_mindir_pm->AddPass(std::make_shared<opt::DropoutUnifyMindIR1>());
unify_mindir_pm->AddPass(std::make_shared<opt::DropoutGradUnifyMindIR>());
unify_mindir_pm->AddPass(std::make_shared<opt::BatchNormGradUnifyMindIR>());
unify_mindir_pm->AddPass(std::make_shared<opt::NeighborExchangeUnifyMindIR>());
unify_mindir_pm->AddPass(std::make_shared<opt::AllToAllUnifyMindIR>());
optimizer->AddPassManager(unify_mindir_pm);
(void)optimizer->Optimize(graph);

2
mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.h
View File

@ -32,7 +32,7 @@ class AscendMemoryManager : public MemoryManager {
void ClearGlobalIdleMem() override;
void *MallocMemFromMemPool(size_t size) override;
uint64_t GetDeviceMemSize();
void MallocSomasDynamicMem(const session::KernelGraph *graph);
void MallocSomasDynamicMem(const session::KernelGraph *graph) override;
uint8_t *MallocCommunicationMemFromMemPool(size_t size) override;
protected:

5
mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc
View File

@ -1389,7 +1389,6 @@ vector<CNodePtr> AscendStreamAssign::GetLastInputCnode(const NotNull<KernelGraph
}
vector<CNodePtr> final_inputs;
const uint32_t max = 0;
CNodePtr max_common_cnode = nullptr;
for (const auto &item : result) {
if (IsHcom(item.second.first)) {
@ -1400,9 +1399,7 @@ vector<CNodePtr> AscendStreamAssign::GetLastInputCnode(const NotNull<KernelGraph
final_inputs.emplace_back(item.second.first);
}
} else {
if (item.second.second > max) {
max_common_cnode = item.second.first;
}
max_common_cnode = item.second.first;
}
}

2
mindspore/ccsrc/runtime/device/memory_manager.cc
View File

@ -30,7 +30,7 @@ size_t MemoryManager::GetCommonAlignSize(size_t input_size) {
return (input_size + kMemAlignSize + kAlignBytes - 1) / kMemAlignSize * kMemAlignSize;
}
size_t MemoryManager::GetCommunicationAlignSize(size_t input_size) const {
size_t MemoryManager::GetCommunicationAlignSize(size_t input_size) {
return (input_size + kMemAlignSize - 1) / kMemAlignSize * kMemAlignSize + 2 * kMemAlignSize;
}

2
mindspore/ccsrc/runtime/device/memory_manager.h
View File

@ -58,7 +58,7 @@ class MemoryManager {
virtual std::vector<void *> MallocContinuousMemFromMemPool(size_t total_size, std::vector<size_t> size_list);
static size_t GetCommonAlignSize(size_t input_size);
size_t GetCommunicationAlignSize(size_t input_size) const;
static size_t GetCommunicationAlignSize(size_t input_size);
protected:
virtual uint8_t *MallocStaticMem(size_t size, bool communication_mem, uint32_t graph_id = kInvalidGraphId) = 0;

119
mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.cc Normal file
View File

@ -0,0 +1,119 @@
/**
* Copyright 2021 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 "runtime/hccl_adapter/all_to_all_v_calc_param.h"
#include <functional>
#include <map>
#include <string>
#include "backend/session/anf_runtime_algorithm.h"
#include "transform/graph_ir/util.h"
#include "runtime/device/memory_manager.h"
namespace mindspore::hccl {
namespace {
bool IsInTheOrder(const std::vector<int64_t> &vec) {
for (size_t i = 1; i < vec.size(); ++i) {
if (vec[i] <= vec[i - 1]) {
return false;
}
}
return true;
}
} // namespace
AllToAllvCalcParam::AllToAllvCalcParam(const CNodeWeakPtr &cnode, uint32_t rank_size)
: node_(cnode),
rank_size_(rank_size),
send_counts_(rank_size, 0),
sdispls_(rank_size, 0),
recv_counts_(rank_size, 0),
rdispls_(rank_size, 0) {}
void AllToAllvCalcParam::CalcOpParam() {
CNodePtr cnode = node_.lock();
MS_EXCEPTION_IF_NULL(cnode);
size_t input_num = AnfAlgo::GetInputTensorNum(cnode);
size_t output_num = AnfAlgo::GetOutputTensorNum(cnode);
std::vector<size_t> input_aligned_mem_size(input_num);
std::vector<size_t> output_aligned_mem_size(output_num);
std::vector<size_t> input_real_mem_size(input_num);
std::vector<size_t> output_real_mem_size(output_num);
for (size_t i = 0; i < input_num; ++i) {
auto ms_shape = AnfAlgo::GetInputDeviceShape(cnode, i);
auto type_size = transform::TransformUtil::GetDataTypeSize(AnfAlgo::GetInputDeviceDataType(cnode, i));
size_t origin_mem_size = std::accumulate(ms_shape.begin(), ms_shape.end(), type_size, std::multiplies<size_t>());
size_t aligned_mem_size = device::MemoryManager::GetCommonAlignSize(origin_mem_size);
input_aligned_mem_size[i] = aligned_mem_size / type_size;
input_real_mem_size[i] = origin_mem_size / type_size;
}
for (size_t i = 0; i < output_num; ++i) {
auto ms_shape = AnfAlgo::GetOutputDeviceShape(cnode, i);
auto type_size = transform::TransformUtil::GetDataTypeSize(AnfAlgo::GetOutputDeviceDataType(cnode, i));
size_t origin_mem_size = std::accumulate(ms_shape.begin(), ms_shape.end(), type_size, std::multiplies<size_t>());
size_t aligned_mem_size = device::MemoryManager::GetCommonAlignSize(origin_mem_size);
output_aligned_mem_size[i] = aligned_mem_size / type_size;
output_real_mem_size[i] = origin_mem_size / type_size;
}
CalcMemOffset(input_aligned_mem_size, input_real_mem_size, kAttrSendRankIds, &send_counts_, &sdispls_);
CalcMemOffset(output_aligned_mem_size, output_real_mem_size, kAttrRecvRankIds, &recv_counts_, &rdispls_);
}
void AllToAllvCalcParam::CalcMemOffset(const std::vector<size_t> &mem_sizes, const std::vector<size_t> &real_sizes,
const std::string &rank_ids_attr, std::vector<int64_t> *counts,
std::vector<int64_t> *displs) {
CNodePtr cnode = node_.lock();
MS_EXCEPTION_IF_NULL(cnode);
auto rank_ids = AnfAlgo::GetNodeAttr<std::vector<int64_t>>(cnode, rank_ids_attr);
if (mem_sizes.size() != rank_ids.size() || real_sizes.size() != rank_ids.size()) {
MS_LOG(EXCEPTION) << "Invalid addr num " << mem_sizes.size() << " and " << real_sizes.size()
<< " must be equal to rank ids size " << rank_ids.size();
}
if (!IsInTheOrder(rank_ids)) {
std::vector<size_t> mem_offset(mem_sizes.size(), 0);
for (size_t i = 1; i < mem_sizes.size(); ++i) {
mem_offset[i] = mem_offset[i - 1] + mem_sizes[i];
}
for (size_t i = 0; i < rank_ids.size(); ++i) {
if (rank_ids[i] < 0 || static_cast<size_t>(rank_ids[i]) >= rank_size_) {
MS_LOG(EXCEPTION) << "Invalid rank id " << rank_ids[i] << " at index " << i << " as rank size " << rank_size_;
}
(*counts)[rank_ids[i]] = real_sizes[i];
(*displs)[rank_ids[i]] = mem_offset[i];
}
return;
}
std::map<int64_t, size_t> rank_id_map;
for (size_t i = 0; i < rank_ids.size(); ++i) {
if (rank_ids[i] < 0 || static_cast<size_t>(rank_ids[i]) >= rank_size_) {
MS_LOG(EXCEPTION) << "Invalid rank id " << rank_ids[i] << " at index " << i << " as rank size " << rank_size_;
}
rank_id_map.emplace(rank_ids[i], i);
}
size_t offset = 0;
for (uint32_t i = 0; i < rank_size_; ++i) {
(*displs)[i] = offset;
auto iter = rank_id_map.find(i);
if (iter != rank_id_map.end()) {
(*counts)[i] = real_sizes[iter->second];
offset += mem_sizes[iter->second];
} else {
(*counts)[i] = 0;
}
}
}
} // namespace mindspore::hccl

47
mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.h Normal file
View File

@ -0,0 +1,47 @@
/**
* Copyright 2021 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_RUNTIME_HCCL_ADAPTER_ALL_TO_ALL_V_CALC_PARAM_H
#define MINDSPORE_RUNTIME_HCCL_ADAPTER_ALL_TO_ALL_V_CALC_PARAM_H
#include <memory>
#include <vector>
#include <string>
#include "mindspore/core/ir/anf.h"
namespace mindspore::hccl {
class AllToAllvCalcParam {
public:
AllToAllvCalcParam(const CNodeWeakPtr &cnode, uint32_t rank_size);
~AllToAllvCalcParam() = default;
void CalcOpParam();
const std::vector<int64_t> &GetSendCounts() const { return send_counts_; }
const std::vector<int64_t> &GetSendDispls() const { return sdispls_; }
const std::vector<int64_t> &GetRecvCounts() const { return recv_counts_; }
const std::vector<int64_t> &GetRecvDispls() const { return rdispls_; }
private:
void CalcMemOffset(const std::vector<size_t> &mem_sizes, const std::vector<size_t> &real_sizes,
const std::string &rank_ids_attr, std::vector<int64_t> *counts, std::vector<int64_t> *displs);
CNodeWeakPtr node_;
uint32_t rank_size_;
std::vector<int64_t> send_counts_;
std::vector<int64_t> sdispls_;
std::vector<int64_t> recv_counts_;
std::vector<int64_t> rdispls_;
};
} // namespace mindspore::hccl
#endif // MINDSPORE_RUNTIME_HCCL_ADAPTER_ALL_TO_ALL_V_CALC_PARAM_H

60
mindspore/ccsrc/runtime/hccl_adapter/converter.cc
View File

@ -24,9 +24,9 @@
#undef google
#include "backend/session/anf_runtime_algorithm.h"
#include "utils/log_adapter.h"
#include "utils/ms_utils.h"
#include "mindspore/core/base/core_ops.h"
#include "transform/graph_ir/util.h"
#include "runtime/hccl_adapter/all_to_all_v_calc_param.h"
static constexpr char kGeOpNameHcclSend[] = "HcomSend";
static constexpr char kGeOpNameHcclReceive[] = "HcomReceive";
@ -34,7 +34,12 @@ static constexpr char kGeOpNameHcclAllRudece[] = "HcomAllReduce";
static constexpr char kGeOpNameHcclAllGather[] = "HcomAllGather";
static constexpr char kGeOpNameHcclBroadcast[] = "HcomBroadcast";
static constexpr char kGeOpNameHcclReduceScatter[] = "HcomReduceScatter";
static constexpr char kGeOpNameHcclAllToAllV[] = "HcomAllToAllV";
static constexpr char kGeNodeAttrUsedStreamNum[] = "used_stream_num";
static constexpr char kGeNodeAttrSendCounts[] = "send_counts";
static constexpr char kGeNodeAttrSendDispls[] = "send_displacements";
static constexpr char kGeNodeAttrRecvCounts[] = "recv_counts";
static constexpr char kGeNodeAttrRecvDispls[] = "recv_displacements";
static constexpr char kStubDataStructureName[] = "any_name_can_work";
static ge::DataType ConvertHcclDTypeToGeDType(HcclDataType datatype) {
@ -105,6 +110,37 @@ static T ConvertAttr(const CNodePtr &cnode, const ge::OpDescPtr &ge_op, const st
return attr;
}
static void SetGeNodeInt64VecAttr(const ge::OpDescPtr &ge_op, const std::string &ge_attr_name,
const std::vector<int64_t> &attr_value) {
MS_EXCEPTION_IF_NULL(ge_op);
for (size_t i = 0; i < attr_value.size(); ++i) {
MS_LOG(INFO) << ge_attr_name << " " << i << " = " << attr_value[i];
}
auto ret = ge::AttrUtils::SetListInt(*ge_op, ge_attr_name, attr_value);
if (!ret) {
MS_LOG(EXCEPTION) << "Set attr " << ge_attr_name << " for ge node failed.";
}
}
static void SetAllToAllvAttr(const CNodePtr &cnode, const ge::OpDescPtr &ge_op, const std::string &group) {
MS_EXCEPTION_IF_NULL(cnode);
MS_EXCEPTION_IF_NULL(ge_op);
if (!IsPrimitiveCNode(cnode, prim::kPrimAllToAllv)) {
return;
}
uint32_t rank_size = 0;
::HcclResult hccl_ret = hccl::HcclAdapter::GetInstance().HcclGetRankSize(group, &rank_size);
if (hccl_ret != ::HcclResult::HCCL_SUCCESS) {
MS_LOG(EXCEPTION) << "Get hccl rank size for group " << group << " failed, ret = " << hccl_ret;
}
mindspore::hccl::AllToAllvCalcParam calc(cnode, rank_size);
calc.CalcOpParam();
SetGeNodeInt64VecAttr(ge_op, kGeNodeAttrSendCounts, calc.GetSendCounts());
SetGeNodeInt64VecAttr(ge_op, kGeNodeAttrSendDispls, calc.GetSendDispls());
SetGeNodeInt64VecAttr(ge_op, kGeNodeAttrRecvCounts, calc.GetRecvCounts());
SetGeNodeInt64VecAttr(ge_op, kGeNodeAttrRecvDispls, calc.GetRecvDispls());
}
std::string GetGeNodeName(const CNodePtr &cnode) {
MS_EXCEPTION_IF_NULL(cnode);
if (IsPrimitiveCNode(cnode, prim::kPrimAllReduce)) {
@ -119,6 +155,8 @@ std::string GetGeNodeName(const CNodePtr &cnode) {
return kGeOpNameHcclSend;
} else if (IsPrimitiveCNode(cnode, prim::kPrimReceive)) {
return kGeOpNameHcclReceive;
} else if (IsPrimitiveCNode(cnode, prim::kPrimAllToAllv)) {
return kGeOpNameHcclAllToAllV;
}
MS_LOG(EXCEPTION) << "Unknown hccl node type " << cnode->DebugString();
@ -132,14 +170,25 @@ std::tuple<ge::NodePtr, ge::ComputeGraphPtr> GenerateStubGeNode(const AnfNodePtr
ge::OpDescPtr op_desc = std::make_shared<ge::OpDesc>(kStubDataStructureName, ge_node_name);
MS_EXCEPTION_IF_NULL(op_desc);
for (size_t i = 1; i < cnode->size(); ++i) {
size_t input_num = AnfAlgo::GetInputTensorNum(cnode);
size_t output_num = AnfAlgo::GetOutputTensorNum(cnode);
for (size_t i = 0; i < input_num; ++i) {
std::vector<int64_t> ge_shape;
auto ms_shape = AnfAlgo::GetInputDeviceShape(cnode, i - 1);
auto ms_shape = AnfAlgo::GetInputDeviceShape(cnode, i);
std::transform(ms_shape.begin(), ms_shape.end(), std::back_inserter(ge_shape),
[](size_t in) { return static_cast<int64_t>(in); });
op_desc->AddInputDesc(
ge::GeTensorDesc(ge::GeShape(ge_shape), ge::Format::FORMAT_NCHW,
transform::TransformUtil::ConvertDataType(AnfAlgo::GetInputDeviceDataType(cnode, i - 1))));
transform::TransformUtil::ConvertDataType(AnfAlgo::GetInputDeviceDataType(cnode, i))));
}
for (size_t i = 0; i < output_num; ++i) {
std::vector<int64_t> ge_shape;
auto ms_shape = AnfAlgo::GetOutputDeviceShape(cnode, i);
std::transform(ms_shape.begin(), ms_shape.end(), std::back_inserter(ge_shape),
[](size_t in) { return static_cast<int64_t>(in); });
op_desc->AddOutputDesc(
ge::GeTensorDesc(ge::GeShape(ge_shape), ge::Format::FORMAT_NCHW,
transform::TransformUtil::ConvertDataType(AnfAlgo::GetOutputDeviceDataType(cnode, i))));
}
// set node data type
@ -151,11 +200,12 @@ std::tuple<ge::NodePtr, ge::ComputeGraphPtr> GenerateStubGeNode(const AnfNodePtr
// set node attr
(void)ConvertAttr<int64_t>(cnode, op_desc, kAttrRankSize, ge::HCOM_ATTR_RANK_SIZE);
(void)ConvertAttr<std::string>(cnode, op_desc, kAttrGroup, ge::HCOM_ATTR_GROUP);
auto group = ConvertAttr<std::string>(cnode, op_desc, kAttrGroup, ge::HCOM_ATTR_GROUP);
(void)ConvertAttr<int64_t>(cnode, op_desc, kAttrSrcRank, ge::HCOM_ATTR_SRC_RANK);
(void)ConvertAttr<int64_t>(cnode, op_desc, kAttrDestRank, ge::HCOM_ATTR_DEST_RANK);
(void)ConvertAttr<int64_t>(cnode, op_desc, kAttrSrTag, ge::HCOM_ATTR_SR_TAG);
(void)ConvertAttr<std::vector<int64_t>>(cnode, op_desc, kAttrShape, ge::HCOM_ATTR_SHAPE);
SetAllToAllvAttr(cnode, op_desc, group);
ge::ComputeGraphPtr ge_graph = std::make_shared<ge::ComputeGraph>(kStubDataStructureName);
MS_EXCEPTION_IF_NULL(ge_graph);

7
mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.cc
View File

@ -91,6 +91,7 @@ void HcclAdapter::InitPlugin() {
hccl_exec_initialize_ = DlsymFuncObj(HcomExecInitialize, plugin_handle_);
hccl_exec_finalize_ = DlsymFuncObj(HcomExecFinalize, plugin_handle_);
hccl_exec_enqueue_op_ = DlsymFuncObj(HcomExecEnqueueOperation, plugin_handle_);
hccl_exec_enqueue_all_to_all_v_ = DlsymFuncObj(HcomExecEnqueueAllToAllV, plugin_handle_);
}
void HcclAdapter::FinalizePlugin() {
@ -113,6 +114,7 @@ void HcclAdapter::FinalizePlugin() {
hccl_exec_initialize_ = nullptr;
hccl_exec_finalize_ = nullptr;
hccl_exec_enqueue_op_ = nullptr;
hccl_exec_enqueue_all_to_all_v_ = nullptr;
(void)dlclose(plugin_handle_);
plugin_handle_ = nullptr;
}
@ -403,4 +405,9 @@ HcclResult HcclAdapter::HcclExecEnqueueOp(const ::HcomOperation &op_info, const
MS_EXCEPTION_IF_NULL(hccl_exec_enqueue_op_);
return hccl_exec_enqueue_op_(op_info, callback);
}
HcclResult HcclAdapter::HcclExecAllToAllv(const ::HcomAllToAllVParams &params, const HExecCallBack &callback) const {
MS_EXCEPTION_IF_NULL(hccl_exec_enqueue_all_to_all_v_);
return hccl_exec_enqueue_all_to_all_v_(params, callback);
}
} // namespace mindspore::hccl

2
mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.h
View File

@ -63,6 +63,7 @@ class HcclAdapter {
// for enqueue op
HcclResult HcclExecEnqueueOp(const ::HcomOperation &op_info, const HExecCallBack &callback) const;
HcclResult HcclExecAllToAllv(const ::HcomAllToAllVParams &params, const HExecCallBack &callback) const;
private:
HcclAdapter() = default;
@ -99,6 +100,7 @@ class HcclAdapter {
HcomExecInitializeFunObj hccl_exec_initialize_ = nullptr;
HcomExecFinalizeFunObj hccl_exec_finalize_ = nullptr;
HcomExecEnqueueOperationFunObj hccl_exec_enqueue_op_ = nullptr;
HcomExecEnqueueAllToAllVFunObj hccl_exec_enqueue_all_to_all_v_ = nullptr;
HcclComm hccl_comm_ = nullptr;

2
mindspore/ccsrc/runtime/hccl_adapter/plugin/hccl_plugin.h
View File

@ -22,6 +22,7 @@
#include <functional>
#include "external/ge/ge_api_types.h"
#include "hccl/hccl.h"
#include "hccl/hcom.h"
#include "utils/dlopen_macro.h"
constexpr const char *kHcclOpsKernelInfoStore = "ops_kernel_info_hccl";
@ -55,4 +56,5 @@ ORIGIN_METHOD(HcomGetRankSize, HcclResult, const char *, uint32_t *);
ORIGIN_METHOD(HcomExecInitialize, HcclResult);
ORIGIN_METHOD(HcomExecFinalize, HcclResult);
ORIGIN_METHOD(HcomExecEnqueueOperation, HcclResult, ::HcomOperation, HExecCallBack);
ORIGIN_METHOD(HcomExecEnqueueAllToAllV, HcclResult, ::HcomAllToAllVParams, HExecCallBack);
#endif // MINDSPORE_RUNTIME_HCCL_ADAPTER_PLUGIN_HCCL_PLUGIN_H

6
mindspore/ccsrc/utils/utils.h
View File

@ -63,6 +63,7 @@ constexpr auto kAtomicAddrCleanOpName = "AtomicAddrClean";
constexpr auto kGetNextOpName = "GetNext";
constexpr auto kInitDatasetQueueOpName = "InitDataSetQueue";
constexpr auto kEndOfSequence = "EndOfSequence";
constexpr auto kAllToAllVOpName = "AllToAllv";
constexpr auto kAllReduceOpName = "AllReduce";
constexpr auto kAllGatherOpName = "AllGather";
constexpr auto kHostAllGatherOpName = "HostAllGather";
@ -407,6 +408,11 @@ constexpr auto kAttrChildGraph = "child_graph";
constexpr auto kAttrInputNums = "inputNums";
constexpr auto kAttrT = "T";
constexpr auto kAttrNum = "num";
constexpr auto kAttrRecvType = "recv_type";
constexpr auto kAttrConcatDim = "concat_dim";
constexpr auto kAttrSplitCount = "split_count";
constexpr auto kAttrSendRankIds = "send_rank_ids";
constexpr auto kAttrRecvRankIds = "recv_rank_ids";
constexpr auto kAttrRankSize = "rank_size";
constexpr auto kAttrPadDimSize = "pad_dim_size";
constexpr auto kAttrPaddings = "paddings";

2
mindspore/core/base/core_ops.h
View File

@ -389,6 +389,8 @@ inline const PrimitivePtr kPrimSend = std::make_shared<Primitive>("Send");
inline const PrimitivePtr kPrimReceive = std::make_shared<Primitive>("Receive");
inline const PrimitivePtr kPrimAllReduce = std::make_shared<Primitive>("AllReduce");
inline const PrimitivePtr kPrimNeighborExchange = std::make_shared<Primitive>("NeighborExchange");
inline const PrimitivePtr kPrimAllToAll = std::make_shared<Primitive>("_AlltoAll");
inline const PrimitivePtr kPrimAllToAllv = std::make_shared<Primitive>("AllToAllv");
inline const PrimitivePtr kPrimAllSwap = std::make_shared<Primitive>("AllSwap");
inline const PrimitivePtr kPrimBroadcast = std::make_shared<Primitive>("Broadcast");
inline const PrimitivePtr kPrimAllGather = std::make_shared<Primitive>("AllGather");

1
tests/ut/cpp/CMakeLists.txt
View File

@ -103,6 +103,7 @@ file(GLOB_RECURSE MINDSPORE_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR}
# dont remove the 4 lines above
"../../../mindspore/ccsrc/debug/data_dump/dump_json_parser.cc"
"../../../mindspore/ccsrc/debug/common.cc"
"../../../mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.cc"
"../../../mindspore/ccsrc/runtime/device/kernel_runtime.cc"
"../../../mindspore/ccsrc/runtime/device/memory_manager.cc"
"../../../mindspore/ccsrc/runtime/device/kernel_runtime_manager.cc"

282
tests/ut/cpp/device/hccl_adapter_test.cc Normal file
View File

@ -0,0 +1,282 @@
/**
* Copyright 2021 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 <memory>
#include "common/common_test.h"
#include "runtime/hccl_adapter/all_to_all_v_calc_param.h"
#include "backend/session/anf_runtime_algorithm.h"
#include "mindspore/core/ir/dtype/type_id.h"
namespace mindspore::hccl {
class TestHcclAdapter : public UT::Common {
public:
TestHcclAdapter() {}
protected:
CNodePtr CreateAllToAllvNode(const FuncGraphPtr &graph, const std::vector<AnfNodePtr> inputs,
const std::vector<int64_t> &send_rank_ids, const std::vector<int64_t> &recv_rank_ids) {
MS_EXCEPTION_IF_NULL(graph);
std::vector<AnfNodePtr> all_to_all_v_input = {NewValueNode(std::make_shared<Primitive>(kAllToAllVOpName))};
all_to_all_v_input.insert(all_to_all_v_input.end(), inputs.begin(), inputs.end());
auto all_to_all_v = graph->NewCNode(all_to_all_v_input);
MS_EXCEPTION_IF_NULL(all_to_all_v);
AnfAlgo::SetNodeAttr(kAttrSendRankIds, MakeValue<std::vector<int64_t>>(send_rank_ids), all_to_all_v);
AnfAlgo::SetNodeAttr(kAttrRecvRankIds, MakeValue<std::vector<int64_t>>(recv_rank_ids), all_to_all_v);
AnfAlgo::SetNodeAttr(kAttrGroup, MakeValue<std::string>("default_group"), all_to_all_v);
return all_to_all_v;
}
void SetOutputs(const CNodePtr &cnode, const std::vector<std::vector<size_t>> &shape,
const std::vector<TypeId> &data_type) {
AnfAlgo::SetOutputInferTypeAndShape(data_type, shape, cnode.get());
kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
builder.SetFusionType(kernel::FusionType::OPAQUE);
builder.SetProcessor(kernel::Processor::AICORE);
builder.SetKernelType(TBE_KERNEL);
builder.SetInputsFormat(std::vector<std::string>(cnode->size() - 1, format_));
builder.SetOutputsFormat(std::vector<std::string>(shape.size(), format_));
builder.SetInputsDeviceType(std::vector<TypeId>(cnode->size() - 1, type_));
builder.SetOutputsDeviceType(std::vector<TypeId>(shape.size(), type_));
cnode->set_kernel_info(std::make_shared<device::KernelInfo>());
AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), cnode.get());
}
std::vector<AnfNodePtr> CreateInputs(const FuncGraphPtr &graph, const std::vector<std::vector<size_t>> &shape,
const std::vector<TypeId> &data_type) {
MS_EXCEPTION_IF_NULL(graph);
if (shape.size() != data_type.size()) {
return {};
}
std::vector<AnfNodePtr> res;
for (size_t i = 0; i < shape.size(); ++i) {
auto node = graph->NewCNode(std::vector<AnfNodePtr>{NewValueNode(std::make_shared<Primitive>("AnyNameOp"))});
AnfAlgo::SetOutputInferTypeAndShape(std::vector<TypeId>{data_type[i]}, std::vector<std::vector<size_t>>{shape[i]},
node.get());
kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
builder.SetFusionType(kernel::FusionType::OPAQUE);
builder.SetProcessor(kernel::Processor::AICORE);
builder.SetKernelType(TBE_KERNEL);
builder.SetInputsFormat({format_});
builder.SetOutputsFormat({format_});
builder.SetInputsDeviceType({type_});
builder.SetOutputsDeviceType({type_});
node->set_kernel_info(std::make_shared<device::KernelInfo>());
AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), node.get());
res.emplace_back(node);
}
return res;
}
TypeId type_ = TypeId::kNumberTypeInt32;
std::string format_ = "NCHW";
};
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_only_send) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {1};
std::vector<int64_t> recv_rank_ids = {};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}}, {type_}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {}, {}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({0, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 0}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({0, 0}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 0}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_only_recv) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {};
std::vector<int64_t> recv_rank_ids = {0, 1};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {}, {}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}, {1}}, {type_, type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({0, 0}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 0}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({1, 1}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 128}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_only_send) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 4;
std::vector<int64_t> send_rank_ids = {1, 2, 3};
std::vector<int64_t> recv_rank_ids = {};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}, {1}, {1}}, {type_, type_, type_}), send_rank_ids,
recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {}, {}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({0, 1, 1, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 0, 128, 256}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({0, 0, 0, 0}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 0, 0, 0}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_only_send_2) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 4;
std::vector<int64_t> send_rank_ids = {1, 3};
std::vector<int64_t> recv_rank_ids = {};
auto alltoall =
CreateAllToAllvNode(graph, CreateInputs(graph, {{1}, {1}}, {type_, type_}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {}, {}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({0, 1, 0, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 0, 128, 128}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({0, 0, 0, 0}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 0, 0, 0}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_exchange) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {1};
std::vector<int64_t> recv_rank_ids = {1};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}}, {type_}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}}, {type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({0, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 0}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({0, 1}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 0}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_send_to_self) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {0};
std::vector<int64_t> recv_rank_ids = {0};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}}, {type_}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}}, {type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({1, 0}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 128}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({1, 0}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 128}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_all_to_all) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 4;
std::vector<int64_t> send_rank_ids = {0, 1, 2, 3};
std::vector<int64_t> recv_rank_ids = {0, 1, 2, 3};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}, {1}, {1}, {1}}, {type_, type_, type_, type_}),
send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}, {1}, {1}, {1}}, {type_, type_, type_, type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({1, 1, 1, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 128, 256, 384}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({1, 1, 1, 1}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 128, 256, 384}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_all_in_all_in_wrong_order) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 4;
std::vector<int64_t> send_rank_ids = {0, 1, 2, 3};
std::vector<int64_t> recv_rank_ids = {3, 1, 0, 2};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}, {1}, {1}, {1}}, {type_, type_, type_, type_}),
send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}, {1}, {1}, {1}}, {type_, type_, type_, type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({1, 1, 1, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 128, 256, 384}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({1, 1, 1, 1}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({256, 128, 384, 0}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_only_send_in_wrong_order) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 4;
std::vector<int64_t> send_rank_ids = {3, 1, 2};
std::vector<int64_t> recv_rank_ids = {};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {{1}, {1}, {1}}, {type_, type_, type_}), send_rank_ids,
recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {}, {}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_NO_THROW(calc.CalcOpParam());
EXPECT_EQ(calc.GetSendCounts(), std::vector<int64_t>({0, 1, 1, 1}));
EXPECT_EQ(calc.GetSendDispls(), std::vector<int64_t>({0, 128, 256, 0}));
EXPECT_EQ(calc.GetRecvCounts(), std::vector<int64_t>({0, 0, 0, 0}));
EXPECT_EQ(calc.GetRecvDispls(), std::vector<int64_t>({0, 0, 0, 0}));
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_invalid_rank_id) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {};
std::vector<int64_t> recv_rank_ids = {0, 2};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {}, {}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}, {1}}, {type_, type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_ANY_THROW(calc.CalcOpParam());
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_invalid_rank_id_2) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {};
std::vector<int64_t> recv_rank_ids = {0};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {}, {}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}, {1}}, {type_, type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_ANY_THROW(calc.CalcOpParam());
}
TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_wrong_order_and_invalid_rank_id) {
auto graph = std::make_shared<FuncGraph>();
ASSERT_TRUE(graph != nullptr);
uint32_t rank_size = 2;
std::vector<int64_t> send_rank_ids = {};
std::vector<int64_t> recv_rank_ids = {2, 0};
auto alltoall = CreateAllToAllvNode(graph, CreateInputs(graph, {}, {}), send_rank_ids, recv_rank_ids);
ASSERT_TRUE(alltoall != nullptr);
ASSERT_NO_THROW(SetOutputs(alltoall, {{1}, {1}}, {type_, type_}));
AllToAllvCalcParam calc(alltoall, rank_size);
ASSERT_ANY_THROW(calc.CalcOpParam());
}
} // namespace mindspore::hccl

85
tests/ut/cpp/pre_activate/ascend/mindir/all_to_all_unify_mindir_test.cc Normal file
View File

@ -0,0 +1,85 @@
/**
* Copyright 2021 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 "common/backend_common_test.h"
#include "frontend/operator/ops.h"
#include "debug/anf_ir_dump.h"
#include "common/py_func_graph_fetcher.h"
#include "backend/session/anf_runtime_algorithm.h"
#include "backend/optimizer/common/optimizer.h"
#include "backend/optimizer/common/pass_manager.h"
#include "backend/optimizer/pass/convert_const_input_to_attr.h"
#include "utils/utils.h"
#include "utils/ms_utils.h"
namespace mindspore {
namespace opt {
class TestAllToAllUnifyMindIr : public BackendCommon {
public:
TestAllToAllUnifyMindIr() : getPyFun_("gtest_input.pre_activate.all_to_all_unify_mindir_test", true) {}
~TestAllToAllUnifyMindIr() override = default;
public:
UT::PyFuncGraphFetcher getPyFun_;
};
TEST_F(TestAllToAllUnifyMindIr, test_neighbor_exchange) {
FuncGraphPtr g = getPyFun_.CallAndParseRet("test_neighbor_exchange", "before");
ASSERT_TRUE(g != nullptr);
std::vector<int64_t> shp_x{2, 3};
auto x_abstract = std::make_shared<abstract::AbstractTuple>(
AbstractBasePtrList{std::make_shared<abstract::AbstractTensor>(kFloat32, shp_x)});
AbstractBasePtrList args_spec_list{x_abstract};
auto func_graph = GetKernelGraph(g, args_spec_list);
ASSERT_TRUE(func_graph != nullptr);
bool has_all_to_all_v_node = false;
for (const auto &n : TopoSort(func_graph->get_return())) {
ASSERT_FALSE(IsPrimitiveCNode(n, prim::kPrimNeighborExchange));
if (IsPrimitiveCNode(n, prim::kPrimAllToAllv)) {
has_all_to_all_v_node = true;
}
}
ASSERT_TRUE(has_all_to_all_v_node);
}
TEST_F(TestAllToAllUnifyMindIr, test_all_to_all) {
FuncGraphPtr g = getPyFun_.CallAndParseRet("test_all_to_all", "before");
ASSERT_TRUE(g != nullptr);
std::vector<int64_t> shp_x{4, 2, 224, 224};
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp_x);
AbstractBasePtrList args_spec_list{x_abstract};
auto func_graph = GetKernelGraph(g, args_spec_list);
ASSERT_TRUE(func_graph != nullptr);
bool has_all_to_all_v_node = false;
bool has_concat_node = false;
bool has_split_v_node = false;
for (const auto &n : TopoSort(func_graph->get_return())) {
ASSERT_FALSE(IsPrimitiveCNode(n, prim::kPrimAllToAll));
if (IsPrimitiveCNode(n, prim::kPrimAllToAllv)) {
has_all_to_all_v_node = true;
}
if (IsPrimitiveCNode(n, prim::kPrimConcat)) {
has_concat_node = true;
}
if (IsPrimitiveCNode(n, prim::kPrimSplitV)) {
has_split_v_node = true;
}
}
ASSERT_TRUE(has_all_to_all_v_node);
ASSERT_TRUE(has_concat_node);
ASSERT_TRUE(has_split_v_node);
}
} // namespace opt
} // namespace mindspore

46
tests/ut/cpp/python_input/gtest_input/pre_activate/all_to_all_unify_mindir_test.py Normal file
View File

@ -0,0 +1,46 @@
# Copyright 2021 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.
# ============================================================================
import mindspore as ms
from mindspore.ops.operations._inner_ops import NeighborExchange
from mindspore.ops.operations.comm_ops import _AlltoAll
class FnDict:
def __init__(self):
self.fnDict = {}
def __call__(self, fn):
self.fnDict[fn.__name__] = fn
def __getitem__(self, name):
return self.fnDict[name]
def test_neighbor_exchange(tag):
fns = FnDict()
neighbor = NeighborExchange(send_rank_ids=[0], recv_rank_ids=[1], recv_shapes=([2, 3],), send_shapes=([2, 2],),
recv_type=ms.float32)
@fns
def before(x):
return neighbor(x)
return fns[tag]
def test_all_to_all(tag):
fns = FnDict()
altoall = _AlltoAll(split_count=8, split_dim=2, concat_dim=3)
@fns
def before(x):
return altoall(x)
return fns[tag]