From 91b09a8e0d312fba841738a9cddacc7501991c6e Mon Sep 17 00:00:00 2001 From: zhupuxu Date: Thu, 12 May 2022 15:00:10 +0800 Subject: [PATCH] fix bug for atomic Signed-off-by: zhupuxu --- .../kernel/tbe/dynamic_tbe_kernel_mod.cc | 20 ++++---- .../kernel/tbe/tiling/op_tiling_adapter.cc | 49 +++++++++++++------ .../kernel/tbe/tiling/op_tiling_adapter.h | 8 +-- tests/ut/cpp/stub/tiling/op_tiling_adapter.cc | 11 +++-- 4 files changed, 59 insertions(+), 29 deletions(-) diff --git a/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/dynamic_tbe_kernel_mod.cc b/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/dynamic_tbe_kernel_mod.cc index 5a2d07f45ae..7c70b6f7be4 100644 --- a/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/dynamic_tbe_kernel_mod.cc +++ b/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/dynamic_tbe_kernel_mod.cc @@ -78,15 +78,20 @@ int DynamicTbeKernelMod::Resize(const BaseOperatorPtr &base_operator, const std: MS_LOG(EXCEPTION) << "The node is not dynamic shape: " << cnode->fullname_with_scope(); } + // update output size after InferShape. + // avoid atomic_clean memory violation, we need dynamic atomic_clean op. + AscendKernelMod::UpdateOutputSizeList(); + // compute tiling of atomic_clean op. device::tiling::OpTilingCalculateAdapter converter; - ::ge::ComputeGraphPtr ge_graph = std::make_shared<::ge::ComputeGraph>("atomic_clean"); const std::map &depend_tensor_map = inputsOnHost; + ::ge::ComputeGraphPtr ge_graph = std::make_shared<::ge::ComputeGraph>("default"); + ::ge::NodePtr ge_node = nullptr; if (!atomic_clean_nodes_.empty()) { atomic_compile_info_ = ParseCompileJson(atomic_clean_nodes_[0].lock()); - optiling::utils::OpRunInfo atomic_op_info(-1, true, 0); - auto ge_node = converter.AnfNodeToGeNodeAdapter(cnode, &ge_graph, depend_tensor_map, op_compile_info_); + ge_node = converter.AnfNodeToGeNodeAdapter(cnode, &ge_graph, depend_tensor_map, op_compile_info_); MS_EXCEPTION_IF_NULL(ge_node); + optiling::utils::OpRunInfo atomic_op_info(-1, true, 0); auto ret = optiling::OpAtomicCalculateV2(*ge_node, atomic_op_info); if (ret != ::ge::GRAPH_SUCCESS) { MS_LOG(EXCEPTION) << "The node: " << cnode->fullname_with_scope() << " compute atomic tiling failed!"; @@ -98,10 +103,6 @@ int DynamicTbeKernelMod::Resize(const BaseOperatorPtr &base_operator, const std: } } - // update output size after InferShape. - // avoid atomic_clean memory violation, we need dynamic atomic_clean op. - AscendKernelMod::UpdateOutputSizeList(); - need_skip_execute_ = AnfAlgo::IsDynamicShapeSkipExecute(cnode); if (need_skip_execute_) { return 0; @@ -125,8 +126,9 @@ int DynamicTbeKernelMod::Resize(const BaseOperatorPtr &base_operator, const std: // start compute tiling optiling::utils::OpRunInfo op_run_info_v2(-1, true, 0); MS_LOG(INFO) << "Start compute tiling of: " << cnode->fullname_with_scope(); - ge_graph = std::make_shared<::ge::ComputeGraph>("default"); - auto ge_op = converter.AnfNodeToGeOperatorAdapter(cnode, &ge_graph, depend_tensor_map, op_compile_info_); + auto ge_op = (ge_node != nullptr) + ? converter.GeNodeToGeOperatorAdapter(ge_node) + : converter.AnfNodeToGeOperatorAdapter(cnode, &ge_graph, depend_tensor_map, op_compile_info_); auto ret = optiling::OpParaCalculateV2(ge_op, op_run_info_v2); if (ret != ::ge::GRAPH_SUCCESS) { MS_LOG(EXCEPTION) << "The node: " << cnode->fullname_with_scope() << " compute tiling failed!"; diff --git a/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.cc b/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.cc index 70895790ffc..680f94fc5b0 100644 --- a/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.cc +++ b/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.cc @@ -29,6 +29,7 @@ #include "common/ge_inner_error_codes.h" #include "graph/utils/op_desc_utils.h" #include "kernel/ascend_kernel_mod.h" +#include "graph/utils/tensor_utils.h" namespace mindspore { namespace device { @@ -164,25 +165,40 @@ void OpTilingCalculateAdapter::ConvertAtomicCompileInfo(const CNodePtr &node, :: MS_EXCEPTION_IF_NULL(*op_desc); auto kernel_mod = dynamic_cast(AnfAlgo::GetKernelMod(node)); MS_EXCEPTION_IF_NULL(kernel_mod); - // clean output - if (common::AnfAlgo::HasNodeAttr(kAttrAtomicOutputIndexs, node)) { - vector output_indexs; - auto help = common::AnfAlgo::GetNodeAttr>(node, kAttrAtomicOutputIndexs); - std::transform(help.begin(), help.end(), std::back_inserter(output_indexs), SizeToLong); - ::ge::AttrUtils::SetListInt(*(*op_desc), ::ge::ATOMIC_ATTR_OUTPUT_INDEX, output_indexs); + bool has_output = common::AnfAlgo::HasNodeAttr(kAttrAtomicOutputIndexs, node); + bool has_workspace = common::AnfAlgo::HasNodeAttr(kAttrAtomicWorkspaceIndexs, node); + constexpr size_t kAlignBytes = 32 - 1; + // set atomic compile info + if (has_output || has_workspace) { std::string atomic_compile_info = kernel_mod->GetAtomicCompileInfo(); std::string atomic_info_key = std::to_string(std::hash()(atomic_compile_info)); (void)::ge::AttrUtils::SetStr(*(*op_desc), ATOMIC_COMPILE_INFO_KEY, atomic_info_key); (void)::ge::AttrUtils::SetStr(*(*op_desc), ATOMIC_COMPILE_INFO_JSON, atomic_compile_info); } + + // clean output + if (has_output) { + vector output_indexs; + auto help = common::AnfAlgo::GetNodeAttr>(node, kAttrAtomicOutputIndexs); + std::transform(help.begin(), help.end(), std::back_inserter(output_indexs), SizeToLong); + ::ge::AttrUtils::SetListInt(*(*op_desc), ::ge::ATOMIC_ATTR_OUTPUT_INDEX, output_indexs); + auto output_mem_size = kernel_mod->GetOutputSizeList(); + for (auto index : output_indexs) { + auto output_size = + static_cast((output_mem_size.at(index) + kMemAlignSize + kAlignBytes) / kMemAlignSize * kMemAlignSize); + auto output = (*op_desc)->MutableOutputDesc(index); + MS_EXCEPTION_IF_NULL(output); + ::ge::TensorUtils::SetSize(*output, output_size); + } + } + // clean workspace - if (common::AnfAlgo::HasNodeAttr(kAttrAtomicWorkspaceIndexs, node)) { + if (has_workspace) { auto workspace_men_sizes = kernel_mod->GetWorkspaceSizeList(); std::vector workspace_list; std::transform(workspace_men_sizes.begin(), workspace_men_sizes.end(), std::back_inserter(workspace_list), SizeToLong); (void)(*op_desc)->SetWorkspaceBytes(workspace_list); - constexpr size_t kAlignBytes = 32 - 1; std::map> workspace_info; std::map clean_size_list; auto workspace_indexes = common::AnfAlgo::GetNodeAttr>(node, kAttrAtomicWorkspaceIndexs); @@ -342,12 +358,7 @@ void OpTilingCalculateAdapter::InitOpIoName(const CNodePtr &node) { return ge_node; } -::ge::Operator OpTilingCalculateAdapter::AnfNodeToGeOperatorAdapter( - const CNodePtr &node, ::ge::ComputeGraphPtr *ge_graph, const std::map &depend_tensor_map, - const std::string &op_compile_info) { - MS_EXCEPTION_IF_NULL(node); - MS_EXCEPTION_IF_NULL(*ge_graph); - auto ge_node = CreateGeNode(node, ge_graph, depend_tensor_map, op_compile_info); +::ge::Operator OpTilingCalculateAdapter::GeNodeToGeOperatorAdapter(const ::ge::NodePtr &ge_node) { MS_EXCEPTION_IF_NULL(ge_node); return ::ge::OpDescUtils::CreateOperatorFromNode(ge_node); } @@ -359,6 +370,16 @@ void OpTilingCalculateAdapter::InitOpIoName(const CNodePtr &node) { MS_EXCEPTION_IF_NULL(*ge_graph); return CreateGeNode(node, ge_graph, depend_tensor_map, op_compile_info); } + +::ge::Operator OpTilingCalculateAdapter::AnfNodeToGeOperatorAdapter( + const CNodePtr &node, ::ge::ComputeGraphPtr *ge_graph, const std::map &depend_tensor_map, + const std::string &op_compile_info) { + MS_EXCEPTION_IF_NULL(node); + MS_EXCEPTION_IF_NULL(*ge_graph); + auto ge_node = CreateGeNode(node, ge_graph, depend_tensor_map, op_compile_info); + MS_EXCEPTION_IF_NULL(ge_node); + return ::ge::OpDescUtils::CreateOperatorFromNode(ge_node); +} } // namespace tiling } // namespace device } // namespace mindspore diff --git a/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.h b/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.h index 05c0a84d200..ad7a4146395 100644 --- a/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.h +++ b/mindspore/ccsrc/plugin/device/ascend/kernel/tbe/tiling/op_tiling_adapter.h @@ -37,14 +37,16 @@ class OpTilingCalculateAdapter { OpTilingCalculateAdapter() = default; ~OpTilingCalculateAdapter() = default; - ::ge::Operator AnfNodeToGeOperatorAdapter(const CNodePtr &node, ::ge::ComputeGraphPtr *ge_graph, - const std::map &depend_tensor_map, - const std::string &op_compile_info); + ::ge::Operator GeNodeToGeOperatorAdapter(const ::ge::NodePtr &ge_node); ::ge::NodePtr AnfNodeToGeNodeAdapter(const CNodePtr &node, ::ge::ComputeGraphPtr *ge_graph, const std::map &depend_tensor_map, const std::string &op_compile_info); + ::ge::Operator AnfNodeToGeOperatorAdapter(const CNodePtr &node, ::ge::ComputeGraphPtr *ge_graph, + const std::map &depend_tensor_map, + const std::string &op_compile_info); + private: void ConvertInputShapeAndType(const CNodePtr &node, ::ge::OpDescPtr *op_desc); void ConvertOutputShapeAndType(const CNodePtr &node, ::ge::OpDescPtr *op_desc); diff --git a/tests/ut/cpp/stub/tiling/op_tiling_adapter.cc b/tests/ut/cpp/stub/tiling/op_tiling_adapter.cc index c4e3ff0a39e..c080660178a 100644 --- a/tests/ut/cpp/stub/tiling/op_tiling_adapter.cc +++ b/tests/ut/cpp/stub/tiling/op_tiling_adapter.cc @@ -97,9 +97,7 @@ void OpTilingCalculateAdapter::AddEdge(const ge::NodePtr &ge_node, void OpTilingCalculateAdapter::InitOpIoName(const CNodePtr &node) {} -ge::Operator OpTilingCalculateAdapter::AnfNodeToGeOperatorAdapter( - const CNodePtr &node, ge::ComputeGraphPtr *ge_graph, const std::map &depend_tensor_map, - const std::string &op_compile_info) { +ge::Operator OpTilingCalculateAdapter::GeNodeToGeOperatorAdapter(const ::ge::NodePtr &ge_node) { ge::Operator op; return op; } @@ -110,6 +108,13 @@ ge::NodePtr OpTilingCalculateAdapter::AnfNodeToGeNodeAdapter( ge::NodePtr ge_node; return ge_node; } + +ge::Operator OpTilingCalculateAdapter::AnfNodeToGeOperatorAdapter(const CNodePtr &node, ::ge::ComputeGraphPtr *ge_graph, + const std::map &depend_tensor_map, + const std::string &op_compile_info) { + ge::Operator op; + return op; +} } // namespace tiling } // namespace device } // namespace mindspore