1 /**
2 * Copyright 2021 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 #include "backend/optimizer/graph_kernel/split_umonad.h"
17
18 #include <vector>
19 #include <string>
20 #include <algorithm>
21 #include <memory>
22
23 #include "base/core_ops.h"
24 #include "utils/utils.h"
25 #include "runtime/device/kernel_info.h"
26 #include "backend/optimizer/common/helper.h"
27
28 namespace mindspore {
29 namespace opt {
DefinePattern() const30 const BaseRef SplitAssign::DefinePattern() const {
31 VarPtr v = std::make_shared<Var>();
32 VarPtr Xs = std::make_shared<Var>();
33 VarPtr Us = std::make_shared<Var>();
34 VarPtr UMonad = std::make_shared<Var>();
35 return VectorRef({v, Xs, Us, UMonad});
36 }
37
CanSplit(const AnfNodePtr & node)38 bool CanSplit(const AnfNodePtr &node) { return IsPrimitiveCNode(node, prim::kPrimAssign); }
39
ProcessNode(const FuncGraphPtr & func_graph,const AnfNodePtr & node,size_t input_idx)40 AnfNodePtr ProcessNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node, size_t input_idx) {
41 MS_EXCEPTION_IF_NULL(node);
42 CNodePtr cnode = node->cast<CNodePtr>();
43 MS_EXCEPTION_IF_NULL(cnode);
44
45 // Get original op's abstract and inputs
46 AbstractBasePtr original_abstract = cnode->abstract()->Clone();
47 auto original_inputs = cnode->inputs();
48
49 // Create depend node
50 AnfNodePtrList depend_inputs = {NewValueNode(prim::kPrimDepend), original_inputs[input_idx], original_inputs.back()};
51 auto depend_cnode = func_graph->NewCNode(depend_inputs);
52 depend_cnode->set_abstract(original_inputs[input_idx]->abstract());
53 depend_cnode->set_kernel_info(std::make_shared<device::KernelInfo>());
54 // Create new node, delete U from inputs.
55 AnfNodePtrList new_inputs = {cnode->input(0)};
56 for (size_t i = 1; i + 1 < cnode->size(); i++) {
57 if (i == input_idx) {
58 new_inputs.push_back(depend_cnode);
59 } else {
60 new_inputs.push_back(cnode->input(i));
61 }
62 }
63 auto new_cnode = func_graph->NewCNode(new_inputs);
64 new_cnode->set_abstract(original_abstract);
65 new_cnode->set_kernel_info(cnode->kernel_info_ptr());
66 return new_cnode;
67 }
68
Process(const FuncGraphPtr & func_graph,const AnfNodePtr & node,const EquivPtr &) const69 const AnfNodePtr SplitAssign::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node, const EquivPtr &) const {
70 MS_EXCEPTION_IF_NULL(node);
71 if (!CanSplit(node)) return node;
72 return ProcessNode(node->func_graph(), node, 1);
73 }
74
Run(const AnfNodePtr & node)75 AnfNodePtr OpUMonadExpander::Run(const AnfNodePtr &node) {
76 auto cnode = node->cast<CNodePtr>();
77 MS_EXCEPTION_IF_NULL(cnode);
78 // assume the UMonad node is the last input
79 if (cnode->size() > 1 && HasAbstractUMonad(cnode->inputs().back())) {
80 auto new_node = ProcessNode(node->func_graph(), node, input_idx_);
81 return DefaultExpander::Run(new_node);
82 }
83 return DefaultExpander::Run(node);
84 }
85 } // namespace opt
86 } // namespace mindspore
87