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

support tot ops

This commit is contained in:
tronzhang 2021-08-17 21:12:08 +08:00
parent 45c7bb382b
commit 134ffa9efe
15 changed files with 648 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
akg

@ -1 +1 @@
Subproject commit 77aaa23fc160756bdf40eef97333305680b3b35c
Subproject commit 5196d6e1e8feb1a5ab074541e0c2acf072a6f9b4

3
mindspore/_extends/graph_kernel/expanders/__init__.py
View File

@ -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

43
mindspore/_extends/graph_kernel/expanders/gather.py Normal file
View File

@ -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

61
mindspore/_extends/graph_kernel/model/graph_split.py
View File

@ -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

5
mindspore/_extends/graph_kernel/model/model.py
View File

@ -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)

27
mindspore/_extends/graph_kernel/model/op_infer.py
View File

@ -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!")

13
mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.cc
View File

@ -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<CNodePtr>()->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<CNodePtr>();
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<CNodePtr>();
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<CNodePtr>();
MS_EXCEPTION_IF_NULL(user_cnode);
user_cnode->set_input(IntToSize(index), load_node);
(void)to_process_order_.emplace_back(composite_node, user_node);
}
}

18
mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean.h
View File

@ -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<std::pair<AnfNodePtr, int>> 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<std::pair<AnfNodePtr, AnfNodePtr>> to_process_order_;
};
using AtomicCleanInsertterPtr = std::shared_ptr<AtomicCleanInsertter>;
} // namespace opt

24
mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc
View File

@ -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<abstract::AbstractTuple>()) {
return out_spec->cast<abstract::AbstractTuplePtr>()->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::vector<AnfNodePt
}
} // namespace
AbstractBasePtr GetOutputAbstract(const AnfNodePtr &node, size_t output_idx) {
auto out_spec = node->abstract();
if (out_spec->isa<abstract::AbstractTuple>()) {
return out_spec->cast<abstract::AbstractTuplePtr>()->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<device::KernelInfo *>(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();

2
mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h
View File

@ -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);

7
mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_optimization.cc
View File

@ -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<DecreaseTransferPrecision>(), level_low_precision);
pm->AddPass(std::make_shared<DecreaseComputePrecision>(), level_low_precision, is_ascend);
// Enable tsa and uss
pm->AddPass(std::make_shared<TsaAtomicAddToFirstTensor>(), OptLevel_1);
pm->AddPass(std::make_shared<UssAtomicAdd>(), OptLevel_1);
return pm;
}

270
mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.cc Normal file
View File

@ -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 <algorithm>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <utility>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <vector>
#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<ParameterPtr>(tsa_cnode->input(1))) {
return false;
}
return true;
}
};
AnfNodePtr TsaAtomicAddToFirstTensor::FindTsaFirstRealInputInGraph(const KernelGraphPtr &main_graph,
const AnfNodePtr &node) {
auto cnode = node->cast<CNodePtr>();
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<ParameterPtr>();
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<ValueNodePtr>(tsa_first_input) || utils::isa<ParameterPtr>(tsa_first_input))) {
return tsa_first_input;
}
// Create composite op's sub-graph.
auto new_sub_graph = std::make_shared<FuncGraph>();
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<device::KernelInfo>());
std::string parameter_format;
TypeId parameter_type;
if (utils::isa<ValueNodePtr>(kernel_with_index.first)) {
auto tensor = GetValueNode<tensor::TensorPtr>(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<Primitive>("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<device::KernelInfo *>(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<std::string> &modified_inputs_format = origin_inputs_format;
std::vector<TypeId> &modified_inputs_type = origin_inputs_type;
std::vector<std::string> new_outputs_format;
std::vector<TypeId> 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<CNodePtr>()->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<std::string>(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<CNodePtr>();
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<session::KernelGraph>(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<AtomicAddChecker> atomic_add_checker =
std::make_shared<TsaChecker>(std::make_shared<Primitive>("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

64
mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.h Normal file
View File

@ -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 <memory>
#include <tuple>
#include <utility>
#include <vector>
#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<TsaAtomicAddToFirstTensor>;
} // namespace opt
} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_TSA_ATOMIC_ADD_TO_FIRST_TENSOR_H_

89
mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.cc Normal file
View File

@ -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 <algorithm>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <utility>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <vector>
#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<session::KernelGraph>(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<AtomicAddChecker> atomic_add_checker =
std::make_shared<UssChecker>(std::make_shared<Primitive>("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

54
mindspore/ccsrc/backend/optimizer/graph_kernel/uss_atomic_add.h Normal file
View File

@ -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 <memory>
#include <tuple>
#include <utility>
#include <vector>
#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<UssAtomicAdd>;
} // namespace opt
} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_USS_ATOMIC_ADD_H_