[luci/export] Introduce clearExportInfo and related

This will introduce clearExportInfo method and related method to clear temporary information added to the graph. ONE-DCO-1.0-Signed-off-by: SaeHie Park <[email protected]>
Samsung · Aug 28, 2024 · ecc7d2b · ecc7d2b
1 parent cc3e6b4
commit ecc7d2b
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 0 deletions.
diff --git a/compiler/luci/export/src/CircleExporterUtils.cpp b/compiler/luci/export/src/CircleExporterUtils.cpp
@@ -293,6 +293,12 @@ void set_tensor_index(loco::Node *node, const CircleTensorIndex &tensor_id)
   node->annot(std::make_unique<CircleTensorIndexAnnotation>(tensor_id));
 }
 
+void clear_tensor_index(loco::Node *node)
+{
+  if (node->annot<CircleTensorIndexAnnotation>() != nullptr)
+    node->annot<CircleTensorIndexAnnotation>(nullptr);
+}
+
 CircleTensorIndex get_tensor_index(loco::Node *node)
 {
   assert(node->annot<CircleTensorIndexAnnotation>() != nullptr);

diff --git a/compiler/luci/export/src/CircleExporterUtils.h b/compiler/luci/export/src/CircleExporterUtils.h
@@ -57,6 +57,7 @@ circle::Padding getOpPadding(const luci::Padding pad);
 using CircleTensorIndex = int32_t;
 
 void set_tensor_index(loco::Node *node, const CircleTensorIndex &tensor_id);
+void clear_tensor_index(loco::Node *node);
 CircleTensorIndex get_tensor_index(loco::Node *node);
 
 } // namespace luci

diff --git a/compiler/luci/export/src/CircleTensorExporter.cpp b/compiler/luci/export/src/CircleTensorExporter.cpp
@@ -694,4 +694,14 @@ void exportOpDefinedTensors(loco::Graph *g, FlatBufferBuilder &builder, Serializ
   }
 }
 
+void clearExportInfo(loco::Graph *g)
+{
+  auto nodes = g->nodes();
+  for (uint32_t n = 0; n < nodes->size(); ++n)
+  {
+    auto node = loco::must_cast<luci::CircleNode *>(nodes->at(n));
+    clear_tensor_index(node);
+  }
+}
+
 } // namespace luci
diff --git a/compiler/luci/export/src/CircleTensorExporter.h b/compiler/luci/export/src/CircleTensorExporter.h
@@ -39,6 +39,11 @@ void prepareModelData(flatbuffers::FlatBufferBuilder &builder, SerializedModelDa
 void exportOpDefinedTensors(loco::Graph *g, flatbuffers::FlatBufferBuilder &builder,
                             SerializedModelData &md, SerializedGraphData &gd);
 
+/**
+ * @brief clear temporary export information annotated to graph nodes
+ */
+void clearExportInfo(loco::Graph *g);
+
 } // namespace luci
 
 #endif // __CIRCLE_TENSOR_EXPORTER_H__