96 lines
3.2 KiB
C++
96 lines
3.2 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "graph/pipeline.h"
|
|
|
|
#include "archive/value_archive.h"
|
|
#include "core/operator.h"
|
|
#include "graph/common.h"
|
|
|
|
namespace mmdeploy::graph {
|
|
|
|
Pipeline::Pipeline(const Value& cfg) : BaseNode(cfg["pipeline"]) {
|
|
input_idx_ = UpdateBindings(inputs(), kWrite);
|
|
for (auto task_config : cfg["pipeline"]["tasks"]) {
|
|
auto name = task_config.value("name", std::string{});
|
|
auto type = task_config.value("type", std::string{});
|
|
if (cfg.contains("context")) {
|
|
task_config["context"].update(cfg["context"]);
|
|
}
|
|
if (auto node = CreateFromRegistry<Node>(task_config); node) {
|
|
nodes_.push_back(std::move(node).value());
|
|
node_input_idx_.push_back(UpdateBindings(nodes_.back()->inputs(), kRead));
|
|
node_output_idx_.push_back(UpdateBindings(nodes_.back()->outputs(), kWrite));
|
|
} else {
|
|
ERROR("could not create {}:{}", name, type);
|
|
throw_exception(eFail);
|
|
}
|
|
}
|
|
output_idx_ = UpdateBindings(outputs(), kRead);
|
|
}
|
|
|
|
void Pipeline::Build(TaskGraph& graph) {
|
|
graph
|
|
.Add([this](Context& ctx) -> Result<void> {
|
|
ctx.current().array().resize(binding_name_to_idx_.size());
|
|
return success();
|
|
})
|
|
->set_name(fmt::format("{}.call", name()));
|
|
;
|
|
for (int index = 0; index < nodes_.size(); ++index) {
|
|
graph.Add([this, index](Context& ctx) { return Call(ctx, index); })
|
|
->set_name(fmt::format("{}.call", nodes_[index]->name()));
|
|
nodes_[index]->Build(graph);
|
|
graph.Add([this, index](Context& ctx) { return Ret(ctx, index); })
|
|
->set_name(fmt::format("{}.ret", nodes_[index]->name()));
|
|
}
|
|
graph
|
|
.Add([this](Context& ctx) -> Result<void> {
|
|
auto vars = std::move(ctx.current()).array();
|
|
return Gather(std::move(vars), output_idx_, ctx.current().array());
|
|
})
|
|
->set_name(fmt::format("{}.ret", name()));
|
|
}
|
|
|
|
std::vector<int> Pipeline::UpdateBindings(const vector<std::string>& names, BindingType type) {
|
|
std::vector<int> idxs;
|
|
for (const auto& name : names) {
|
|
auto it = binding_name_to_idx_.lower_bound(name);
|
|
if (it == binding_name_to_idx_.end() || it->first != name) {
|
|
if (type == kRead) {
|
|
ERROR("unknown binding name: {}", name);
|
|
throw_exception(eEntryNotFound);
|
|
} else {
|
|
auto index = static_cast<int>(binding_name_to_idx_.size());
|
|
it = binding_name_to_idx_.emplace_hint(it, name, index);
|
|
binding_idx_to_name_.emplace(index, name);
|
|
}
|
|
}
|
|
idxs.push_back(it->second);
|
|
}
|
|
return idxs;
|
|
}
|
|
|
|
Result<void> Pipeline::Call(Context& ctx, int idx) {
|
|
OUTCOME_TRY(auto&& args, Gather(ctx.current().array(), node_input_idx_[idx]));
|
|
ctx.push(std::move(args));
|
|
return success();
|
|
}
|
|
|
|
Result<void> Pipeline::Ret(Context& ctx, int idx) {
|
|
auto rets = ctx.pop().array();
|
|
return Scatter(std::move(rets), node_output_idx_[idx], ctx.current().array());
|
|
}
|
|
|
|
class PipelineCreator : public Creator<Node> {
|
|
public:
|
|
const char* GetName() const override { return "Pipeline"; }
|
|
int GetVersion() const override { return 0; }
|
|
std::unique_ptr<Node> Create(const Value& value) override {
|
|
return std::make_unique<Pipeline>(value);
|
|
}
|
|
};
|
|
|
|
REGISTER_MODULE(Node, PipelineCreator);
|
|
|
|
} // namespace mmdeploy::graph
|