[enco] Decouple ANN IR from enco Code (#1780)

With this commit, enco::Code no longer holds ANNContext. Each pass can
access ANNContext through SubnetManager.

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

diff --git a/contrib/enco/core/src/Code.h b/contrib/enco/core/src/Code.h
index 336dc39..91756d5 100644 (file)
--- a/contrib/enco/core/src/Code.h
+++ b/contrib/enco/core/src/Code.h
@@ -37,16 +37,9 @@ public:
   coco::Module *module(void) const { return _module; }
   coco::Data *data(void) const { return _data; }
 
-public:
-  ANNContext *ann(void) { return &_ann; }
-  const ANNContext *ann(void) const { return &_ann; }
-
 private:
   coco::Module *const _module;
   coco::Data *const _data;
-
-private:
-  ANNContext _ann;
 };
 
 } // namespace enco

diff --git a/contrib/enco/core/src/Code.test.cpp b/contrib/enco/core/src/Code.test.cpp
index ba7cf01..8e96e47 100644 (file)
--- a/contrib/enco/core/src/Code.test.cpp
+++ b/contrib/enco/core/src/Code.test.cpp
@@ -27,5 +27,4 @@ TEST(CODE, constructor)
 
   ASSERT_EQ(code.module(), m.get());
   ASSERT_EQ(code.data(), d.get());
-  ASSERT_NE(code.ann(), nullptr);
 }

diff --git a/contrib/enco/core/src/CppCode.cpp b/contrib/enco/core/src/CppCode.cpp
index db984b4..39ac7a0 100644 (file)
--- a/contrib/enco/core/src/CppCode.cpp
+++ b/contrib/enco/core/src/CppCode.cpp
@@ -17,6 +17,7 @@
 #include "CppCode.h"
 
 #include "Transforms/GlobalDataGeneration.h"
+#include "Transforms/Split.h"
 
 #include "CppGen/MemoryContext.h"
 
@@ -86,7 +87,9 @@ std::set<coco::Bag *> intermediates(const enco::Code *code)
     };
 
     auto is_host_block = [code](const coco::Block *blk) {
-      return (blk) ? (code->ann()->find(blk) == nullptr) : false;
+      auto m = blk->module();
+      auto ann_ctx = enco::SubnetManager::context(m);
+      return (blk) ? (ann_ctx->find(blk) == nullptr) : false;
     };
 
     for (const auto &updater : coco::updaters(b))
@@ -131,6 +134,7 @@ namespace enco
 void CppCode::dump(std::ostream &os) const
 {
   auto m = _code->module();
+  auto ann_ctx = enco::SubnetManager::context(m);
 
   NetworkStruct network;
   InvokeFunction invoke;
@@ -160,11 +164,11 @@ void CppCode::dump(std::ostream &os) const
    * };
    *
    */
-  for (uint32_t n = 0; n < _code->ann()->count(); ++n)
+  for (uint32_t n = 0; n < ann_ctx->count(); ++n)
   {
     SubnetStructBuilder builder;
 
-    auto subnet_binder = _code->ann()->nth(n);
+    auto subnet_binder = ann_ctx->nth(n);
     auto subnet_struct_name = pp::fmt("Subnet_", subnet_ctx.size());
     auto subnet_field_name = pp::fmt("_subnet_", subnet_ctx.size());
 
@@ -267,7 +271,7 @@ void CppCode::dump(std::ostream &os) const
     invoke.body.append("{");
     invoke.body.indent();
 
-    if (auto binder = _code->ann()->find(blk))
+    if (auto binder = ann_ctx->find(blk))
     {
       // Generate code that invokes Android NN sub-network
       auto lines = subnet_compiler.compile(binder);

diff --git a/contrib/enco/core/src/Transforms/GlobalDataGeneration.cpp b/contrib/enco/core/src/Transforms/GlobalDataGeneration.cpp
index 123f466..f0c4dc8 100644 (file)
--- a/contrib/enco/core/src/Transforms/GlobalDataGeneration.cpp
+++ b/contrib/enco/core/src/Transforms/GlobalDataGeneration.cpp
@@ -15,6 +15,7 @@
  */
 
 #include "GlobalDataGeneration.h"
+#include "Split.h"
 #include "Dims.h"
 
 #include <nncc/foundation/Memory.h>
@@ -101,11 +102,13 @@ GlobalOffset GlobalData::dims_offset(const coco::Output *out) { return dims_offs
 void generate_global_data(std::ostream &os, enco::Code *code)
 {
   auto m = code->module();
+  auto ann_ctx = enco::SubnetManager::context(m);
+
   auto global = make_unique<Global>(os);
 
-  for (uint32_t n = 0; n < code->ann()->count(); ++n)
+  for (uint32_t n = 0; n < ann_ctx->count(); ++n)
   {
-    auto binder = code->ann()->nth(n);
+    auto binder = ann_ctx->nth(n);
 
     auto emit = [&](const ann::OperandID &id, const ann::Operand *info) {
       if (info->weight())

diff --git a/contrib/enco/core/src/Transforms/Split.cpp b/contrib/enco/core/src/Transforms/Split.cpp
index de2a234..ffc69bd 100644 (file)
--- a/contrib/enco/core/src/Transforms/Split.cpp
+++ b/contrib/enco/core/src/Transforms/Split.cpp
@@ -21,10 +21,31 @@
 #include <coco/IR.h>
 
 #include <nncc/core/ADT/kernel/NHWCLayout.h>
+#include <nncc/foundation/Memory.h>
 
+#include <map>
 #include <stdexcept>
 #include <functional>
 
+using nncc::foundation::make_unique;
+
+namespace
+{
+
+std::map<const coco::Module *, std::unique_ptr<ANNContext>> _subnet_contexts;
+
+} // namespace
+
+namespace enco
+{
+
+const ANNContext *SubnetManager::context(const coco::Module *m)
+{
+  return _subnet_contexts.at(m).get();
+}
+
+} // namespace enco
+
 namespace
 {
 
@@ -584,11 +605,15 @@ namespace enco
 
 void SplitPass::runOnCode(enco::Code *code) const
 {
-  ANNGroupBuilder group_builder{code->ann()};
+  auto ann_ctx = make_unique<ANNContext>();
+
+  ANNGroupBuilder group_builder{ann_ctx.get()};
   group_builder.build(code);
 
   ANNModuleBuilder module_builder;
-  module_builder.build(code->ann());
+  module_builder.build(ann_ctx.get());
+
+  _subnet_contexts[code->module()] = std::move(ann_ctx);
 }
 
 } // namespace enco

diff --git a/contrib/enco/core/src/Transforms/Split.h b/contrib/enco/core/src/Transforms/Split.h
index 0c6a095..3b93d61 100644 (file)
--- a/contrib/enco/core/src/Transforms/Split.h
+++ b/contrib/enco/core/src/Transforms/Split.h
@@ -22,6 +22,11 @@
 namespace enco
 {
 
+struct SubnetManager
+{
+  static const ANNContext *context(const coco::Module *m);
+};
+
 /**
  * @brief Split instructions into a set of phases
  */