[coco] Record Kernel Index to Elem ID relation (#804)

This commit revises KernelObject class to record Kernel Index to Elem ID
mapping relation in it.

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

diff --git a/contrib/coco/core/include/coco/IR/KernelObject.h b/contrib/coco/core/include/coco/IR/KernelObject.h
index 7383874..f2681f0 100644 (file)
--- a/contrib/coco/core/include/coco/IR/KernelObject.h
+++ b/contrib/coco/core/include/coco/IR/KernelObject.h
@@ -2,9 +2,12 @@
 #define __COCO_IR_KERNEL_OBJECT_H__
 
 #include "coco/IR/Object.h"
+#include "coco/IR/ElemID.h"
 
 #include <nncc/core/ADT/kernel/Shape.h>
 
+#include <vector>
+
 namespace coco
 {
 
@@ -24,6 +27,13 @@ public:
 
 private:
   nncc::core::ADT::kernel::Shape const _shape;
+
+public:
+  ElemID &at(uint32_t n, uint32_t ch, uint32_t row, uint32_t col);
+  const ElemID &at(uint32_t n, uint32_t ch, uint32_t row, uint32_t col) const;
+
+private:
+  std::vector<ElemID> _map;
 };
 
 } // namespace coco

diff --git a/contrib/coco/core/src/IR/KernelObject.cpp b/contrib/coco/core/src/IR/KernelObject.cpp
index a1e4cda..4217371 100644 (file)
--- a/contrib/coco/core/src/IR/KernelObject.cpp
+++ b/contrib/coco/core/src/IR/KernelObject.cpp
@@ -1,11 +1,30 @@
 #include "coco/IR/KernelObject.h"
 
+#include <nncc/core/ADT/kernel/NCHWLayout.h>
+
+namespace
+{
+
+static nncc::core::ADT::kernel::NCHWLayout l{};
+
+} // namespace
+
 namespace coco
 {
 
 KernelObject::KernelObject(const nncc::core::ADT::kernel::Shape &shape) : _shape{shape}
 {
-  // DO NOTHING
+  _map.resize(nncc::core::ADT::kernel::num_elements(shape));
+}
+
+ElemID &KernelObject::at(uint32_t n, uint32_t ch, uint32_t row, uint32_t col)
+{
+  return _map.at(l.offset(_shape, n, ch, row, col));
+}
+
+const ElemID &KernelObject::at(uint32_t n, uint32_t ch, uint32_t row, uint32_t col) const
+{
+  return _map.at(l.offset(_shape, n, ch, row, col));
 }
 
 } // namespace coco

diff --git a/contrib/coco/core/src/IR/KernelObject.test.cpp b/contrib/coco/core/src/IR/KernelObject.test.cpp
index 4cc6e44..b96125e 100644 (file)
--- a/contrib/coco/core/src/IR/KernelObject.test.cpp
+++ b/contrib/coco/core/src/IR/KernelObject.test.cpp
@@ -13,3 +13,45 @@ TEST(IR_KERNEL_OBJECT, ctor_should_set_size)
   ASSERT_EQ(o.shape().height(), shape.height());
   ASSERT_EQ(o.shape().width(), shape.width());
 }
+
+TEST(IR_KERNEL_OBJECT, at)
+{
+  const uint32_t N = 1;
+  const uint32_t C = 1;
+  const uint32_t H = 3;
+  const uint32_t W = 3;
+
+  const nncc::core::ADT::kernel::Shape shape{N, C, H, W};
+  coco::KernelObject o{shape};
+
+  coco::KernelObject *mutable_ptr = &o;
+  const coco::KernelObject *immutable_ptr = &o;
+
+  for (uint32_t n = 0; n < N; ++n)
+  {
+    for (uint32_t ch = 0; ch < C; ++ch)
+    {
+      for (uint32_t row = 0; row < H; ++row)
+      {
+        for (uint32_t col = 0; col < W; ++col)
+        {
+          mutable_ptr->at(n, ch, row, col) = coco::ElemID{16};
+        }
+      }
+    }
+  }
+
+  for (uint32_t n = 0; n < N; ++n)
+  {
+    for (uint32_t ch = 0; ch < C; ++ch)
+    {
+      for (uint32_t row = 0; row < H; ++row)
+      {
+        for (uint32_t col = 0; col < W; ++col)
+        {
+          ASSERT_EQ(immutable_ptr->at(n, ch, row, col).value(), 16);
+        }
+      }
+    }
+  }
+}