diff options
Diffstat (limited to 'runtimes/neurun/src/linear/Linear.cc')
| -rw-r--r-- | runtimes/neurun/src/linear/Linear.cc | 152 |
1 files changed, 139 insertions, 13 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 |