onert::compiler::train::TrainingCompiler Class Reference

Class to compile NN package. More...

#include <TrainingCompiler.h>

Collaboration diagram for onert::compiler::train::TrainingCompiler:

Public Member Functions
	TrainingCompiler (std::unique_ptr< ir::NNPkg > nnpkg, CompilerOptions *copts, const ir::train::TrainingInfo &training_info)
	Construct a new TrainingCompiler object for an nnpkg.
	TrainingCompiler (void)=delete
	Construct a TrainingCompiler object.
	~TrainingCompiler ()=default
	Destroy the TrainingCompiler object.
std::unique_ptr< CompilerArtifact >	compile (void)
	Do compilation with the options.
Public Member Functions inherited from onert::compiler::ICompiler
virtual	~ICompiler ()=default
	Virtual ICompiler destructor.

Detailed Description

Class to compile NN package.

Definition at line 36 of file TrainingCompiler.h.

Constructor & Destructor Documentation

TrainingCompiler() [1/2]

onert::compiler::train::TrainingCompiler::TrainingCompiler ( std::unique_ptr< ir::NNPkg >  nnpkg, CompilerOptions *  copts, const ir::train::TrainingInfo &  training_info  )

explicit

Construct a new TrainingCompiler object for an nnpkg.

Parameters

[in]	nnpkg	nnpkg to compile
[in]	copts	compiler options
[in]	training_info	training information

Definition at line 44 of file TrainingCompiler.cc.

  : _model{nnpkg->primary_model()}, _options{copts}, _training_info{training_info}
{
  if (nnpkg->model_count() > 1)
    throw std::runtime_error("TrainingCompiler does not support multiple models yet");
 
  if (nnpkg->primary_model()->subgraphs_count() > 1)
    throw std::runtime_error("TrainingCompiler does not support multiple subgraphs yet");
}

TrainingCompiler() [2/2]

onert::compiler::train::TrainingCompiler::TrainingCompiler ( void  )

delete

Construct a TrainingCompiler object.

~TrainingCompiler()

onert::compiler::train::TrainingCompiler::~TrainingCompiler ( )

default

Destroy the TrainingCompiler object.

Member Function Documentation

compile()

std::unique_ptr< CompilerArtifact > onert::compiler::train::TrainingCompiler::compile ( void  )

virtual

Do compilation with the options.

Returns

std::unique_ptr<CompilerArtifact> Executors as a result of compilation

Implements onert::compiler::ICompiler.

Definition at line 55 of file TrainingCompiler.cc.

