DRAFT CFe to support huge 2G model

on-going draft to support huge size > 2G model.

Signed-off-by: SaeHie Park <[email protected]>

compiler/circledump/driver/Driver.cpp

@@ -50,7 +50,11 @@ int entry(int argc, char **argv)
 
   std::cout << "Dump: " << circle_path << std::endl << std::endl;
 
-  std::cout << circlemodel << std::endl;
+  circledump::ModelEx modelex;
+  modelex.model = circlemodel;
+  modelex.rawdata = &modelData;
+
+  std::cout << modelex << std::endl;
 
   return 0;
 }

compiler/circledump/include/circledump/Dump.h

@@ -24,10 +24,16 @@
 namespace circledump
 {
 
-void dump_model(std::ostream &os, const circle::Model *model);
+struct ModelEx
+{
+  const circle::Model *model;
+  const std::vector<char> *rawdata;
+};
+
+void dump_model(std::ostream &os, const circledump::ModelEx &model);
 
 } // namespace circledump
 
-std::ostream &operator<<(std::ostream &os, const circle::Model *model);
+std::ostream &operator<<(std::ostream &os, const circledump::ModelEx &model);
 
 #endif // __CIRCLEDUMP_DUMP_H__

compiler/circledump/src/Dump.cpp

@@ -341,9 +341,9 @@ void dump_sub_graph(std::ostream &os, mio::circle::Reader &reader)
   os << std::endl;
 }
 
