!14649 [GraphKernel] fix atomic add judge condition for Ascend

From: @looop5
Reviewed-by: @gaoxiong1,@ckey_dou
Signed-off-by: @ckey_dou
pull/14649/MERGE
mindspore-ci-bot 4 years ago committed by Gitee
commit 7268661822

@ -180,23 +180,11 @@ bool AtomicAddCheckerGPU::SuitableForAtomicAdd(const AnfNodePtr &node) {
}
bool AtomicAddCheckerAscend::SuitableForAtomicAdd(const AnfNodePtr &node) {
auto input = node->cast<CNodePtr>()->input(kFirstDataInputIndex);
auto src_shape_vec = GetShape(input);
std::set<int64_t> axis_set = GetUniqReduceAxes(node);
auto dst_shape_vec = AnfAlgo::GetOutputDeviceShape(node, 0);
// case 1: all reduce
if (src_shape_vec.size() == axis_set.size()) {
return true;
}
// case 2: non-reduce axes with dimension 1
for (size_t i = 0; i < src_shape_vec.size(); ++i) {
if (axis_set.find(i) == axis_set.end()) {
if (src_shape_vec[i] != 1) {
return false;
}
}
}
return true;
// all reduce
// non-reduce axes with dimension 1
return std::all_of(dst_shape_vec.cbegin(), dst_shape_vec.cend(), [](const size_t &dim) { return dim == 1; });
}
void AtomicCleanInsertter::CorrectKernelBuildInfo(const AnfNodePtr &composite_node, const AnfNodePtr &new_input) {

Loading…
Cancel
Save