1 /** 2 * Copyright 2024 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 #ifndef MINDSPORE_CCSRC_FRONTEND_OPTIMIZER_IRPASS_J_NODE_AND_USER_REMATCH_H_ 18 #define MINDSPORE_CCSRC_FRONTEND_OPTIMIZER_IRPASS_J_NODE_AND_USER_REMATCH_H_ 19 20 #include <memory> 21 22 #include "frontend/optimizer/optimizer.h" 23 #include "frontend/optimizer/anf_visitor.h" 24 25 namespace mindspore { 26 namespace opt { 27 namespace irpass { 28 // The J node should have only one user. 29 // %0 = J(net) 30 // %1 = %0(x) 31 // %2 = %0(y) 32 // => 33 // %0 = J(net) 34 // %1 = %0(x) 35 // %2 = J(net) 36 // %3 = %2(y) 37 class JNodeAndUserRematch : public AnfVisitor { 38 public: operator()39 AnfNodePtr operator()(const OptimizerPtr &optimizer, const AnfNodePtr &node) override { 40 auto cnode = dyn_cast<CNode>(node); 41 if (cnode == nullptr || cnode->empty()) { 42 return nullptr; 43 } 44 auto input0 = node->cast<CNodePtr>()->input(0); 45 if (!IsPrimitiveCNode(input0, prim::kPrimJ)) { 46 return nullptr; 47 } 48 auto j_cnode = input0->cast<CNodePtr>(); 49 MS_EXCEPTION_IF_NULL(j_cnode); 50 auto manager = optimizer->manager(); 51 MS_EXCEPTION_IF_NULL(manager); 52 const auto &users = manager->node_users()[j_cnode]; 53 if (users.size() <= 1) { 54 return nullptr; 55 } 56 57 auto fg = node->func_graph(); 58 MS_EXCEPTION_IF_NULL(fg); 59 auto new_j_cnode = fg->NewCNodeInOrder(j_cnode->inputs()); 60 new_j_cnode->CloneCNodeInfo(j_cnode); 61 auto inputs = cnode->inputs(); 62 inputs[0] = new_j_cnode; 63 return fg->NewCNodeInOrder(inputs); 64 } 65 }; 66 } // namespace irpass 67 } // namespace opt 68 } // namespace mindspore 69 #endif // MINDSPORE_CCSRC_FRONTEND_OPTIMIZER_IRPASS_J_NODE_AND_USER_REMATCH_H_ 70