-void dump_model(std::ostream &os, const circle::Model *model)
+void dump_model(std::ostream &os, const circle::Model *model, const std::vector<char> *rawdata)
 {
-  mio::circle::Reader reader(model);
+  mio::circle::Reader reader(model, rawdata);
 
   uint32_t num_subgraph = reader.num_subgraph();
 
@@ -378,13 +378,17 @@ void dump_model(std::ostream &os, const circle::Model *model)
   os << std::endl;
 
   // dump buffer
-  os << "Buffers: B(index) (length) values, if any" << std::endl;
+  os << "Buffers: B(index) (length) values, if any; (length *) for ext_offset" << std::endl;
   for (uint32_t i = 0; i < buffers->size(); ++i)
   {
+    bool ext_offset;
     const uint8_t *buff_data;
-    size_t size = reader.buffer_info(i, &buff_data);
+    size_t size = reader.buffer_info(i, &buff_data, ext_offset);
 
-    os << "B(" << i << ") (" << size << ") ";
+    os << "B(" << i << ") (" << size;
+    if (ext_offset)
+      os << " *";
+    os << ") ";
     if (buff_data != nullptr)
     {
       dump_buffer(os, buff_data, size, 16);
@@ -460,8 +464,8 @@ void dump_model(std::ostream &os, const circle::Model *model)
 
 } // namespace circledump
 
-std::ostream &operator<<(std::ostream &os, const circle::Model *model)
+std::ostream &operator<<(std::ostream &os, const circledump::ModelEx &modelex)
 {
-  circledump::dump_model(os, model);
+  circledump::dump_model(os, modelex.model, modelex.rawdata);
   return os;
 }

compiler/luci/export/include/luci/CircleExporter.h

@@ -36,14 +36,10 @@ class CircleExporter
     virtual ~Contract() = default;
 
   public: // Client -> Exporter
-    // Input Graph (to be exported)
-    // Exporter expects a loco graph that consists of Circle nodes
-    virtual loco::Graph *graph(void) const = 0;
-
     // Input Module (to be exported)
     // Exporter expects a luci module that consists of loco graphs
     // TODO make this pure virtual
-    virtual luci::Module *module(void) const;
+    virtual luci::Module *module(void) const = 0;
 
   public: // Exporter -> Client
     // Exporter calls store for export data

compiler/luci/export/include/luci/CircleFileExpContract.h

@@ -40,7 +40,6 @@ struct CircleFileExpContract : public luci::CircleExporter::Contract
   virtual ~CircleFileExpContract() = default;
 
 public:
-  loco::Graph *graph(void) const final { return nullptr; }
   luci::Module *module(void) const final { return _module; }
 
 public:

compiler/luci/export/src/CircleExporter.cpp

@@ -26,9 +26,6 @@
 namespace luci
 {
 
-// TODO remove this
-Module *CircleExporter::Contract::module(void) const { return nullptr; }
-
 CircleExporter::CircleExporter()
 {
   // NOTHING TO DO
@@ -48,17 +45,7 @@ bool CircleExporter::invoke(Contract *contract) const
     return contract->store(ptr, size);
   }
 
-  auto graph = contract->graph();
-  if (graph == nullptr)
-    return false;
-
-  CircleExporterImpl impl(graph);
-
-  const char *ptr = impl.getBufferPointer();
-  const size_t size = impl.getBufferSize();
-
-  // we just send one time
-  return contract->store(ptr, size);
+  return false;
 }
 
 } // namespace luci

compiler/luci/export/src/CircleExporter.test.cpp

@@ -33,12 +33,12 @@ class SampleGraphContract : public luci::CircleExporter::Contract
   SampleGraphContract() : luci::CircleExporter::Contract(), _buffer(new std::vector<char>)
   {
     // create needed entities
-    _g = loco::make_graph();
-    auto graph_input = _g->inputs()->create();
-    auto graph_output = _g->outputs()->create();
-    input_node = _g->nodes()->create<luci::CircleInput>();
-    output_node = _g->nodes()->create<luci::CircleOutput>();
-    relu_node = _g->nodes()->create<luci::CircleRelu>();
+    auto g = loco::make_graph();
+    auto graph_input = g->inputs()->create();
+    auto graph_output = g->outputs()->create();
+    input_node = g->nodes()->create<luci::CircleInput>();
+    output_node = g->nodes()->create<luci::CircleOutput>();
+    relu_node = g->nodes()->create<luci::CircleRelu>();
 
     // link nodes and link them to graph
     relu_node->features(input_node);
@@ -57,9 +57,12 @@ class SampleGraphContract : public luci::CircleExporter::Contract
 
     graph_output->shape({1, 2, 3, 4});
     graph_output->dtype(loco::DataType::FLOAT32);
+
+    _m = std::unique_ptr<luci::Module>{new luci::Module};
+    _m->add(std::move(g));
   }
 
-  loco::Graph *graph(void) const override { return _g.get(); }
+  luci::Module *module(void) const override { return _m.get(); }
 
 public:
   bool store(const char *ptr, const size_t size) const override
@@ -77,7 +80,7 @@ class SampleGraphContract : public luci::CircleExporter::Contract
   luci::CircleRelu *relu_node;
 
 private:
-  std::unique_ptr<loco::Graph> _g;
+  std::unique_ptr<luci::Module> _m;
   std::unique_ptr<std::vector<char>> _buffer;
 };
 

compiler/luci/export/src/CircleExporterImpl.cpp

@@ -15,14 +15,18 @@
  */
 
 #include "CircleExporterImpl.h"
-#include "Optimize.h"
 #include "CircleExportMetadata.h"
 #include "CircleTensorExporter.h"
 #include "CircleOperationExporter.h"
 #include "CircleExporterUtils.h"
+#include "ProgressReporter.h"
 
 #include <luci/IR/CircleNodes.h>
+#include <luci/Pass/CircleShapeInferencePass.h>
+#include <luci/Pass/CircleTypeInferencePass.h>
 
+#include <loco.h>
+#include <logo/Phase.h>
 #include <oops/InternalExn.h>
 #include <mio/circle/schema_generated.h>
 #include <flatbuffers/flatbuffers.h>
@@ -100,13 +104,35 @@ encodeOperatorCodes(FlatBufferBuilder &builder, std::unordered_map<luci::OpCode,
 
 } // namespace
 
