diff options
Diffstat (limited to 'runtimes/neurun/src/linear')
-rw-r--r-- | runtimes/neurun/src/linear/Linear.cc | 152 | ||||
-rw-r--r-- | runtimes/neurun/src/linear/Linear.h | 25 |
2 files changed, 158 insertions, 19 deletions
diff --git a/runtimes/neurun/src/linear/Linear.cc b/runtimes/neurun/src/linear/Linear.cc index 2ffcbdb93..6452bbd49 100644 --- a/runtimes/neurun/src/linear/Linear.cc +++ b/runtimes/neurun/src/linear/Linear.cc @@ -14,19 +14,26 @@ * limitations under the License. */ +#include <algorithm> + #include "Linear.h" #include "graph/Graph.h" #include "graph/operation/LowerInfo.h" -#include "backend/IStageGenerator.h" +#include "backend/interface/IStageGenerator.h" +#include "backend/interface/IConfig.h" +#include "compiler/SubTensorInfo.h" +#include "compiler/TensorInfo.h" + +#include "util/logging.h" namespace neurun { namespace linear { -Linear::Linear(const graph::Graph &graph) +Linear::Linear(const graph::Graph &graph) : _graph(graph) { // Linearize with topological sort // @@ -36,38 +43,157 @@ Linear::Linear(const graph::Graph &graph) // 3. Reverse the order of nodes graph::Graph::PostDfsConstIterator().iterate( - graph, [&](const neurun::graph::operation::Node &node) { _operations.emplace_back(&node); }); + graph, [&](const model::operation::Index &index, const model::operation::Node &node) { + const auto lower_info = graph.getLowerInfo(index); + _operations.emplace_back(&node, lower_info); + }); std::reverse(std::begin(_operations), std::end(_operations)); } -void Linear::accept(graph::operation::NodeVisitor &&visitor) const +void Linear::accept(model::operation::NodeVisitor &&visitor) const { for (const auto op : _operations) { - op->accept(std::move(visitor)); + op.node->accept(std::move(visitor)); } } -backend::TensorBuilderSet Linear::markTensors() const +backend::TensorBuilderSet Linear::planTensors() { + using ITensorBuilderPtr = std::shared_ptr<backend::ITensorBuilder>; + using FnOnTensorBuilder = + std::function<void(const model::operand::Index &ind, ITensorBuilderPtr)>; + + const auto &operands = _graph.operands(); + auto iterTensorBuilders = [&operands](const model::operand::Index &ind, FnOnTensorBuilder fn) { + const auto &obj = operands.at(ind); + for (auto backend : obj.lower_info()->def_backends()) + { + auto tensor_builder = backend->tensor_builder(); + fn(ind, tensor_builder); + } + }; + backend::TensorBuilderSet tensor_builders; + + std::unordered_map<model::operand::Index, uint32_t> uses_map; + std::vector<model::operand::Index> constants; + + _graph.operands().iterate( + [&](const model::operand::Index &ind, const model::operand::Object &obj) { + uses_map[ind] = obj.getUses().size(); + + // If a tensor is a constant, increase the use of the tensor. + // It makes the tensor not be dealloced. + if (obj.getUsage() == model::operand::OperandUsage::CONSTANT) + { + constants.push_back(ind); + uses_map[ind]++; + } + + for (auto backend : obj.lower_info()->def_backends()) + { + bool isSubTensor = false; + auto tensor_builder = backend->tensor_builder(); + + if (backend->config()->SupportSubTensorAlloc()) + { + const auto parentInfo = obj.parent_info(); + if (parentInfo != nullptr) + { + isSubTensor = true; + } + } + + if (isSubTensor) + { + const compiler::SubTensorInfo info(obj); + tensor_builder->registerSubTensorInfo(ind, info); + } + else + { + const auto info = compiler::TensorInfo(obj.shape(), obj.typeInfo()); + tensor_builder->registerTensorInfo(ind, info); + } + + // Prepare tensor builders to be returned + tensor_builders.insert(tensor_builder); + } + }); + + // If a tensor is model output, increase the use of the tensor. + // This aim is same to above one. + for (const auto &ind : _graph.getOutputs()) + { + uses_map[ind]++; + } + + // Allocate constant operands first + VERBOSE(LINEAR) << "TENSORS as CONSTANT" << std::endl; + for (const auto &ind : constants) + { + iterTensorBuilders(ind, [](const model::operand::Index &ind, ITensorBuilderPtr tensor_builder) { + tensor_builder->notifyFirstUse(ind); + }); + } + + // Allocate Model's inputs + VERBOSE(LINEAR) << "TENSORS as MODEL INPUT" << std::endl; + for (const auto &ind : _graph.getInputs()) + { + iterTensorBuilders(ind, [](const model::operand::Index &ind, ITensorBuilderPtr tensor_builder) { + tensor_builder->notifyFirstUse(ind); + }); + } + + // At each operation, + // 1. Scan USE of inputs. Decrease the USE and deallocate if the USE is 0 + // 2. Scan DEF of outputs. If the DEF, allocate it + VERBOSE(LINEAR) << "TENSORS" << std::endl; for (const auto op : _operations) { - const auto tensor_builder = op->lower_info()->backend().stage_gen()->tensor_builder(); - for (const auto &ind : op->getInputs()) + for (const auto &ind : op.node->getOutputs()) { - tensor_builder->mark(ind); - tensor_builders.insert(tensor_builder); + const auto &obj = operands.at(ind); + if (obj.getDef().size()) + { + iterTensorBuilders(ind, + [](const model::operand::Index &ind, ITensorBuilderPtr tensor_builder) { + tensor_builder->notifyFirstUse(ind); + }); + } } - for (const auto &ind : op->getOutputs()) + + for (const auto &ind : op.node->getInputs()) { - tensor_builder->mark(ind); - tensor_builders.insert(tensor_builder); + uses_map[ind]--; + if (uses_map[ind] == 0) + { + iterTensorBuilders(ind, + [](const model::operand::Index &ind, ITensorBuilderPtr tensor_builder) { + tensor_builder->notifyLastUse(ind); + }); + } } } + + // Now, model outputs should be not deallocated + assert(std::all_of(_graph.getOutputs().begin(), _graph.getOutputs().end(), + [&uses_map](const model::operand::Index &ind) { return uses_map[ind] > 0; })); + + // Set subtensor information + // Todo: move this phase outside as optimization phase return tensor_builders; } +void Linear::iterate(const std::function<void(const Element &element)> &fn) const +{ + for (const auto op : _operations) + { + fn(op); + } +} + } // namespace linear } // namespace neurun diff --git a/runtimes/neurun/src/linear/Linear.h b/runtimes/neurun/src/linear/Linear.h index ffbc68ecb..fb3f539d4 100644 --- a/runtimes/neurun/src/linear/Linear.h +++ b/runtimes/neurun/src/linear/Linear.h @@ -19,8 +19,8 @@ #include <vector> -#include "graph/operation/Node.h" -#include "backend/ITensorBuilder.h" +#include "model/operation/Node.h" +#include "backend/interface/ITensorBuilder.h" namespace neurun { @@ -46,6 +46,17 @@ namespace neurun namespace linear { +struct Element +{ + const model::operation::Node *node; + const graph::operation::LowerInfo *lower_info; + + Element(const model::operation::Node *node, const graph::operation::LowerInfo *lower_info) + : node{node}, lower_info{lower_info} + { + } +}; + class Linear { public: @@ -55,14 +66,16 @@ public: Linear(const Linear &linear) = delete; public: - void accept(graph::operation::NodeVisitor &&visitor) const; + void accept(model::operation::NodeVisitor &&visitor) const; // TODO Should not return TensorBuilderSet - virtual backend::TensorBuilderSet markTensors() const; + backend::TensorBuilderSet planTensors(); + + void iterate(const std::function<void(const Element &element)> &fn) const; -public: private: - std::vector<const graph::operation::Node *> _operations; + const graph::Graph &_graph; + std::vector<Element> _operations; }; } // namespace linear |