Introduce 'IndexEnumerator' (#1502)

This commit introduces 'tensor::IndexEnumerator' and uses it to simplify
the implementation of 'tensor::IndexIterator'.

Unlike 'tensor::IndexIterator', 'tensor::IndexEnumerator' can be used in for loop
similarly as C++ standard iterator.

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

diff --git a/include/util/tensor/IndexEnumerator.h b/include/util/tensor/IndexEnumerator.h
new file mode 100644 (file)
index 0000000..30325cb
--- /dev/null
+++ b/include/util/tensor/IndexEnumerator.h
@@ -0,0 +1,101 @@
+/*
+ * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __NNFW_UTIL_TENSOR_INDEX_ENUMERATOR_H__
+#define __NNFW_UTIL_TENSOR_INDEX_ENUMERATOR_H__
+
+#include "util/tensor/Shape.h"
+#include "util/tensor/Index.h"
+
+namespace nnfw
+{
+namespace util
+{
+namespace tensor
+{
+
+class IndexEnumerator
+{
+public:
+  explicit IndexEnumerator(const Shape &shape) : _shape(shape), _index(shape.rank()), _cursor(0)
+  {
+    const size_t rank = _shape.rank();
+
+    for (size_t axis = 0; axis < rank; ++axis)
+    {
+      _index.at(axis) = 0;
+    }
+
+    for (_cursor = 0; _cursor < rank; ++_cursor)
+    {
+      if (_index.at(_cursor) < _shape.dim(_cursor))
+      {
+        break;
+      }
+    }
+  }
+
+public:
+  IndexEnumerator(IndexEnumerator &&) = delete;
+  IndexEnumerator(const IndexEnumerator &) = delete;
+
+public:
+  bool valid(void) const { return _cursor < _shape.rank(); }
+
+public:
+  const Index &curr(void) const { return _index; }
+
+public:
+  void advance(void)
+  {
+    const size_t rank = _shape.rank();
+
+    // Find axis to be updated
+    while((_cursor < rank) && !(_index.at(_cursor) + 1 < _shape.dim(_cursor)))
+    {
+      ++_cursor;
+    }
+
+    if(_cursor == rank)
+    {
+      return;
+    }
+
+    // Update index
+    _index.at(_cursor) += 1;
+
+    for (size_t axis = 0; axis < _cursor; ++axis)
+    {
+      _index.at(axis) = 0;
+    }
+
+    // Update cursor
+    _cursor = 0;
+  }
+
+public:
+  const Shape _shape;
+
+private:
+  size_t _cursor;
+  Index _index;
+};
+
+} // namespace tensor
+} // namespace util
+} // namespace nnfw
+
+#endif // __NNFW_UTIL_TENSOR_INDEX_ENUMERATOR_H__

diff --git a/include/util/tensor/IndexIterator.h b/include/util/tensor/IndexIterator.h
index 56a8c7d..cbe8951 100644 (file)
--- a/include/util/tensor/IndexIterator.h
+++ b/include/util/tensor/IndexIterator.h
@@ -19,6 +19,7 @@
 
 #include "util/tensor/Shape.h"
 #include "util/tensor/Index.h"
+#include "util/tensor/IndexEnumerator.h"
 
 namespace nnfw
 {
@@ -43,44 +44,9 @@ public:
 public:
   template <typename Callable> IndexIterator &iter(Callable fn)
   {
-    Index index(_shape.rank());
-
-    for (size_t d = 0; d < _shape.rank(); ++d)
-    {
-      index.at(d) = 0;
-    }
-
-    size_t cursor = 0;
-
-    while (cursor < _shape.rank())
+    for (IndexEnumerator e{_shape}; e.valid(); e.advance())
     {
-      fn(index);
-
-      if (index.at(cursor) + 1 < _shape.dim(cursor))
-      {
-        index.at(cursor) += 1;
-      }
-      else
-      {
-        while ((cursor < _shape.rank()) && (index.at(cursor) + 1 == _shape.dim(cursor)))
-        {
-          ++cursor;
-        }
-
-        if (cursor == _shape.rank())
-        {
-          break;
-        }
-
-        index.at(cursor) += 1;
-
-        for (size_t d = 0; d < cursor; ++d)
-        {
-          index.at(d) = 0;
-        }
-
-        cursor = 0;
-      }
+      fn(e.curr());
     }
 
     return (*this);

diff --git a/libs/util/examples/tensor_index_iterator.cpp b/libs/util/examples/tensor_index_iterator.cpp
index 01346ee..284e04a 100644 (file)
--- a/libs/util/examples/tensor_index_iterator.cpp
+++ b/libs/util/examples/tensor_index_iterator.cpp
@@ -16,10 +16,47 @@
 
 #include "util/tensor/IndexIterator.h"
 
+#include <array>
+
 #include <iostream>
+#include <algorithm>
+
+#include <cassert>
+
+void test_iterate(void)
+{
+  const nnfw::util::tensor::Shape shape{3, 4, 7};
+
+  std::array<int, 3 * 4 * 7> array;
+
+  array.fill(0);
+
+  using nnfw::util::tensor::iterate;
+  using nnfw::util::tensor::Index;
+
+  iterate(shape) << [&](const Index &index) {
+    assert(index.rank() == shape.rank());
+
+    const size_t rank = index.rank();
+
+    uint32_t offset = index.at(0);
+
+    for (size_t axis = 1; axis < rank; ++axis)
+    {
+      offset *= shape.dim(axis);
+      offset += index.at(axis);
+    }
+
+    array[offset] += 1;
+  };
+
+  assert(std::all_of(array.begin(), array.end(), [](int num) { return num == 1; }));
+}
 
 int main(int argc, char **argv)
 {
+  test_iterate();
+
   nnfw::util::tensor::Shape shape{3, 4, 3, 4};
 
   std::cout << "Iterate over tensor{3, 4, 3, 4}" << std::endl;