1 /**
2 * Copyright 2020-2022 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "frontend/optimizer/graph_transform.h"
18 #include <vector>
19 #include <algorithm>
20 #include "mindspore/core/ops/sequence_ops.h"
21 #include "ir/graph_utils.h"
22
23 namespace mindspore {
24 /* namespace to support opt */
25 namespace opt {
ContainSparseTensor(const abstract::AbstractBasePtr & abs)26 bool ContainSparseTensor(const abstract::AbstractBasePtr &abs) {
27 MS_EXCEPTION_IF_NULL(abs);
28 if (abs->isa<abstract::AbstractSparseTensor>()) {
29 return true;
30 }
31 if (abs->isa<abstract::AbstractTuple>()) {
32 auto vec = abs->cast_ptr<abstract::AbstractTuple>()->elements();
33 return std::any_of(vec.begin(), vec.end(), ContainSparseTensor);
34 }
35 return false;
36 }
37
IsSequence(const AnfNodePtr & param)38 bool IsSequence(const AnfNodePtr ¶m) {
39 return param->abstract() != nullptr && param->abstract()->isa<abstract::AbstractSequence>();
40 }
41
IsSequenceExpandable(const AbstractBasePtr & abs)42 bool IsSequenceExpandable(const AbstractBasePtr &abs) {
43 if (abs == nullptr) {
44 return false;
45 }
46 auto abs_seq = abs->cast<abstract::AbstractSequencePtr>();
47 if (abs_seq == nullptr) {
48 return false;
49 }
50 // Not expand SparseTensor.
51 if (common::AnfAlgo::CheckAbsSparseTensor(abs)) {
52 return false;
53 }
54 // Not expand dynamic len sequence.
55 if (abs_seq->dynamic_len()) {
56 return false;
57 }
58 // If arg is a empty tuple or list, should not expand it.
59 if (abs_seq->empty()) {
60 return false;
61 }
62 // Not expand node which is arg of dynamic len parameter.
63 return !abs_seq->dyn_len_arg();
64 }
65
ParamContainSparseTensor(const AnfNodePtr & param)66 bool ParamContainSparseTensor(const AnfNodePtr ¶m) {
67 // If SparseTensor, Tuple(SparseTensor,...) or Tuple(...,(..., SparseTensor)), return false and skip this pass.
68 return param->abstract() != nullptr && ContainSparseTensor(param->abstract());
69 }
70
FuncGraphHasSequenceInput(const FuncGraphPtr & fg)71 bool FuncGraphHasSequenceInput(const FuncGraphPtr &fg) {
72 return std::any_of(fg->parameters().cbegin(), fg->parameters().cend(), IsSequence);
73 }
74
FuncGraphHasConstantSequenceInput(const FuncGraphPtr & fg)75 bool FuncGraphHasConstantSequenceInput(const FuncGraphPtr &fg) {
76 return std::any_of(fg->parameters().cbegin(), fg->parameters().cend(),
77 [](const AnfNodePtr ¶m) { return IsSequenceExpandable(param->abstract()); });
78 }
79
TransformSequenceArgument(const FuncGraphPtr & fg,const AnfNodePtr & node,const abstract::AbstractSequencePtr & abs)80 std::vector<AnfNodePtr> TransformSequenceArgument(const FuncGraphPtr &fg, const AnfNodePtr &node,
81 const abstract::AbstractSequencePtr &abs) {
82 auto &elements = abs->elements();
83 std::vector<AnfNodePtr> sequence_node_expanded;
84 for (size_t i = 0; i < elements.size(); i++) {
85 auto idx = NewValueNode(SizeToLong(i));
86 auto abstract_scalar = std::make_shared<abstract::AbstractScalar>(std::make_shared<Int64Imm>(SizeToLong(i)));
87 idx->set_abstract(abstract_scalar);
88 auto elem_node = fg->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx});
89 elem_node->set_abstract(elements[i]);
90 if (IsSequenceExpandable(elements[i])) {
91 auto nodes = TransformSequenceArgument(fg, elem_node, elements[i]->cast<abstract::AbstractSequencePtr>());
92 (void)sequence_node_expanded.insert(sequence_node_expanded.cend(), nodes.cbegin(), nodes.cend());
93 } else {
94 sequence_node_expanded.push_back(elem_node);
95 }
96 }
97 return sequence_node_expanded;
98 }
99 } // namespace opt
100 } // namespace mindspore
101