diff --git a/mindspore/ccsrc/backend/kernel_compiler/hccl/hccl_kernel_metadata.cc b/mindspore/ccsrc/backend/kernel_compiler/hccl/hccl_kernel_metadata.cc index 6e3700a8908..d9a46984163 100755 --- a/mindspore/ccsrc/backend/kernel_compiler/hccl/hccl_kernel_metadata.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/hccl/hccl_kernel_metadata.cc @@ -67,7 +67,7 @@ void HcclMetadataInfo(const CNodePtr &kernel_node, std::vector &, const std::vector &, + const std::vector &, void *) { + return true; +} + +MS_HCCL_REG_KERNEL(AllToAllv, HcomAllToAllKernel); +} // namespace mindspore::kernel diff --git a/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_all_to_all.h b/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_all_to_all.h new file mode 100644 index 00000000000..70721c504d1 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_all_to_all.h @@ -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 +#include +#include "backend/kernel_compiler/hccl/hccl_kernel.h" + +namespace mindspore::kernel { +class HcomAllToAllKernel : public HcclKernel { + public: + HcomAllToAllKernel(); + ~HcomAllToAllKernel() override; + bool Launch(const std::vector &, const std::vector &, const std::vector &, + void *) override; +}; +} // namespace mindspore::kernel +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HCCL_HCOM_ALL_TO_ALL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_util.h b/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_util.h index a487b1a8fa1..13427e852b7 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_util.h +++ b/mindspore/ccsrc/backend/kernel_compiler/hccl/hcom_util.h @@ -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 CONST_OP_HCOM_DATA_TYPE_MAP = { diff --git a/mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.cc b/mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.cc new file mode 100644 index 00000000000..08cf18be42c --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.cc @@ -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 +#include +#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(); + 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(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(all_to_all, kAttrSplitCount); + int64_t split_dim = AnfAlgo::GetNodeAttr(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 split_input = {NewValueNode(std::make_shared(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 dtypes(split_count, dtype); + std::vector> shapes(split_count, shape); + AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split_v.get()); + AnfAlgo::SetNodeAttr(kAttrSplitDim, MakeValue(split_dim), split_v); + AnfAlgo::SetNodeAttr(kAttrNumSplit, MakeValue(split_count), split_v); + AnfAlgo::SetNodeAttr(kAttrSizeSplits, MakeValue(std::vector(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(all_to_all, kAttrSplitCount); + std::string group = AnfAlgo::GetNodeAttr(all_to_all, kAttrGroup); + std::vector split_outputs; + CreateMultipleOutputsOfAnfNode(graph, split, split_count, &split_outputs); + if (split_outputs.empty()) { + MS_LOG(EXCEPTION) << "Create tuple get item failed."; + } + std::vector all_to_all_v_input = {NewValueNode(std::make_shared(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 dtypes(split_count, single_type); + std::vector> shapes(split_count, single_shape); + AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, all_to_all_v.get()); + uint32_t rank_size = GetRankSize(group); + std::vector rank_ids(rank_size, 0); + for (uint32_t i = 0; i < rank_size; ++i) { + rank_ids[i] = static_cast(i); + } + + AnfAlgo::SetNodeAttr(kAttrSendRankIds, MakeValue>(rank_ids), all_to_all_v); + AnfAlgo::SetNodeAttr(kAttrRecvRankIds, MakeValue>(rank_ids), all_to_all_v); + AnfAlgo::SetNodeAttr(kAttrGroup, MakeValue(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(all_to_all, kAttrSplitCount); + int64_t concat_dim = AnfAlgo::GetNodeAttr(all_to_all, kAttrConcatDim); + std::vector 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 concat_input = {NewValueNode(std::make_shared(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(concat_dim), concat); + AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(split_count), concat); + std::vector 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()}); +} + +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()}); +} + +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(); + 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 diff --git a/mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.h b/mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.h new file mode 100644 index 00000000000..103bd8ff7ee --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/ascend/mindir/all_to_all_unify_mindir.h @@ -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 +#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_ diff --git a/mindspore/ccsrc/backend/session/ascend_session.cc b/mindspore/ccsrc/backend/session/ascend_session.cc index 0615d999ef0..eae542e2164 100644 --- a/mindspore/ccsrc/backend/session/ascend_session.cc +++ b/mindspore/ccsrc/backend/session/ascend_session.cc @@ -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()); unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); optimizer->AddPassManager(unify_mindir_pm); (void)optimizer->Optimize(graph); diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.h b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.h index 35b8711cb56..ee8804915ea 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.h +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.h @@ -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: diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc b/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc index c9f3f6932f7..e1a773864c8 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc @@ -1389,7 +1389,6 @@ vector AscendStreamAssign::GetLastInputCnode(const NotNull 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 AscendStreamAssign::GetLastInputCnode(const NotNull max) { - max_common_cnode = item.second.first; - } + max_common_cnode = item.second.first; } } diff --git a/mindspore/ccsrc/runtime/device/memory_manager.cc b/mindspore/ccsrc/runtime/device/memory_manager.cc index 08646bebe23..e097f5c9820 100644 --- a/mindspore/ccsrc/runtime/device/memory_manager.cc +++ b/mindspore/ccsrc/runtime/device/memory_manager.cc @@ -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; } diff --git a/mindspore/ccsrc/runtime/device/memory_manager.h b/mindspore/ccsrc/runtime/device/memory_manager.h index 769b07cf9cc..0972e5cb8b4 100644 --- a/mindspore/ccsrc/runtime/device/memory_manager.h +++ b/mindspore/ccsrc/runtime/device/memory_manager.h @@ -58,7 +58,7 @@ class MemoryManager { virtual std::vector MallocContinuousMemFromMemPool(size_t total_size, std::vector 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; diff --git a/mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.cc b/mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.cc new file mode 100644 index 00000000000..d9628df3682 --- /dev/null +++ b/mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.cc @@ -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 +#include +#include +#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 &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 input_aligned_mem_size(input_num); + std::vector output_aligned_mem_size(output_num); + std::vector input_real_mem_size(input_num); + std::vector 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 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 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 &mem_sizes, const std::vector &real_sizes, + const std::string &rank_ids_attr, std::vector *counts, + std::vector *displs) { + CNodePtr cnode = node_.lock(); + MS_EXCEPTION_IF_NULL(cnode); + auto rank_ids = AnfAlgo::GetNodeAttr>(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 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(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 rank_id_map; + for (size_t i = 0; i < rank_ids.size(); ++i) { + if (rank_ids[i] < 0 || static_cast(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 diff --git a/mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.h b/mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.h new file mode 100644 index 00000000000..8bd37b9669f --- /dev/null +++ b/mindspore/ccsrc/runtime/hccl_adapter/all_to_all_v_calc_param.h @@ -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 +#include +#include +#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 &GetSendCounts() const { return send_counts_; } + const std::vector &GetSendDispls() const { return sdispls_; } + const std::vector &GetRecvCounts() const { return recv_counts_; } + const std::vector &GetRecvDispls() const { return rdispls_; } + + private: + void CalcMemOffset(const std::vector &mem_sizes, const std::vector &real_sizes, + const std::string &rank_ids_attr, std::vector *counts, std::vector *displs); + CNodeWeakPtr node_; + uint32_t rank_size_; + std::vector send_counts_; + std::vector sdispls_; + std::vector recv_counts_; + std::vector rdispls_; +}; +} // namespace mindspore::hccl +#endif // MINDSPORE_RUNTIME_HCCL_ADAPTER_ALL_TO_ALL_V_CALC_PARAM_H diff --git a/mindspore/ccsrc/runtime/hccl_adapter/converter.cc b/mindspore/ccsrc/runtime/hccl_adapter/converter.cc index a6280043fa1..4b2c069f98f 100644 --- a/mindspore/ccsrc/runtime/hccl_adapter/converter.cc +++ b/mindspore/ccsrc/runtime/hccl_adapter/converter.cc @@ -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 &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 GenerateStubGeNode(const AnfNodePtr ge::OpDescPtr op_desc = std::make_shared(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 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(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 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(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 GenerateStubGeNode(const AnfNodePtr // set node attr (void)ConvertAttr(cnode, op_desc, kAttrRankSize, ge::HCOM_ATTR_RANK_SIZE); - (void)ConvertAttr(cnode, op_desc, kAttrGroup, ge::HCOM_ATTR_GROUP); + auto group = ConvertAttr(cnode, op_desc, kAttrGroup, ge::HCOM_ATTR_GROUP); (void)ConvertAttr(cnode, op_desc, kAttrSrcRank, ge::HCOM_ATTR_SRC_RANK); (void)ConvertAttr(cnode, op_desc, kAttrDestRank, ge::HCOM_ATTR_DEST_RANK); (void)ConvertAttr(cnode, op_desc, kAttrSrTag, ge::HCOM_ATTR_SR_TAG); (void)ConvertAttr>(cnode, op_desc, kAttrShape, ge::HCOM_ATTR_SHAPE); + SetAllToAllvAttr(cnode, op_desc, group); ge::ComputeGraphPtr ge_graph = std::make_shared(kStubDataStructureName); MS_EXCEPTION_IF_NULL(ge_graph); diff --git a/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.cc b/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.cc index 2fe05293e58..24fb30f82f6 100644 --- a/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.cc +++ b/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.cc @@ -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 ¶ms, 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 diff --git a/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.h b/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.h index 2dff5a46ab0..f3c39937405 100644 --- a/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.h +++ b/mindspore/ccsrc/runtime/hccl_adapter/hccl_adapter.h @@ -63,6 +63,7 @@ class HcclAdapter { // for enqueue op HcclResult HcclExecEnqueueOp(const ::HcomOperation &op_info, const HExecCallBack &callback) const; + HcclResult HcclExecAllToAllv(const ::HcomAllToAllVParams ¶ms, 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; diff --git a/mindspore/ccsrc/runtime/hccl_adapter/plugin/hccl_plugin.h b/mindspore/ccsrc/runtime/hccl_adapter/plugin/hccl_plugin.h index f3f521cf078..a4b5fa3b0ae 100644 --- a/mindspore/ccsrc/runtime/hccl_adapter/plugin/hccl_plugin.h +++ b/mindspore/ccsrc/runtime/hccl_adapter/plugin/hccl_plugin.h @@ -22,6 +22,7 @@ #include #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 diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h index ff3241f6049..7e884d52645 100644 --- a/mindspore/ccsrc/utils/utils.h +++ b/mindspore/ccsrc/utils/utils.h @@ -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"; diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index 5657d8f2df9..ab7c128ffbf 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -389,6 +389,8 @@ inline const PrimitivePtr kPrimSend = std::make_shared("Send"); inline const PrimitivePtr kPrimReceive = std::make_shared("Receive"); inline const PrimitivePtr kPrimAllReduce = std::make_shared("AllReduce"); inline const PrimitivePtr kPrimNeighborExchange = std::make_shared("NeighborExchange"); +inline const PrimitivePtr kPrimAllToAll = std::make_shared("_AlltoAll"); +inline const PrimitivePtr kPrimAllToAllv = std::make_shared("AllToAllv"); inline const PrimitivePtr kPrimAllSwap = std::make_shared("AllSwap"); inline const PrimitivePtr kPrimBroadcast = std::make_shared("Broadcast"); inline const PrimitivePtr kPrimAllGather = std::make_shared("AllGather"); diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index f8cebf1cc0e..58288960327 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -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" diff --git a/tests/ut/cpp/device/hccl_adapter_test.cc b/tests/ut/cpp/device/hccl_adapter_test.cc new file mode 100644 index 00000000000..6c3b6fdeb56 --- /dev/null +++ b/tests/ut/cpp/device/hccl_adapter_test.cc @@ -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 +#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 inputs, + const std::vector &send_rank_ids, const std::vector &recv_rank_ids) { + MS_EXCEPTION_IF_NULL(graph); + std::vector all_to_all_v_input = {NewValueNode(std::make_shared(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>(send_rank_ids), all_to_all_v); + AnfAlgo::SetNodeAttr(kAttrRecvRankIds, MakeValue>(recv_rank_ids), all_to_all_v); + AnfAlgo::SetNodeAttr(kAttrGroup, MakeValue("default_group"), all_to_all_v); + return all_to_all_v; + } + + void SetOutputs(const CNodePtr &cnode, const std::vector> &shape, + const std::vector &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(cnode->size() - 1, format_)); + builder.SetOutputsFormat(std::vector(shape.size(), format_)); + builder.SetInputsDeviceType(std::vector(cnode->size() - 1, type_)); + builder.SetOutputsDeviceType(std::vector(shape.size(), type_)); + cnode->set_kernel_info(std::make_shared()); + AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), cnode.get()); + } + + std::vector CreateInputs(const FuncGraphPtr &graph, const std::vector> &shape, + const std::vector &data_type) { + MS_EXCEPTION_IF_NULL(graph); + if (shape.size() != data_type.size()) { + return {}; + } + std::vector res; + for (size_t i = 0; i < shape.size(); ++i) { + auto node = graph->NewCNode(std::vector{NewValueNode(std::make_shared("AnyNameOp"))}); + AnfAlgo::SetOutputInferTypeAndShape(std::vector{data_type[i]}, std::vector>{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()); + 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(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {1}; + std::vector 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({0, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 0})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({0, 0})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 0})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_only_recv) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {}; + std::vector 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({0, 0})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 0})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({1, 1})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 128})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_only_send) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 4; + std::vector send_rank_ids = {1, 2, 3}; + std::vector 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({0, 1, 1, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 0, 128, 256})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({0, 0, 0, 0})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 0, 0, 0})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_only_send_2) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 4; + std::vector send_rank_ids = {1, 3}; + std::vector 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({0, 1, 0, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 0, 128, 128})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({0, 0, 0, 0})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 0, 0, 0})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_exchange) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {1}; + std::vector 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({0, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 0})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({0, 1})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 0})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_send_to_self) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {0}; + std::vector 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({1, 0})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 128})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({1, 0})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 128})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_4p_all_to_all) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 4; + std::vector send_rank_ids = {0, 1, 2, 3}; + std::vector 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({1, 1, 1, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 128, 256, 384})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({1, 1, 1, 1})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({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(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 4; + std::vector send_rank_ids = {0, 1, 2, 3}; + std::vector 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({1, 1, 1, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 128, 256, 384})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({1, 1, 1, 1})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({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(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 4; + std::vector send_rank_ids = {3, 1, 2}; + std::vector 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({0, 1, 1, 1})); + EXPECT_EQ(calc.GetSendDispls(), std::vector({0, 128, 256, 0})); + EXPECT_EQ(calc.GetRecvCounts(), std::vector({0, 0, 0, 0})); + EXPECT_EQ(calc.GetRecvDispls(), std::vector({0, 0, 0, 0})); +} + +TEST_F(TestHcclAdapter, test_all_to_all_v_calc_param_2p_invalid_rank_id) { + auto graph = std::make_shared(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {}; + std::vector 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(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {}; + std::vector 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(); + ASSERT_TRUE(graph != nullptr); + uint32_t rank_size = 2; + std::vector send_rank_ids = {}; + std::vector 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 diff --git a/tests/ut/cpp/pre_activate/ascend/mindir/all_to_all_unify_mindir_test.cc b/tests/ut/cpp/pre_activate/ascend/mindir/all_to_all_unify_mindir_test.cc new file mode 100644 index 00000000000..9338e293e09 --- /dev/null +++ b/tests/ut/cpp/pre_activate/ascend/mindir/all_to_all_unify_mindir_test.cc @@ -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 shp_x{2, 3}; + auto x_abstract = std::make_shared( + AbstractBasePtrList{std::make_shared(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 shp_x{4, 2, 224, 224}; + auto x_abstract = std::make_shared(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 diff --git a/tests/ut/cpp/python_input/gtest_input/pre_activate/all_to_all_unify_mindir_test.py b/tests/ut/cpp/python_input/gtest_input/pre_activate/all_to_all_unify_mindir_test.py new file mode 100644 index 00000000000..08d49cee307 --- /dev/null +++ b/tests/ut/cpp/python_input/gtest_input/pre_activate/all_to_all_unify_mindir_test.py @@ -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]