[locomotiv] Support S32 type for ConstGen (#3656)

* [locomotiv] Support S32 type for ConstGen

This will enable S32 type for ConstGen node

Signed-off-by: SaeHie Park <saehie.park@samsung.com>
* fix format

diff --git a/contrib/locomotiv/src/Node/ConstGen.cpp b/contrib/locomotiv/src/Node/ConstGen.cpp
index d55a90e..ad6efe6 100644 (file)
--- a/contrib/locomotiv/src/Node/ConstGen.cpp
+++ b/contrib/locomotiv/src/Node/ConstGen.cpp
@@ -71,6 +71,22 @@ void NodeExecution::execute(loco::ConstGen *constgen)
 
   switch (constgen->dtype())
   {
+  case loco::DataType::S32:
+  {
+    assert(volume == constgen->size<loco::DataType::S32>());
+
+    auto buf = make_buffer<int32_t, LexicalLayout>(shape);
+
+    for (IndexEnumerator e{shape}; e.valid(); e.advance())
+    {
+      const auto &index = e.current();
+      uint32_t offset = ::offset_by_index(shape, index);
+      buf.at(index) = constgen->at<loco::DataType::S32>(offset);
+    }
+
+    data = locomotiv::make_data(buf);
+    break;
+  }
   case loco::DataType::FLOAT32:
   {
     assert(volume == constgen->size<loco::DataType::FLOAT32>());

diff --git a/contrib/locomotiv/src/Node/ConstGen.test.cpp b/contrib/locomotiv/src/Node/ConstGen.test.cpp
index e099d56..5128be9 100644 (file)
--- a/contrib/locomotiv/src/Node/ConstGen.test.cpp
+++ b/contrib/locomotiv/src/Node/ConstGen.test.cpp
@@ -30,6 +30,38 @@ using nncc::core::ADT::tensor::Shape;
 using nncc::core::ADT::tensor::LexicalLayout;
 using nncc::core::ADT::tensor::make_buffer;
 
+TEST(NodeExecution_ConstGen, s32)
+{
+  // Make ConstGen node
+  loco::ConstGen constgen;
+
+  constgen.dtype(loco::DataType::S32);
+  constgen.shape({2, 3});
+  constgen.size<loco::DataType::S32>(6);
+
+  constgen.at<loco::DataType::S32>(0) = 0;  // Set 0,0
+  constgen.at<loco::DataType::S32>(1) = 1;  // Set 0,1
+  constgen.at<loco::DataType::S32>(2) = 2;  // Set 0,2
+  constgen.at<loco::DataType::S32>(3) = -3; // Set 1,0
+  constgen.at<loco::DataType::S32>(4) = -4; // Set 1,1
+  constgen.at<loco::DataType::S32>(5) = -5; // Set 1,2
+
+  // run execution
+  locomotiv::NodeExecution::get().run(&constgen);
+
+  // test
+  auto data = locomotiv::annot_data(&constgen);
+  ASSERT_NE(data, nullptr);
+  ASSERT_EQ(data->dtype(), loco::DataType::S32);
+  ASSERT_EQ(*data->shape(), Shape({2, 3}));
+  ASSERT_FLOAT_EQ(data->as_s32_bufptr()->at(Index{0, 0}), 0);
+  ASSERT_FLOAT_EQ(data->as_s32_bufptr()->at(Index{0, 1}), 1);
+  ASSERT_FLOAT_EQ(data->as_s32_bufptr()->at(Index{0, 2}), 2);
+  ASSERT_FLOAT_EQ(data->as_s32_bufptr()->at(Index{1, 0}), -3);
+  ASSERT_FLOAT_EQ(data->as_s32_bufptr()->at(Index{1, 1}), -4);
+  ASSERT_FLOAT_EQ(data->as_s32_bufptr()->at(Index{1, 2}), -5);
+}
+
 TEST(NodeExecution_ConstGen, f32)
 {
   // Make ConstGen node