From 134ffa9efefd3f70dd741201e12760ff1274bf28 Mon Sep 17 00:00:00 2001 From: tronzhang Date: Tue, 17 Aug 2021 21:12:08 +0800 Subject: [PATCH] support tot ops --- akg | 2 +- .../graph_kernel/expanders/__init__.py | 3 +- .../_extends/graph_kernel/expanders/gather.py | 43 +++ .../graph_kernel/model/graph_split.py | 61 ++++ .../_extends/graph_kernel/model/model.py | 5 + .../_extends/graph_kernel/model/op_infer.py | 27 +- .../graph_kernel/add_atomic_clean.cc | 13 +- .../optimizer/graph_kernel/add_atomic_clean.h | 18 +- .../graph_kernel/graph_kernel_helper.cc | 24 +- .../graph_kernel/graph_kernel_helper.h | 2 + .../graph_kernel/graph_kernel_optimization.cc | 7 + .../tsa_atomic_add_to_first_tensor.cc | 270 ++++++++++++++++++ .../tsa_atomic_add_to_first_tensor.h | 64 +++++ .../optimizer/graph_kernel/uss_atomic_add.cc | 89 ++++++ .../optimizer/graph_kernel/uss_atomic_add.h | 54 ++++ 15 files changed, 648 insertions(+), 34 deletions(-) create mode 100644 mindspore/_extends/graph_kernel/expanders/gather.py create mode 100644 mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.cc create mode 100644 mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h create mode 100644 mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.cc create mode 100644 mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.h diff --git a/akg b/akg index 77aaa23fc16..5196d6e1e8f 160000 --- a/akg +++ b/akg @@ -1 +1 @@ -Subproject commit 77aaa23fc160756bdf40eef97333305680b3b35c +Subproject commit 5196d6e1e8feb1a5ab074541e0c2acf072a6f9b4 diff --git a/mindspore/_extends/graph_kernel/expanders/__init__.py b/mindspore/_extends/graph_kernel/expanders/__init__.py index 11fcd76080a..fb229621c89 100644 --- a/mindspore/_extends/graph_kernel/expanders/__init__.py +++ b/mindspore/_extends/graph_kernel/expanders/__init__.py @@ -28,6 +28,7 @@ from .erfc import Erfc from .fused_adam import FusedAdam from .fused_adam_weight_decay import FusedAdamWeightDecay from .fused_mul_add import FusedMulAdd +from .gather import Gather from .gelu import GeLU from .gelu_grad import GeLUGrad from .gkdropout import GkDropout @@ -55,8 +56,8 @@ from .softmax_cross_entropy_with_logits import SoftmaxCrossEntropyWithLogits from .softmax_grad_ext import SoftmaxGradExt from .sqrt_grad import SqrtGrad from .square import Square -from .square_sum_v1 import SquareSumV1 from .squared_difference import SquaredDifference +from .square_sum_v1 import SquareSumV1 from .square_sum_all import SquareSumAll from .squeeze import Squeeze from .tanh_grad import TanhGrad diff --git a/mindspore/_extends/graph_kernel/expanders/gather.py b/mindspore/_extends/graph_kernel/expanders/gather.py new file mode 100644 index 00000000000..2eec945edc9 --- /dev/null +++ b/mindspore/_extends/graph_kernel/expanders/gather.py @@ -0,0 +1,43 @@ +# 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. +# =========================================================================== +"""generate json desc for gather""" +from ._utils import Expander, ExpanderInfoValidator as VLD + + +@VLD.check_all_formats_same +@VLD.check_attrs('axis') +class Gather(Expander): + """Expand Gather""" + + def _expand(self, graph_builder): + inputs, indices = self.inputs + axis = self.attrs['axis'] + if axis < 0: + axis += len(inputs.shape) + if len(indices.shape) == 1: + result = graph_builder.emit('Gather', [inputs, indices], attrs={'axis': axis}) + else: + ori_indices_shape = indices.shape + indices_shape_one_dim = 1 + for dim in ori_indices_shape: + indices_shape_one_dim *= dim + new_indices_shape = [indices_shape_one_dim] + reshape_indices = graph_builder.emit('Reshape', [indices], attrs={'shape': new_indices_shape}) + tmp_result = graph_builder.emit('Gather', [inputs, reshape_indices], attrs={'axis': axis}) + output_shape = inputs.shape.copy() + output_shape[axis:axis] = ori_indices_shape + del output_shape[axis + len(ori_indices_shape)] + result = graph_builder.emit('Reshape', [tmp_result], attrs={'shape': output_shape}) + return result diff --git a/mindspore/_extends/graph_kernel/model/graph_split.py b/mindspore/_extends/graph_kernel/model/graph_split.py index f267b928de6..6bcb74baa03 100644 --- a/mindspore/_extends/graph_kernel/model/graph_split.py +++ b/mindspore/_extends/graph_kernel/model/graph_split.py @@ -483,6 +483,8 @@ class GraphSplitByPattern: def _find_cheap_regions(dom): sub = self.to_subgraph(dom) inputs, outputs = sub.deduce_parameters() + if not inputs: + return list() cheap_regions = [] for output in outputs: # tensor should have user other than user_area to be fused @@ -788,6 +790,63 @@ class GraphSplitGpu(GraphSplitByPattern): fused.append(a) return fused, True + def _gather_output(dom): + gather_prims = ("Gather", "GatherNd") + if not dom.dom_op().prim in gather_prims: + return None + + def _count_target_prim(ops, target_list): + count = 0 + for op in ops: + if op.prim in target_list: + count += 1 + return count + + def _shape_consistent(start_prims, end_prims, source, target): + start_ops = [] + for op in source.ops: + if op.prim in start_prims: + start_ops.append(op) + + total_ops = source.ops + target.ops + for start_op in start_ops: + consisten_shape = start_op.output.shape + visited = [] + op_queue = [start_op] + while op_queue: + tmp_queue = [] + for op in op_queue: + if op in visited: + continue + if op.prim in end_prims or not op in total_ops: + continue + if (op.prim in start_prims and op != start_op) or consisten_shape != op.output.shape: + return False + for to_op in op.output.to_ops: + tmp_queue.append(to_op) + visited.append(op) + op_queue = tmp_queue + return True + + appected_areas = {"TensorScatterAdd", "UnsortedSegmentSum"} + for a, _ in dom.out_relations.items(): + if _shape_consistent(gather_prims, appected_areas, dom, a) and \ + _count_target_prim(a.ops + dom.ops, appected_areas) < 2 and dom.check_acyclic(a): + return [a], False + return None + + def _broadcast_opaque(dom): + fuse_arg = {"TensorScatterAdd": slice(1, None), "UnsortedSegmentSum": slice(0, 2)} + arg_idx = fuse_arg.get(dom.dom_op().prim, -1) + if arg_idx == -1 or len(dom.ops) != 1: + return None + fuse_tensor = dom.dom_op().inputs[arg_idx] + for a, _ in dom.in_relations.items(): + if a.pattern <= PrimLib.BROADCAST and dom.check_acyclic(a) and \ + any([op.output in fuse_tensor for op in a.ops]): + return [a], True + return None + def _fuse_loop(): changed = True while changed: @@ -799,6 +858,8 @@ class GraphSplitGpu(GraphSplitByPattern): changed = self.fuse(_broadcast_depth) or changed changed = self.fuse(_broadcast_width) or changed changed = self.fuse(_strided_slice) or changed + changed = self.fuse(_broadcast_opaque) or changed + changed = self.fuse(_gather_output) or changed changed = self.fuse(_reduce_output) or changed if enable_stitch_fusion: changed = self.fuse(_reduce_stitch) or changed diff --git a/mindspore/_extends/graph_kernel/model/model.py b/mindspore/_extends/graph_kernel/model/model.py index 06a1c18a2ad..c64aefc8713 100644 --- a/mindspore/_extends/graph_kernel/model/model.py +++ b/mindspore/_extends/graph_kernel/model/model.py @@ -17,6 +17,7 @@ class Utils: """Model utils""" + def __init__(self): pass @@ -229,6 +230,10 @@ class PrimLib: 'Atan': Prim(ELEMWISE), 'Atan2': Prim(ELEMWISE), 'Expm1': Prim(ELEMWISE), + 'TensorScatterAdd': Prim(OPAQUE), + 'Gather': Prim(OPAQUE), + 'GatherNd': Prim(OPAQUE), + 'UnsortedSegmentSum': Prim(OPAQUE), } default_primtive = Prim(UNKNOWN) diff --git a/mindspore/_extends/graph_kernel/model/op_infer.py b/mindspore/_extends/graph_kernel/model/op_infer.py index bf442d07a2c..760f6586efe 100644 --- a/mindspore/_extends/graph_kernel/model/op_infer.py +++ b/mindspore/_extends/graph_kernel/model/op_infer.py @@ -169,7 +169,7 @@ class _Reduce(OpInfer): """Common infer for reduction operators""" def _check(self): - super()._check() + super(_Reduce, self)._check() # check reduce axis in the range [-len, len) shape_len = len(self.inputs[0].shape) axis = self.attrs['reduce_axis'] @@ -451,3 +451,28 @@ class UnPadAkg(OpInfer): raise GKException("Input dimension and pad mismatch: {}d vs {}d".format(n, len(unpad_after))) out_shape = [shape[i] - unpad_after[i] for i in range(n)] return out_shape + + +class Gather(OpInfer): + """Gather infer""" + + def _infer_shape(self): + input_shape = self.inputs[0].shape + indices_shape = self.inputs[1].shape + axis = self.attrs['axis'] + output_shape = input_shape + indices_shape_one_dim = 1 + for dim in indices_shape: + indices_shape_one_dim *= dim + output_shape[axis] = indices_shape_one_dim + return output_shape + + def _infer_type(self): + return self.inputs[0].dtype + + def _infer_format(self): + return self.inputs[0].data_format + + def _check_type(self): + if self.inputs[1].dtype != "int32": + raise GKException("Indices dtype must be int32!") diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.cc index e7367de351f..57f6c712749 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.cc @@ -114,21 +114,21 @@ bool AtomicAddChecker::FindCandidate(const AnfNodePtr &anf_node) { // Rule: Only one ReduceSum inside sub-graph. auto real_return_node = sub_graph->get_return()->input(kFirstDataInputIndex); if (IsPrimitiveCNode(real_return_node, prim::kPrimMakeTuple)) { - size_t reduce_cnt = 0; + size_t target_cnt = 0; const auto &inputs = real_return_node->cast()->inputs(); for (size_t i = 1; i < inputs.size(); ++i) { - if (IsPrimitiveCNode(inputs[i], prim::kPrimReduceSum)) { + if (IsPrimitiveCNode(inputs[i], target_type_)) { atomic_add_info_.atomic_add_node = inputs[i]->cast(); atomic_add_info_.reduce_real_output_index = i - 1; - reduce_cnt++; + target_cnt++; } } - if (reduce_cnt != 1) { + if (target_cnt != 1) { return false; } atomic_add_info_.real_output_num = inputs.size() - 1; - } else if (IsPrimitiveCNode(real_return_node, prim::kPrimReduceSum)) { + } else if (IsPrimitiveCNode(real_return_node, target_type_)) { atomic_add_info_.atomic_add_node = real_return_node->cast(); atomic_add_info_.real_output_num = 1; } else { @@ -385,6 +385,7 @@ CNodePtr AtomicCleanInsertter::InsertUpdateState(const KernelGraphPtr &main_grap u->set_abstract(kUMonad->ToAbstract()); AnfNodePtrList update_state_inputs = {NewValueNode(prim::kPrimUpdateState), u, composite_node}; auto update_state_cnode = main_graph->NewCNode(update_state_inputs); + update_state_cnode->set_abstract(kUMonad->ToAbstract()); main_graph->AddNode(update_state_cnode); return update_state_cnode; } @@ -509,11 +510,11 @@ void AtomicCleanInsertter::ProcessOriginCNodeUser(const KernelGraphPtr &main_gra // update_state_node, broadcat_node and load_node to keep order. AnfNodePtrList load_inputs = {NewValueNode(prim::kPrimLoad), broadcast_to_node, update_state_node}; auto load_node = main_graph->NewCNode(load_inputs); + load_node->set_abstract(broadcast_to_node->abstract()); main_graph->AddNode(load_node); auto user_cnode = user_node->cast(); MS_EXCEPTION_IF_NULL(user_cnode); user_cnode->set_input(IntToSize(index), load_node); - (void)to_process_order_.emplace_back(composite_node, user_node); } } diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.h b/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.h index e7cf1e4209d..4a868ceeb5a 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.h +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.h @@ -46,9 +46,8 @@ class AtomicAddChecker { virtual bool SuitableForAtomicAdd(const AnfNodePtr &node) { return false; } virtual bool FindCandidate(const AnfNodePtr &anf_node); virtual bool CanActivateAtomicAdd(const AnfNodePtr &anf_node); - - private: AtomicAddInfo atomic_add_info_; + PrimitivePtr target_type_{prim::kPrimReduceSum}; }; class AtomicAddCheckerGPU : public AtomicAddChecker { @@ -78,29 +77,28 @@ class AtomicCleanInsertter : public Pass { protected: virtual void CorrectKernelBuildInfo(const AnfNodePtr &composite_node, const AnfNodePtr &new_input); virtual void ProcessOriginCNode(const AnfNodePtr &composite_node, const AnfNodePtr &new_input); + virtual CNodePtr CreateAtomicCleanCompositeNode(const KernelGraphPtr &main_graph, TypeId dst_type); void AddDepend(const FuncGraphPtr &main_graph, const AnfNodePtr &clean_node, const AnfNodePtr &composite_node, const AnfNodePtr &user_node, int index) const; void InsertAtomicClean(const KernelGraphPtr &main_graph, const AnfNodePtr &anf_node, const FuncGraphManagerPtr &mng); CNodePtr InsertUpdateState(const KernelGraphPtr &main_graph, const CNodePtr &composite_node) const; - CNodePtr atomic_add_node_{nullptr}; - - private: void CorrectAbstract(const AnfNodePtr &composite_node) const; - CNodePtr CreateAtomicCleanCompositeNode(const KernelGraphPtr &main_graph, TypeId dst_type); void CreateInplaceAssignNodeAndCorrectReturn(const FuncGraphPtr &sub_graph, const AnfNodePtr &new_parameter); void ProcessOriginCNodeUser(const KernelGraphPtr &main_graph, const AnfNodePtr &composite_node, const AnfNodePtr &broadcast_to_node, const AnfNodePtr &update_state_node, const FuncGraphManagerPtr &mng); + + CNodePtr atomic_add_node_{nullptr}; + size_t reduce_real_output_index_{0}; + size_t real_output_num_{0}; + + private: std::vector> FindOriginCNodeUsers(const KernelGraphPtr &main_graph, const AnfNodePtr &composite_node, const FuncGraphManagerPtr &mng, bool correct_index) const; bool IsExistStructuralObstacle(const KernelGraphPtr &main_graph, const AnfNodePtr &node, const FuncGraphManagerPtr &mng); - - size_t reduce_real_output_index_{0}; - size_t real_output_num_{0}; - std::vector> to_process_order_; }; using AtomicCleanInsertterPtr = std::shared_ptr; } // namespace opt diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc index 40ec3d2593d..00bd2afbb30 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc @@ -66,14 +66,6 @@ bool IsMakeTupleOut(const AnfNodePtr &out, AnfNodePtrList *real_outs) { return false; } -AbstractBasePtr GetOutputAbstract(const AnfNodePtr &node, size_t output_idx) { - auto out_spec = node->abstract(); - if (out_spec->isa()) { - return out_spec->cast()->elements()[output_idx]; - } - return out_spec; -} - AnfNodePtrList EliminateMakeTuple(const FuncGraphPtr &fg, const FuncGraphManagerPtr &mng) { AnfNodePtrList outs; auto out_node = fg->output(); @@ -186,6 +178,14 @@ void ReplaceTensorWithScalar(const FuncGraphPtr &fg, const std::vectorabstract(); + if (out_spec->isa()) { + return out_spec->cast()->elements()[output_idx]; + } + return out_spec; +} + bool ConvertNonscalarTensorToParameter(const FuncGraphPtr &fg, AnfNodePtrList *inputs_ptr) { MS_EXCEPTION_IF_NULL(inputs_ptr); auto nodes = TopoSort(fg->get_return()); @@ -612,13 +612,7 @@ void ResetKernelInfo(const AnfNodePtr &node, KernelType kernel_type) { #endif } -std::string GetFormat(const AnfNodePtr &node) { - auto kernel_info = dynamic_cast(node->kernel_info()); - MS_EXCEPTION_IF_NULL(kernel_info); - auto kernel_build_info = kernel_info->select_kernel_build_info(); - MS_EXCEPTION_IF_NULL(kernel_build_info); - return kernel_build_info->GetOutputFormat(0); -} +std::string GetFormat(const AnfNodePtr &node) { return AnfAlgo::GetOutputFormat(node, 0); } TypePtr GetType(const AnfNodePtr &node) { const auto &abstract = node->abstract(); diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h index 68bdae3638d..9ed232d92ff 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h @@ -133,6 +133,8 @@ ValueNodePtr CreateScalarTensorValueNode(const DataInfo &info, T value, size_t d return new_value_node; } +AbstractBasePtr GetOutputAbstract(const AnfNodePtr &node, size_t output_idx); + // functions to graphkernel model graphkernel::LiteGraphPtr AnfGraph2LiteGraph(const FuncGraphPtr &func_graph); FuncGraphPtr LiteGraph2AnfGraph(const graphkernel::LiteGraphPtr &lite_graph, AnfNodePtrList *outputs = nullptr); diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_optimization.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_optimization.cc index 30e160ee01d..af6d6246dfe 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_optimization.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_optimization.cc @@ -44,6 +44,8 @@ #include "backend/optimizer/graph_kernel/axis_normalizer.h" #include "backend/optimizer/graph_kernel/decrease_compute_precision.h" #include "backend/optimizer/graph_kernel/decrease_transfer_precision.h" +#include "backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h" +#include "backend/optimizer/graph_kernel/uss_atomic_add.h" #include "backend/optimizer/pass/getitem_tuple.h" #include "backend/optimizer/graph_kernel/graph_kernel_pass_manager.h" #include "backend/optimizer/graph_kernel/rewrite_output_shape.h" @@ -162,6 +164,11 @@ PassManagerPtr GraphKernelOptimizer::HighLevelOpt2() const { auto level_low_precision = GetPassLevelByFlag(context::GraphKernelFlags::GetInstance().enable_low_precision); pm->AddPass(std::make_shared(), level_low_precision); pm->AddPass(std::make_shared(), level_low_precision, is_ascend); + + // Enable tsa and uss + pm->AddPass(std::make_shared(), OptLevel_1); + pm->AddPass(std::make_shared(), OptLevel_1); + return pm; } diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.cc new file mode 100644 index 00000000000..f65db5cdbbb --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.cc @@ -0,0 +1,270 @@ +/** + * 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/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "base/core_ops.h" +#include "ir/tensor.h" +#include "utils/utils.h" +#include "utils/log_adapter.h" +#include "backend/kernel_compiler/kernel.h" +#include "backend/kernel_compiler/common_utils.h" +#include "backend/optimizer/graph_kernel/graph_kernel_helper.h" +#include "backend/session/anf_runtime_algorithm.h" +#include "backend/session/kernel_graph.h" +#include "debug/anf_ir_dump.h" + +namespace mindspore { +namespace opt { + +class TsaChecker : public AtomicAddChecker { + public: + explicit TsaChecker(const PrimitivePtr &target) { target_type_ = target; } + virtual ~TsaChecker() = default; + + private: + bool CanActivateAtomicAdd(const AnfNodePtr &anf_node) override { + if (!FindCandidate(anf_node)) { + return false; + } + + auto tsa_cnode = atomic_add_info_.atomic_add_node; + if (!utils::isa(tsa_cnode->input(1))) { + return false; + } + + return true; + } +}; + +AnfNodePtr TsaAtomicAddToFirstTensor::FindTsaFirstRealInputInGraph(const KernelGraphPtr &main_graph, + const AnfNodePtr &node) { + auto cnode = node->cast(); + MS_EXCEPTION_IF_NULL(cnode); + auto sub_graph = AnfAlgo::GetCNodeFuncGraphPtr(cnode); + auto mng_sub = sub_graph->manager(); + if (mng_sub == nullptr) { + mng_sub = Manage(sub_graph, false); + sub_graph->set_manager(mng_sub); + } + + auto first_input = atomic_add_node_->input(1)->cast(); + MS_EXCEPTION_IF_NULL(first_input); + auto parameters = sub_graph->parameters(); + bool hit = false; + for (size_t i = 0; i < parameters.size(); ++i) { + if (parameters[i] == first_input) { + tsa_first_input_index_ = i; + hit = true; + break; + } + } + if (!hit) { + MS_LOG(EXCEPTION) << "Cannot find tensor scatter add first input in sub-graph parameters!"; + } + + return cnode->input(tsa_first_input_index_ + 1); // CNode input have a primitive, so add 1. +} + +AnfNodePtr TsaAtomicAddToFirstTensor::ProcessTsaFirstNode(const KernelGraphPtr &main_graph, const AnfNodePtr &node) { + auto mng = main_graph->manager(); + if (mng == nullptr) { + mng = Manage(main_graph, true); + main_graph->set_manager(mng); + } + // find first input of tsa + auto tsa_first_input = FindTsaFirstRealInputInGraph(main_graph, node); + auto users = mng->node_users()[tsa_first_input]; + + if (users.size() == 1 && !(utils::isa(tsa_first_input) || utils::isa(tsa_first_input))) { + return tsa_first_input; + } + + // Create composite op's sub-graph. + auto new_sub_graph = std::make_shared(); + auto parameter = new_sub_graph->add_parameter(); + auto kernel_with_index = AnfAlgo::VisitKernel(tsa_first_input, 0); + parameter->set_abstract(GetOutputAbstract(kernel_with_index.first, kernel_with_index.second)); + parameter->set_kernel_info(std::make_shared()); + std::string parameter_format; + TypeId parameter_type; + if (utils::isa(kernel_with_index.first)) { + auto tensor = GetValueNode(kernel_with_index.first); + MS_EXCEPTION_IF_NULL(tensor); + parameter_format = kOpFormat_DEFAULT; + parameter_type = tensor->data_type(); + } else { + parameter_format = AnfAlgo::GetOutputFormat(kernel_with_index.first, kernel_with_index.second); + parameter_type = AnfAlgo::GetOutputDeviceDataType(kernel_with_index.first, kernel_with_index.second); + } + + kernel::KernelBuildInfo::KernelBuildInfoBuilder para_info_builder; + para_info_builder.SetOutputsFormat({parameter_format}); + para_info_builder.SetOutputsDeviceType({parameter_type}); + para_info_builder.SetKernelType(KernelType::AKG_KERNEL); + para_info_builder.SetProcessor(kernel::GetProcessorFromContext()); + AnfAlgo::SetSelectKernelBuildInfo(para_info_builder.Build(), parameter.get()); + + // Create inner op. + auto identity_node = + CreateCNode({NewValueNode(std::make_shared("Reshape")), parameter}, new_sub_graph, + {.format = GetFormat(parameter), .shape = GetShape(parameter), .type = GetType(parameter)}); + SetNodeAttrSafely("shape", MakeValue(GetDeviceShape(parameter)), identity_node); + + // Makeup sub-graph. + new_sub_graph->set_output(identity_node); + auto new_composite_node = main_graph->NewCNode({NewValueNode(new_sub_graph), tsa_first_input}); + new_composite_node->set_abstract(identity_node->abstract()); + SetNewKernelInfo(new_composite_node, new_sub_graph, {tsa_first_input}, {identity_node}); + auto graph_attr = ExtractGraphKernelName(TopoSort(new_sub_graph->get_return()), "", "tsa_identity"); + new_sub_graph->set_attr(FUNC_GRAPH_ATTR_GRAPH_KERNEL, MakeValue(graph_attr)); + new_sub_graph->set_attr("composite_type", MakeValue("tsa_identity")); + + return new_composite_node; +} + +void TsaAtomicAddToFirstTensor::CorrectKernelBuildInfo(const AnfNodePtr &composite_node, + const AnfNodePtr &modified_input) { + // Change kernel build info with modify input + auto kernel_info = static_cast(composite_node->kernel_info()); + MS_EXCEPTION_IF_NULL(kernel_info); + const auto &origin_kernel_build_info = kernel_info->GetMutableSelectKernelBuildInfo(); + auto origin_inputs_format = origin_kernel_build_info->GetAllInputFormats(); + auto origin_outputs_format = origin_kernel_build_info->GetAllOutputFormats(); + auto origin_inputs_type = origin_kernel_build_info->GetAllInputDeviceTypes(); + auto origin_outputs_type = origin_kernel_build_info->GetAllOutputDeviceTypes(); + auto origin_processor = origin_kernel_build_info->processor(); + + std::vector &modified_inputs_format = origin_inputs_format; + std::vector &modified_inputs_type = origin_inputs_type; + std::vector new_outputs_format; + std::vector new_outputs_type; + for (size_t i = 0; i < origin_outputs_format.size(); ++i) { + if (real_output_num_ > 1 && i == reduce_real_output_index_) { + continue; + } + new_outputs_format.push_back(origin_outputs_format[i]); + new_outputs_type.push_back(origin_outputs_type[i]); + } + + auto kernel_with_index = AnfAlgo::VisitKernel(modified_input, 0); + modified_inputs_format[tsa_first_input_index_] = + AnfAlgo::GetOutputFormat(kernel_with_index.first, kernel_with_index.second); + modified_inputs_type[tsa_first_input_index_] = + AnfAlgo::GetOutputDeviceDataType(kernel_with_index.first, kernel_with_index.second); + + kernel::KernelBuildInfo::KernelBuildInfoBuilder new_info_builder; + new_info_builder.SetInputsFormat(modified_inputs_format); + new_info_builder.SetInputsDeviceType(modified_inputs_type); + new_info_builder.SetOutputsFormat(new_outputs_format); + new_info_builder.SetOutputsDeviceType(new_outputs_type); + new_info_builder.SetProcessor(origin_processor); + new_info_builder.SetKernelType(KernelType::AKG_KERNEL); + new_info_builder.SetFusionType(kernel::FusionType::OPAQUE); + auto new_selected_info = new_info_builder.Build(); + AnfAlgo::SetSelectKernelBuildInfo(new_selected_info, composite_node.get()); +} + +void TsaAtomicAddToFirstTensor::ProcessOriginCNode(const AnfNodePtr &composite_node, const AnfNodePtr &outter_node) { + auto sub_graph = AnfAlgo::GetCNodeFuncGraphPtr(composite_node); + auto mng_sub = sub_graph->manager(); + if (mng_sub == nullptr) { + mng_sub = Manage(sub_graph, false); + sub_graph->set_manager(mng_sub); + } + + // modify input + composite_node->cast()->set_input(tsa_first_input_index_ + 1, outter_node); + CreateInplaceAssignNodeAndCorrectReturn(sub_graph, sub_graph->parameters()[tsa_first_input_index_]); + + CorrectAbstract(composite_node); + CorrectKernelBuildInfo(composite_node, outter_node); + + auto old_graph_name = GetValue(sub_graph->get_attr(FUNC_GRAPH_ATTR_GRAPH_KERNEL)); + auto new_graph_name = ExtractGraphKernelName(TopoSort(sub_graph->get_return()), "", "tensor_scatter_add_modified"); + sub_graph->set_attr(FUNC_GRAPH_ATTR_GRAPH_KERNEL, MakeValue(new_graph_name)); + MS_LOG(INFO) << "Convert " << old_graph_name << " to tensor scatter add graph " << new_graph_name; +} + +void TsaAtomicAddToFirstTensor::ProcessTsa(const KernelGraphPtr &main_graph, const AnfNodePtr &anf_node, + const FuncGraphManagerPtr &mng) { + auto origin_composite_node = anf_node->cast(); + MS_EXCEPTION_IF_NULL(origin_composite_node); + + // Create identity node. + auto outter_node = ProcessTsaFirstNode(main_graph, anf_node); + + // Insert extra input(broadcast node output) to composite node, and make origin TensorScatterAdd inplaceassign to it. + // Note: if it's single output, this will increase total memory because of a fake out. + ProcessOriginCNode(origin_composite_node, outter_node); + + // Insert update_state_node to keep execution order. + auto update_state_node = InsertUpdateState(main_graph, origin_composite_node); + + // Replace origin ReduceSum's user with atomic clean output + ProcessOriginCNodeUser(main_graph, origin_composite_node, outter_node, update_state_node, mng); + MS_LOG(INFO) << "Target node: " << origin_composite_node->fullname_with_scope() + << ", outer node: " << outter_node->fullname_with_scope(); +} + +bool TsaAtomicAddToFirstTensor::Run(const FuncGraphPtr &func_graph) { + auto kernel_graph = std::dynamic_pointer_cast(func_graph); + MS_EXCEPTION_IF_NULL(kernel_graph); + auto mng = kernel_graph->manager(); + if (mng == nullptr) { + mng = Manage(kernel_graph, true); + kernel_graph->set_manager(mng); + } + + bool changed = false; + std::shared_ptr atomic_add_checker = + std::make_shared(std::make_shared("TensorScatterAdd")); + if (atomic_add_checker == nullptr) { + return changed; + } + + auto topo_nodes = TopoSort(kernel_graph->get_return()); + for (const auto &node : topo_nodes) { + if (!atomic_add_checker->Check(node)) { + continue; + } + auto atomic_add_info = atomic_add_checker->GetAtomicAddInfo(); + atomic_add_node_ = atomic_add_info.atomic_add_node; + reduce_real_output_index_ = atomic_add_info.reduce_real_output_index; + real_output_num_ = atomic_add_info.real_output_num; + ProcessTsa(kernel_graph, node, mng); + changed = true; + } + + if (changed) { + mng->RemoveRoots(); + mng->KeepRoots({func_graph}); + } + + return changed; +} +} // namespace opt +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h b/mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h new file mode 100644 index 00000000000..1804c5a4234 --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h @@ -0,0 +1,64 @@ +/** + * 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_OPTIMIZER_GRAPH_KERNEL_TSA_ATOMIC_ADD_TO_FIRST_TENSOR_H_ +#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_TSA_ATOMIC_ADD_TO_FIRST_TENSOR_H_ + +#include +#include +#include +#include +#include "backend/optimizer/common/optimizer.h" +#include "backend/optimizer/graph_kernel/add_atomic_clean.h" +#include "backend/session/kernel_graph.h" + +namespace mindspore { +namespace opt { +/* + * output = SubGraph(input_x, indices, update) { + * %0 = TensorScatterAdd(%para1, %para2, %para3) + * return %0 + * } + * ----------------------------------------------------------------> + * // Initialize output with input_x. + * output = Reshape(input_x) + * fake_out = SubGraph'(output, indices, update) { + * %0 = TensorScatterAdd(%para1, %para2, %para3) + * %1 = InplaceAssign(%para1, %0, %0) // attrs{"fake_output":true} + * return %1 + * } + */ +class TsaAtomicAddToFirstTensor : public AtomicCleanInsertter { + public: + TsaAtomicAddToFirstTensor() : AtomicCleanInsertter("tensor_scatter_add_atomic_add_to_first_tensor") {} + ~TsaAtomicAddToFirstTensor() override = default; + + bool Run(const FuncGraphPtr &func_graph) override; + + private: + void ProcessOriginCNode(const AnfNodePtr &composite_node, const AnfNodePtr &new_input) override; + void CorrectKernelBuildInfo(const AnfNodePtr &composite_node, const AnfNodePtr &new_input) override; + void ProcessTsa(const KernelGraphPtr &main_graph, const AnfNodePtr &anf_node, const FuncGraphManagerPtr &mng); + AnfNodePtr ProcessTsaFirstNode(const KernelGraphPtr &main_graph, const AnfNodePtr &node); + AnfNodePtr FindTsaFirstRealInputInGraph(const KernelGraphPtr &main_graph, const AnfNodePtr &node); + + size_t tsa_first_input_index_{0}; // sub-graph parameter index. +}; +using TsaAtomicAddToFirstTensorPtr = std::shared_ptr; +} // namespace opt +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_TSA_ATOMIC_ADD_TO_FIRST_TENSOR_H_ diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.cc new file mode 100644 index 00000000000..2c94628dd44 --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.cc @@ -0,0 +1,89 @@ +/** + * 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/optimizer/graph_kernel/uss_atomic_add.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "base/core_ops.h" +#include "ir/tensor.h" +#include "utils/utils.h" +#include "utils/log_adapter.h" +#include "backend/kernel_compiler/kernel.h" +#include "backend/kernel_compiler/common_utils.h" +#include "backend/optimizer/graph_kernel/graph_kernel_helper.h" +#include "backend/session/anf_runtime_algorithm.h" +#include "backend/session/kernel_graph.h" +#include "debug/anf_ir_dump.h" + +namespace mindspore { +namespace opt { +class UssChecker : public AtomicAddChecker { + public: + explicit UssChecker(const PrimitivePtr &target) { target_type_ = target; } + virtual ~UssChecker() = default; + + private: + bool CanActivateAtomicAdd(const AnfNodePtr &anf_node) override { return FindCandidate(anf_node); } +}; + +bool UssAtomicAdd::Run(const FuncGraphPtr &func_graph) { + auto kernel_graph = std::dynamic_pointer_cast(func_graph); + MS_EXCEPTION_IF_NULL(kernel_graph); + auto mng = kernel_graph->manager(); + if (mng == nullptr) { + mng = Manage(kernel_graph, true); + kernel_graph->set_manager(mng); + } + + bool changed = false; + std::shared_ptr atomic_add_checker = + std::make_shared(std::make_shared("UnsortedSegmentSum")); + if (atomic_add_checker == nullptr) { + return changed; + } + + auto topo_nodes = TopoSort(kernel_graph->get_return()); + for (const auto &node : topo_nodes) { + if (!atomic_add_checker->Check(node)) { + continue; + } + auto atomic_add_info = atomic_add_checker->GetAtomicAddInfo(); + atomic_add_node_ = atomic_add_info.atomic_add_node; + reduce_real_output_index_ = atomic_add_info.reduce_real_output_index; + real_output_num_ = atomic_add_info.real_output_num; + InsertAtomicClean(kernel_graph, node, mng); + changed = true; + } + + if (changed) { + mng->RemoveRoots(); + mng->KeepRoots({func_graph}); + } + + return changed; +} + +} // namespace opt +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.h b/mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.h new file mode 100644 index 00000000000..b9dd6ed1235 --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.h @@ -0,0 +1,54 @@ +/** + * 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_OPTIMIZER_GRAPH_KERNEL_USS_ATOMIC_ADD_H_ +#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_USS_ATOMIC_ADD_H_ + +#include +#include +#include +#include +#include "backend/optimizer/common/optimizer.h" +#include "backend/optimizer/graph_kernel/add_atomic_clean.h" +#include "backend/session/kernel_graph.h" + +namespace mindspore { +namespace opt { +/* + * output = SubGraph(input_x, segment_ids) { + * %0 = UnsortedSegmentSum(%para1, %para2) + * return %0 + * } + * ----------------------------------------------------------------> + * // Clean output with zero. + * output = broadcast_to(0.0) // attrs{"shape": [shape of origin output.]} + * fake_out = SubGraph'(input_x, segment_ids, output) { + * %0 = UnsortedSegmentSum(%para1, %para2) + * %1 = InplaceAssign(%para3, %0, %0) // attrs{"fake_output":true} + * return %1 + * } + */ +class UssAtomicAdd : public AtomicCleanInsertter { + public: + UssAtomicAdd() : AtomicCleanInsertter("unsorted_segment_sum_atomic_add_process") {} + ~UssAtomicAdd() override = default; + bool Run(const FuncGraphPtr &func_graph) override; +}; +using UssAtomicAddPtr = std::shared_ptr; +} // namespace opt +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_USS_ATOMIC_ADD_H_