+namespace
+{
+
+void optimize(loco::Graph *g)
+{
+  logo::Phase phase;
+  {
+    // prepare type and shape before optimization
+    phase.emplace_back(std::make_unique<luci::CircleShapeInferencePass>());
+    phase.emplace_back(std::make_unique<luci::CircleTypeInferencePass>());
+
+    // TODO add more optimization passes (with a knob)
+  }
+
+  logo::PhaseRunner<logo::PhaseStrategy::Restart> phase_runner{g};
+
+  luci::ProgressReporter prog(g, logo::PhaseStrategy::Restart);
+  phase_runner.attach(&prog);
+  phase_runner.run(phase);
+}
+
+} // namespace
+
 namespace luci
 {
 
 using namespace circle;
 using namespace flatbuffers;
 
-CircleExporterImpl::CircleExporterImpl(loco::Graph *graph) { exportGraph(graph); }
 CircleExporterImpl::CircleExporterImpl(Module *module) { exportModule(module); }
 
 ::flatbuffers::Offset<::circle::SubGraph>
@@ -121,61 +147,6 @@ CircleExporterImpl::exportSubgraph(SerializedGraphData &gd)
   return subgraph;
 }
 
-void CircleExporterImpl::exportGraph(loco::Graph *graph)
-{
-  // do graph optimization
-  optimize(graph);
-
-  _builder.Clear();
-
-  SerializedModelData md;
-  SerializedGraphData gd;
-
-  // This version is taken from comment in fbs
-  constexpr uint32_t version = 0;
-
-  // set Subgraph name
-  gd._name = graph->name();
-
-  // TODO set this value properly
-  gd._data_format = circle::DataFormat::DataFormat_CHANNELS_LAST;
-
-  // prepare model data
-  prepareModelData(_builder, md);
-
-  // parse graph into SerializedModelData structure
-  exportOpDefinedTensors(graph, _builder, md, gd);
-
-  // NOTE Invoke these register functions only after each node is annotated with its tensor_index
-  registerGraphInputTensors(graph, gd);
-  registerGraphOutputTensors(graph, gd);
-
-  exportNodes(graph, _builder, md, gd);
-
-  // encode operator codes
-  auto operator_codes = encodeOperatorCodes(_builder, md._operator_codes);
-
-  // Subgraphs
-  Offset<SubGraph> subgraph = exportSubgraph(gd);
-  auto subgraphs = _builder.CreateVector(std::vector<Offset<SubGraph>>{subgraph});
-
-  // Description
-  std::string description_str = "nnpackage";
-  auto description = _builder.CreateString(description_str);
-
-  // Metadata
-  auto metadata_vec = createCircleMetadataVector(_builder, md);
-  auto metadata = _builder.CreateVector(std::vector<Offset<Metadata>>(metadata_vec));
-
-  // create array of buffers
-  auto buffers = _builder.CreateVector(md._buffers);
-
-  // Model
-  auto model_offset = CreateModel(_builder, version, operator_codes, subgraphs, description,
-                                  buffers, 0 /* metadata_buffer */, metadata);
-  FinishModelBuffer(_builder, model_offset);
-}
-
 void CircleExporterImpl::exportModule(Module *module)
 {
   assert(module->size() > 0);

compiler/luci/export/src/CircleExporterImpl.h

@@ -38,7 +38,6 @@ class CircleExporterImpl
   CircleExporterImpl() = delete;
   ~CircleExporterImpl() = default;
 
-  explicit CircleExporterImpl(loco::Graph *graph);
   explicit CircleExporterImpl(Module *module);
 
   /**
@@ -59,12 +58,6 @@ class CircleExporterImpl
    */
   flatbuffers::Offset<circle::SubGraph> exportSubgraph(SerializedGraphData &gd);
 
-  /**
-   * @brief root function that writes graph into internal buffer
-   * @param graph
-   */
-  void exportGraph(loco::Graph *graph);
-
   /**
    * @brief root function that writes Module into internal buffer
    * @param module