[core] Add 'num_elements' on kernel shape (#253)

This commit adds 'num_elements' for kernel shape similarly as
'num_elements' for tensor shape.

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

diff --git a/libs/core/include/nncc/core/ADT/kernel/Shape.h b/libs/core/include/nncc/core/ADT/kernel/Shape.h
index 6376aae..18d3dcb 100644 (file)
--- a/libs/core/include/nncc/core/ADT/kernel/Shape.h
+++ b/libs/core/include/nncc/core/ADT/kernel/Shape.h
@@ -37,6 +37,11 @@ private:
   uint32_t _width;
 };
 
+inline uint64_t num_elements(const Shape &shape)
+{
+  return shape.count() * shape.depth() * shape.height() * shape.width();
+}
+
 } // namespace kernel
 } // namespace ADT
 } // namespace core

diff --git a/libs/core/src/ADT/kernel/Shape.test.cpp b/libs/core/src/ADT/kernel/Shape.test.cpp
index f050964..59cf69f 100644 (file)
--- a/libs/core/src/ADT/kernel/Shape.test.cpp
+++ b/libs/core/src/ADT/kernel/Shape.test.cpp
@@ -16,3 +16,16 @@ TEST(ADT_KERNEL_SHAPE, ctor)
   ASSERT_EQ(shape.height(), H);
   ASSERT_EQ(shape.width(), W);
 }
+
+TEST(ADT_KERNEL_SHAPE, num_elements)
+{
+  const uint32_t N = 1;
+  const uint32_t C = 3;
+  const uint32_t H = 4;
+  const uint32_t W = 5;
+
+  using nncc::core::ADT::kernel::Shape;
+  using nncc::core::ADT::kernel::num_elements;
+
+  ASSERT_EQ(num_elements(Shape{N, C, H, W}), N * C * H * W);
+}