[locomotiv] Introduce NodeDataImpl annotation erase (#3400)

This commit introduces erase_annot_data() function to erase already
annotated NodeDataImpl. It also includes test for added function.

Signed-off-by: Cheongyo Bahk <ch.bahk@samsung.com>

diff --git a/contrib/locomotiv/src/NodeDataImpl.cpp b/contrib/locomotiv/src/NodeDataImpl.cpp
index 9d1fc18..55d7b5a 100644 (file)
--- a/contrib/locomotiv/src/NodeDataImpl.cpp
+++ b/contrib/locomotiv/src/NodeDataImpl.cpp
@@ -50,4 +50,6 @@ const NodeData *annot_data(loco::Node *node)
   return data;
 }
 
+void erase_annot_data(loco::Node *node) { node->annot<NodeDataImpl>(nullptr); }
+
 } // namespace locomotiv

diff --git a/contrib/locomotiv/src/NodeDataImpl.h b/contrib/locomotiv/src/NodeDataImpl.h
index 7802434..012a8c8 100644 (file)
--- a/contrib/locomotiv/src/NodeDataImpl.h
+++ b/contrib/locomotiv/src/NodeDataImpl.h
@@ -52,6 +52,9 @@ void annot_data(loco::Node *node, std::unique_ptr<NodeData> &&data);
 /// @brief Wrapper to get NodeDataImpl annotation of node
 const NodeData *annot_data(loco::Node *node);
 
+/// @brief Erase already annotated NodeDataImpl
+void erase_annot_data(loco::Node *node);
+
 } // namespace locomotiv
 
 #endif // _LOCOMOTIV_NODEDATAIMPL_H_

diff --git a/contrib/locomotiv/src/NodeDataImpl.test.cpp b/contrib/locomotiv/src/NodeDataImpl.test.cpp
index 43a3233..b859560 100644 (file)
--- a/contrib/locomotiv/src/NodeDataImpl.test.cpp
+++ b/contrib/locomotiv/src/NodeDataImpl.test.cpp
@@ -51,4 +51,8 @@ TEST(NodeDataImpl, as_annotation)
   ASSERT_EQ(obtained->dtype(), loco::DataType::FLOAT32);
   ASSERT_EQ(*(obtained->shape()), shape);
   ASSERT_FLOAT_EQ(obtained->as_f32_bufptr()->at(Index{0}), 3.14f);
+
+  // Erase annotation
+  locomotiv::erase_annot_data(node);
+  ASSERT_EQ(locomotiv::annot_data(node), nullptr);
 }