[neurun] Introduce set operators for LayoutSet (#2658)

Introduce Union, Intersect and Minus operators to operand::LayoutSet.

Signed-off-by: Hanjoung Lee <hanjoung.lee@samsung.com>

diff --git a/runtimes/neurun/src/graph/operand/LayoutSet.cc b/runtimes/neurun/src/graph/operand/LayoutSet.cc
index 0407cee..47bb590 100644 (file)
--- a/runtimes/neurun/src/graph/operand/LayoutSet.cc
+++ b/runtimes/neurun/src/graph/operand/LayoutSet.cc
@@ -31,6 +31,39 @@ LayoutSet::LayoutSet(std::initializer_list<Layout> layouts)
   }
 }
 
+LayoutSet LayoutSet::operator|(const LayoutSet &other) const
+{
+  auto ret = *this;
+  for (auto layout : other)
+  {
+    ret.add(layout);
+  }
+  return ret;
+}
+
+LayoutSet LayoutSet::operator&(const LayoutSet &other) const
+{
+  LayoutSet ret;
+  for (auto layout : other)
+  {
+    if (contains(layout))
+    {
+      ret.add(layout);
+    }
+  }
+  return ret;
+}
+
+LayoutSet LayoutSet::operator-(const LayoutSet &other) const
+{
+  auto ret = *this;
+  for (auto layout : other)
+  {
+    ret.remove(layout);
+  }
+  return ret;
+}
+
 } // namespace operand
 } // namespace graph
 } // namespace neurun

diff --git a/runtimes/neurun/src/graph/operand/LayoutSet.h b/runtimes/neurun/src/graph/operand/LayoutSet.h
index 517753e..928259c 100644 (file)
--- a/runtimes/neurun/src/graph/operand/LayoutSet.h
+++ b/runtimes/neurun/src/graph/operand/LayoutSet.h
@@ -42,6 +42,11 @@ public:
   bool contains(const Layout &layout) const { return _set.find(layout) != _set.end(); }
 
 public:
+  LayoutSet operator|(const LayoutSet &other) const; // Union
+  LayoutSet operator&(const LayoutSet &other) const; // Intersect
+  LayoutSet operator-(const LayoutSet &other) const; // Minus
+
+public:
   std::unordered_set<Layout>::const_iterator begin() const { return _set.begin(); }
   std::unordered_set<Layout>::const_iterator end() const { return _set.end(); }
 

diff --git a/runtimes/neurun/test/graph/operand/LayoutSet.cc b/runtimes/neurun/test/graph/operand/LayoutSet.cc
new file mode 100644 (file)
index 0000000..6ff21cb
--- /dev/null
+++ b/runtimes/neurun/test/graph/operand/LayoutSet.cc
@@ -0,0 +1,27 @@
+#include <gtest/gtest.h>
+
+#include "graph/operand/LayoutSet.h"
+
+using neurun::graph::operand::Layout;
+using neurun::graph::operand::LayoutSet;
+
+TEST(graph_operand_LayoutSet, layout_set_operators)
+{
+  LayoutSet set1{Layout::NCHW};
+  LayoutSet set2{Layout::NHWC};
+  LayoutSet set3 = set1 | set2;
+
+  ASSERT_EQ(set3.size(), 2);
+
+  ASSERT_EQ((set3 - set1).size(), 1);
+  ASSERT_EQ((set3 - set1).contains(Layout::NHWC), true);
+  ASSERT_EQ((set3 - set2).size(), 1);
+  ASSERT_EQ((set3 - set2).contains(Layout::NCHW), true);
+  ASSERT_EQ((set3 - set3).size(), 0);
+
+  ASSERT_EQ((set3 & set1).size(), 1);
+  ASSERT_EQ((set3 & set1).contains(Layout::NCHW), true);
+  ASSERT_EQ((set3 & set2).size(), 1);
+  ASSERT_EQ((set3 & set2).contains(Layout::NHWC), true);
+  ASSERT_EQ((set1 & set2).size(), 0);
+}