• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1From 9eef7081d9d73994498eb0dde7aa0c1975c8902a Mon Sep 17 00:00:00 2001
2From: Zhu Guodong <zhuguodong0001@163.com>
3Date: Fri, 7 Apr 2023 14:48:57 +0800
4Subject: [PATCH] mindir add custom op
5
6Signed-off-by: Zhu Guodong <zhuguodong0001@163.com>
7---
8 mindspore/lite/mindir/include/mindir.h       |  8 +++++
9 mindspore/lite/mindir/include/mindir_types.h |  1 +
10 mindspore/lite/mindir/src/mindir.cc          | 36 ++++++++++++++++++++
11 3 files changed, 45 insertions(+)
12
13diff --git a/mindspore/lite/mindir/include/mindir.h b/mindspore/lite/mindir/include/mindir.h
14index 73cd6898..ae72cc9b 100644
15--- a/mindspore/lite/mindir/include/mindir.h
16+++ b/mindspore/lite/mindir/include/mindir.h
17@@ -30,6 +30,10 @@ struct Model;
18 }  // namespace OHOS
19
20 namespace mindspore {
21+namespace schema {
22+struct Attribute;
23+}
24+
25 namespace lite {
26
27 // ********** Model **********
28@@ -418,6 +422,10 @@ PrimitivePtr MindIR_Unsqueeze_CreatePrimitive(const std::vector<int64_t> &axis);
29 std::vector<int64_t> MindIR_Unsqueeze_GetAxis(ConstPrimitivePtr primitive);
30 void MindIR_Unsqueeze_SetAxis(PrimitivePtr *primitive, const std::vector<int64_t> &axis);
31
32+// ********** Custom **********
33+std::vector<const mindspore::schema::Attribute *> MindIR_Custom_GetAttr(ConstPrimitivePtr primitive);
34+std::string MindIR_Attribute_GetName(const mindspore::schema::Attribute &attr);
35+std::vector<uint8_t> MindIR_Attribute_GetData(const mindspore::schema::Attribute &attr);
36 }  // namespace lite
37 }  // namespace mindspore
38 #endif
39diff --git a/mindspore/lite/mindir/include/mindir_types.h b/mindspore/lite/mindir/include/mindir_types.h
40index 8f2a9c70..ad272f8e 100644
41--- a/mindspore/lite/mindir/include/mindir_types.h
42+++ b/mindspore/lite/mindir/include/mindir_types.h
43@@ -101,6 +101,7 @@ enum NodeType : uint32_t {
44   NODE_TYPE_TOPK_FUSION = 161,
45   NODE_TYPE_TRANSPOSE = 162,
46   NODE_TYPE_UNSQUEEZE = 165,
47+  NODE_TYPE_CUSTOM = 191,
48 };
49
50 enum ResizeMethod : int8_t {
51diff --git a/mindspore/lite/mindir/src/mindir.cc b/mindspore/lite/mindir/src/mindir.cc
52index c2a1cd3f..dd249738 100644
53--- a/mindspore/lite/mindir/src/mindir.cc
54+++ b/mindspore/lite/mindir/src/mindir.cc
55@@ -4254,5 +4254,41 @@ void MindIR_Unsqueeze_SetAxis(PrimitivePtr *primitive, const std::vector<int64_t
56   }
57 }
58
59+// ********** Custom **********
60+std::vector<const mindspore::schema::Attribute *> MindIR_Custom_GetAttr(ConstPrimitivePtr primitive) {
61+  if (primitive == nullptr) {
62+    return {};
63+  }
64+  auto prim = static_cast<const schema::Primitive *>(primitive);
65+  auto value = prim->value_as_Custom();
66+  if (value == nullptr) {
67+    return {};
68+  }
69+  std::vector<const mindspore::schema::Attribute *> result;
70+  if (value->attr() == nullptr) {
71+    return {};
72+  }
73+  for (auto attr: *(value->attr())) {
74+    result.push_back(attr);
75+  }
76+  return result;
77+}
78+
79+std::string MindIR_Attribute_GetName(const mindspore::schema::Attribute &attr) {
80+  if (attr.name() == nullptr) {
81+    return "";
82+  }
83+  return attr.name()->str();
84+}
85+
86+std::vector<uint8_t> MindIR_Attribute_GetData(const mindspore::schema::Attribute &attr) {
87+  if (attr.data() == nullptr) {
88+    return {};
89+  }
90+  std::vector<uint8_t> tmp_data(attr.data()->size());
91+  std::transform(attr.data()->begin(), attr.data()->end(), tmp_data.begin(),
92+                  [](uint8_t item) { return item; });
93+  return tmp_data;
94+}
95 }  // namespace lite
96 }  // namespace mindspore
97--
982.34.1
99
100