1 /**
2 * Copyright 2023 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_SRC_EXTENDRT_UTILS_CUSTOM_ASCEND_UTILS_H_
18 #define MINDSPORE_LITE_SRC_EXTENDRT_UTILS_CUSTOM_ASCEND_UTILS_H_
19 #include <utility>
20 #include <string>
21 #include <vector>
22 #include <tuple>
23 #include <map>
24 #include <memory>
25
26 #include "ir/anf.h"
27 #include "ir/dtype/type.h"
28 #include "ir/func_graph.h"
29 #include "include/api/types.h"
30 #include "include/api/data_type.h"
31 #include "include/api/status.h"
32 #include "mindspore/ccsrc/kernel/kernel.h"
33 #include "include/common/utils/anfalgo.h"
34 #include "mindspore/core/ops/custom.h"
35 #include "mindspore/lite/src/common/common.h"
36
37 namespace mindspore {
38 struct DynKVCacheSaveInfo {
39 bool batch_size_dyn = false;
40 bool seq_length_dyn = false;
41 std::string kv_cache_layout = lite::kKVCacheLayoutBNSD;
42 };
43
SetKVCacheShape(bool dyn_batch,bool dyn_seq,const std::string & layout,const std::vector<int64_t> & org_shape)44 static inline std::vector<int64_t> SetKVCacheShape(bool dyn_batch, bool dyn_seq, const std::string &layout,
45 const std::vector<int64_t> &org_shape) {
46 std::vector<int64_t> shape = org_shape;
47 if (layout == lite::kKVCacheLayoutBNSD && org_shape.size() == kShape4dDims) {
48 if (dyn_batch) {
49 shape[kDim0] = abstract::Shape::kShapeDimAny;
50 }
51 if (dyn_seq) {
52 shape[kDim2] = abstract::Shape::kShapeDimAny;
53 }
54 } else if (layout == lite::kKVCacheLayoutBSH && org_shape.size() == kShape3dDims) {
55 if (dyn_batch) {
56 shape[kDim0] = abstract::Shape::kShapeDimAny;
57 }
58 if (dyn_seq) {
59 shape[kDim1] = abstract::Shape::kShapeDimAny;
60 }
61 }
62 return shape;
63 }
64
65 class MS_API CustomAscendUtils {
66 public:
67 static bool CreateCustomFuncGraph(const FuncGraphPtr &func_graph, const Buffer &model_cache,
68 const std::string &graph_name, const std::map<std::string, ValuePtr> &attr_map,
69 const std::vector<std::string> &ref_datas,
70 const DynKVCacheSaveInfo &dyn_kv_info = {});
71
72 static bool ParseCustomFuncGraph(const FuncGraphPtr &func_graph, tensor::TensorPtr *model_cache,
73 std::string *graph_name, std::map<std::string, ValuePtr> *attr_map,
74 std::vector<std::pair<std::string, tensor::TensorPtr>> *ref_datas,
75 DynKVCacheSaveInfo *dyn_kv_info = nullptr);
76
77 static bool IsCustomFuncGraph(const FuncGraphPtr &func_graph);
78 static ParameterPtr CreateOmParameter(const FuncGraphPtr &func_graph, const Buffer &om_data,
79 const std::string &graph_name);
80
81 private:
82 std::vector<std::pair<AnfNodePtr, size_t>> outputs_;
83
84 CNodePtr CreateCustomNode(const FuncGraphPtr &func_graph, const ParameterPtr &om_parameter,
85 const std::map<std::string, ValuePtr> &attr_map, const std::vector<std::string> &ref_datas);
86 void SetCustomAttrs(const std::shared_ptr<ops::Custom> &prim, const std::map<std::string, ValuePtr> &attr_map);
87 bool SetCustomOutputs(const FuncGraphPtr &func_graph, const CNodePtr &custom_node);
88 bool ModifyGraphByCustomNode(const FuncGraphPtr &func_graph, const CNodePtr &custom_node);
89 CNodePtr CreateMakeTupleGraphOutput(const FuncGraphPtr &func_graph, const CNodePtr &custom_node);
90 static CNodePtr GetCustomNode(const FuncGraphPtr &func_graph);
91 static bool IsParameterValueZero(const tensor::TensorPtr &tensor);
92
93 static void SetZeroValueRefDatas(const ops::PrimitiveCPtr &primc,
94 const std::vector<std::pair<std::string, tensor::TensorPtr>> &ref_infos);
95 static bool GetZeroValueRefDatas(const ops::PrimitiveCPtr &primc,
96 std::vector<std::pair<std::string, tensor::TensorPtr>> *ref_infos);
97 };
98 } // namespace mindspore
99 #endif // MINDSPORE_LITE_SRC_EXTENDRT_UTILS_CUSTOM_ASCEND_UTILS_H_
100