[coco] Record Input-to-Element ID map (#794)

This commit revises Input class to record Tensor Index to Element ID
mapping for input.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>

diff --git a/contrib/coco/core/include/coco/IR/Input.h b/contrib/coco/core/include/coco/IR/Input.h
index 6486234..737c9de 100644 (file)
--- a/contrib/coco/core/include/coco/IR/Input.h
+++ b/contrib/coco/core/include/coco/IR/Input.h
@@ -2,10 +2,13 @@
 #define __COCO_IR_INPUT_H__
 
 #include "coco/IR/Bag.h"
+#include "coco/IR/ElemID.h"
 
 #include <nncc/core/ADT/tensor/Shape.h>
+#include <nncc/core/ADT/tensor/Index.h>
 
 #include <string>
+#include <vector>
 
 namespace coco
 {
@@ -34,6 +37,13 @@ public:
 
 private:
   std::string _name;
+
+private:
+  std::vector<ElemID> _map;
+
+public:
+  ElemID &at(const nncc::core::ADT::tensor::Index &index);
+  const ElemID &at(const nncc::core::ADT::tensor::Index &) const;
 };
 
 } // namespace coco

diff --git a/contrib/coco/core/src/IR/Input.cpp b/contrib/coco/core/src/IR/Input.cpp
index b1a43bd..2de681d 100644 (file)
--- a/contrib/coco/core/src/IR/Input.cpp
+++ b/contrib/coco/core/src/IR/Input.cpp
@@ -1,11 +1,20 @@
 #include "coco/IR/Input.h"
 
+#include <nncc/core/ADT/tensor/LexicalLayout.h>
+
+namespace
+{
+
+const nncc::core::ADT::tensor::LexicalLayout l;
+
+} // namespace
+
 namespace coco
 {
 
 Input::Input(const nncc::core::ADT::tensor::Shape &shape) : _shape{shape}, _bag{nullptr}
 {
-  // DO NOTHING
+  _map.resize(nncc::core::ADT::tensor::num_elements(shape));
 }
 
 void Input::bag(Bag *bag)
@@ -14,4 +23,14 @@ void Input::bag(Bag *bag)
   _bag = bag;
 }
 
+ElemID &Input::at(const nncc::core::ADT::tensor::Index &index)
+{
+  return _map.at(l.offset(_shape, index));
+}
+
+const ElemID &Input::at(const nncc::core::ADT::tensor::Index &index) const
+{
+  return _map.at(l.offset(_shape, index));
+}
+
 } // namespace coco

diff --git a/contrib/coco/core/src/IR/Input.test.cpp b/contrib/coco/core/src/IR/Input.test.cpp
index b93dd7d..705aa18 100644 (file)
--- a/contrib/coco/core/src/IR/Input.test.cpp
+++ b/contrib/coco/core/src/IR/Input.test.cpp
@@ -1,7 +1,12 @@
 #include "coco/IR/Input.h"
 
+#include <nncc/core/ADT/tensor/IndexEnumerator.h>
+
 #include <gtest/gtest.h>
 
+using nncc::core::ADT::tensor::Shape;
+using nncc::core::ADT::tensor::IndexEnumerator;
+
 TEST(IR_INPUT, ctor_should_set_shape)
 {
   const nncc::core::ADT::tensor::Shape shape{1, 3, 3, 1};
@@ -30,3 +35,22 @@ TEST(IR_INPUT, name_update)
   input.name("data");
   ASSERT_EQ(input.name(), "data");
 }
+
+TEST(IR_INPUT, at)
+{
+  const Shape shape{1, 3, 3, 1};
+  coco::Input input{shape};
+
+  coco::Input *mutable_ptr = &input;
+  const coco::Input *immutable_ptr = &input;
+
+  for (IndexEnumerator e{shape}; e.valid(); e.advance())
+  {
+    mutable_ptr->at(e.current()) = coco::ElemID{16};
+  }
+
+  for (IndexEnumerator e{shape}; e.valid(); e.advance())
+  {
+    ASSERT_EQ(immutable_ptr->at(e.current()).value(), 16);
+  }
+}