{
  /***************************************************
   * Prepare compilation phase
   ***************************************************/
  if (!_options)
    throw std::runtime_error{"Empty compile option"};
 
  // Mode check
  // TODO handle option for each model
  if (_options->he_profiling_mode)
  {
    if (!_options->he_scheduler)
      throw std::runtime_error("Heterogeneous scheduler must be enabled during profiling.");
 
    if (_options->executor != "Dataflow")
      throw std::runtime_error("Profiling mode works only with 'Dataflow' executor");
  }
 
  _options->forceInternalOptions();
  _options->verboseOptions();
 
  auto custom_kernel_builder = _model->getKernelBuilder();
 
  _model->iterate([&](const ir::SubgraphIndex &, ir::IGraph &graph) {
    auto &subg = nnfw::misc::polymorphic_downcast<ir::Graph &>(graph);
    // Mandatory passes
    compiler::pass::PassRunner{}
      .append(std::make_unique<compiler::pass::ConstantOutputPass>(subg))
      .append(std::make_unique<compiler::pass::OddOutputPass>(subg))
      .run();
 
    // Optimizations
    compiler::pass::PassRunner{}
      .append(std::make_unique<compiler::pass::UnusedOperandEliminationPass>(subg))
      .run();
  });
 
  std::unordered_map<ir::SubgraphIndex, std::shared_ptr<ir::train::TrainableGraph>>
    trainable_subgraphs;
 
  if (_model->hasOnly<ir::Graph>())
  {
    // Create trainable subgraphs by copy and converting inference model
    _model->iterate([&](const ir::SubgraphIndex &subg_index, const ir::IGraph &graph) {
      const auto &subg = nnfw::misc::polymorphic_downcast<const ir::Graph &>(graph);
      // Create TrainableGraph by copying Graph
      auto trainable_subg = std::make_shared<ir::train::TrainableGraph>(subg);
 
      // Convert operations to trainable operations
      auto converter = TrainableOperationConverter{*trainable_subg, &_training_info};
      ir::OperationIndex min_trainable_op_idx;
      subg.operations().iterate(
        [&](const onert::ir::OperationIndex &op_index, const onert::ir::IOperation &op) {
          auto trainable_op = converter(op);
          if (_training_info.getTrainableOps().find(op_index) !=
              std::end(_training_info.getTrainableOps()))
          {
            trainable_op->enableWeightsUpdate();
            if (op_index.value() < min_trainable_op_idx.value())
            {
              min_trainable_op_idx = op_index;
            }
          }
          [[maybe_unused]] auto gen_index =
            trainable_subg->replaceOperation(op_index, std::move(trainable_op));
          assert(gen_index == op_index);
        });
 
      for (ir::OperationIndex idx{min_trainable_op_idx};
           idx.value() < trainable_subg->operations().size(); idx++)
      {
        trainable_subg->enableBackward(idx);
      }
 
      trainable_subgraphs[subg_index] = std::move(trainable_subg);
    });
  }
  else
  {
    // TODO Support models that have TrainableGraphs
    throw std::runtime_error("TrainingCompiler: Invalid model");
  }
 
  // operation
  _model.reset();
 
  // TODO Handle dump level for each model
  auto dump_level = static_cast<dumper::dot::DotDumper::Level>(_options->graph_dump_level);
  onert::dumper::dot::DotDumper dot_dumper(dump_level);
 
  for (const auto &[subg_index, subg] : trainable_subgraphs)
  {
    dot_dumper.dump(*subg, nnfw::misc::str("before_loss_insertion-", subg_index.value()));
  }
 
  // Apply pass for trainable subgraphs
  for (auto &&[subg_index, trainable_subg] : trainable_subgraphs)
  {
    compiler::pass::PassRunner{}
      .append(std::make_unique<train::pass::LossInsertionPass>(*trainable_subg, &_training_info,
                                                               subg_index))
      .run();
  }
 
  for (const auto &[subg_index, subg] : trainable_subgraphs)
  {
    dot_dumper.dump(*subg, nnfw::misc::str("after_loss_insertion-", subg_index.value()));
  }
 
  for (auto &&[subg_index, subg] : trainable_subgraphs)
  {
    subg->updateGraphDependency();
    subg->verify();
 
    dot_dumper.dump(*subg,
                    nnfw::misc::str("after_initializing_training_usedefs-", subg_index.value()));
  }
 
  // Change input shape according to batch_size
  for (auto &&pair : trainable_subgraphs)
  {
    auto trainable_subg = pair.second;
 
    for (const auto &ind : trainable_subg->getInputs())
    {
      auto &input = trainable_subg->operands().at(ind);
      auto new_shape = input.info().shape();
      // TODO Consider batch size index
      if (new_shape.dim(0) != 1 && new_shape.dim(0) != ir::Shape::kUnspecifiedDim)
        throw std::runtime_error("the first dim is not 1. It is not supported yet.");
      new_shape.dim(0) = _training_info.batchSize();
      input.info().shape(new_shape);
    }
  }
 
  /***************************************************
   * Backend independent analysis & optimization phase
   ***************************************************/
  // Tracing context
  auto tracing_ctx = std::make_unique<util::TracingCtx>();
 
  // Lower: Assign backend
  std::unordered_map<ir::SubgraphIndex, std::unique_ptr<compiler::train::LoweredTrainableGraph>>
    lowered_subgs;
  {
    for (auto &&[subg_index, trainable_subg] : trainable_subgraphs)
    {
      // Lower: Assign backend
      lowered_subgs[subg_index] =
        std::make_unique<compiler::train::LoweredTrainableGraph>(*trainable_subg, *_options);
      // Set tracing_ctx for copied graph
      tracing_ctx->setSubgraphIndex(&(lowered_subgs[subg_index]->graph()),
                                    {ir::ModelIndex{0}, subg_index});
    }
  }
 
  for (const auto &[subg_index, lowered_subg] : lowered_subgs)
  {
    dot_dumper.dump(*lowered_subg, nnfw::misc::str("after_lower_subg-", subg_index.value()));
  }
 
  // Set operands' info for back propagation as default tensor info
  for (const auto &pair : lowered_subgs)
  {
    auto lowered_subg = pair.second.get();
    auto &tgraph = lowered_subg->trainable_graph();
    tgraph.operands().iterate([&](const ir::OperandIndex &index, const ir::Operand &obj) {
      if (!obj.isConstant())
      {
        auto bwd_operand = std::make_unique<ir::Operand>(obj);
        [[maybe_unused]] const auto gen_index =
          tgraph.addBackwardOperand(index, std::move(bwd_operand));
        assert(gen_index == index);
      }
    });
  }
 
  // Shape inference.
  {
    // Run the StaticShapeInfer of primary subg. All child StaticShapeInferers are called
    // recursively
    std::unordered_map<ir::SubgraphIndex, std::unique_ptr<StaticShapeInferer>> inferers =
      createStaticShapeInferers(lowered_subgs);
 
    const auto primary_subg_idx = ir::SubgraphIndex{0};
    inferers.at(primary_subg_idx)->infer();
 
    for (const auto &pair_inferer : inferers)
    {
      const auto inferer = pair_inferer.second.get();
      inferer->dump();
    }
 
    // NOTE StaticBackwardShapeInferer is allocated for each subgraph,
    //      so it does not support models that have controlflow operations yet.
    for (auto &&pair : lowered_subgs)
    {
      auto &lowered_subg = pair.second;
      auto inferer = std::make_unique<StaticBackwardShapeInferer>(lowered_subg.get());
      inferer->infer();
      inferer->dump();
    }
  }
 
  // Shape validation
  for (const auto &pair : lowered_subgs)
  {
    auto &lowered_subg = pair.second;
    compiler::ShapeValidator{lowered_subg->graph()}();
  }
 
  // TODO Validate shapes of the tensors for back propagation
 
  /*************************************************************
   *  Backend independent analysis & optimization phase finished
   *************************************************************/
  auto executors = std::make_shared<exec::train::TrainableExecutors>();
  for (auto &&[subg_index, lowered_subg] : lowered_subgs)
  {
    auto const model_index = ir::ModelIndex{0};
    auto const indexed_ranks = lowered_subg->indexed_ranks();
 
    ir::OperationDumper dumper("Executor generation of Subgraph " +
                               std::to_string(subg_index.value()));
    lowered_subg->graph().operations().iterate(
      [&](const ir::OperationIndex &, const ir::IOperation &op) { op.accept(dumper); });
 
    ExecutorFactoryArgs args;
    args.tracing_ctx = tracing_ctx.get();
    args.options = _options;
    args.model_index = model_index;
    args.custom_kernel_builder = custom_kernel_builder;
    // TODO: Support internal io if necessary
    auto executor = std::unique_ptr<exec::IExecutor>{
      ExecutorFactory::get().create(std::move(lowered_subg), executors, args, _training_info)};
    executor->setIndexedRanks(indexed_ranks);
    executors->emplace(model_index, subg_index, std::move(executor));
  }
 
  /********************************
   * Code generation phase finished
   ********************************/
  return std::make_unique<CompilerArtifact>(executors, std::move(tracing_ctx));
}

References onert::ir::IOperation::accept(), onert::compiler::pass::PassRunner::append(), onert::ir::train::TrainingInfo::batchSize(), onert::compiler::ExecutorFactory::create(), onert::dumper::dot::DotDumper::dump(), onert::compiler::CompilerOptions::executor, onert::compiler::CompilerOptions::forceInternalOptions(), onert::compiler::ExecutorFactory::get(), onert::ir::train::TrainingInfo::getTrainableOps(), onert::compiler::CompilerOptions::graph_dump_level, onert::compiler::CompilerOptions::he_profiling_mode, onert::compiler::CompilerOptions::he_scheduler, onert::compiler::pass::PassRunner::run(), run(), onert::exec::IExecutor::setIndexedRanks(), nnfw::misc::str(), onert::util::Index< T, DummyTag >::value(), and onert::compiler::CompilerOptions::verboseOptions().

---

The documentation for this class was generated from the following files:

• runtime/onert/core/src/compiler/train/TrainingCompiler.h
• runtime/onert/core/src/compiler/train/TrainingCompiler.cc