1 /** 2 * Copyright 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 #ifndef MINDSPORE_LITE_TOOLS_OPTIMIZER_FUSION_TRANSPOSE_GATHER_FUSION_H 18 #define MINDSPORE_LITE_TOOLS_OPTIMIZER_FUSION_TRANSPOSE_GATHER_FUSION_H 19 20 #include <set> 21 #include <vector> 22 #include "include/backend/optimizer/pass.h" 23 24 namespace mindspore { 25 namespace opt { 26 /* 27 * The subgraph such as the following, in some times, the transpose-op can be fused. 28 * Transpose perm(1, 0, 2) 29 * / | \ 30 * Gather Gather Gather axis(0) 31 * / | \ 32 * Transpose Transpose Transpose perm(1, 0, 2) 33 */ 34 class TransposeGatherFusion : public Pass { 35 public: TransposeGatherFusion()36 TransposeGatherFusion() : Pass("TransposeGatherFusion") {} 37 ~TransposeGatherFusion() override = default; 38 bool Run(const FuncGraphPtr &func_graph) override; 39 40 private: 41 int Process(const FuncGraphPtr &func_graph, const CNodePtr &transpose, std::set<AnfNodePtr> *has_visited); 42 void FindNodes(const FuncGraphPtr &func_graph, const CNodePtr &transpose); 43 bool CheckCanFused(const CNodePtr &transpose); 44 bool CheckCommonAttr(const CNodePtr &transpose); 45 bool CheckIsMatch(const std::vector<int> &pre_perm, const std::vector<int> &post_perm, int axis); 46 std::vector<CNodePtr> gather_nodes_; 47 std::vector<std::vector<CNodePtr>> transpose_nodes_; 48 std::vector<int> gather_updated_axes_; 49 std::vector<int *> gather_axes_data_ptr_; 50 }; 51 } // namespace opt 52 } // namespace mindspore 53 54 #endif // MINDSPORE_LITE_TOOLS_OPTIMIZER_FUSION_TRANSPOSE_GATHER_FUSION_H 55