112 lines
2.4 KiB
C++
112 lines
2.4 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#ifndef MMDEPLOY_SRC_EXPERIMENTAL_PIPELINE_IR_H_
|
|
#define MMDEPLOY_SRC_EXPERIMENTAL_PIPELINE_IR_H_
|
|
|
|
#include "core/model.h"
|
|
#include "core/module.h"
|
|
#include "core/registry.h"
|
|
#include "core/status_code.h"
|
|
#include "mpl/span.h"
|
|
#include "utils/formatter.h"
|
|
|
|
#if MMDEPLOY_USE_TASKFLOW
|
|
#include "taskflow/taskflow.hpp"
|
|
#endif
|
|
|
|
namespace mmdeploy {
|
|
|
|
namespace graph {
|
|
|
|
using std::pair;
|
|
using std::string;
|
|
using std::unique_ptr;
|
|
using std::vector;
|
|
|
|
class TaskGraph;
|
|
class Node;
|
|
|
|
class MMDEPLOY_API Context {
|
|
public:
|
|
explicit Context(TaskGraph* graph) : graph_(graph) {}
|
|
|
|
Value& current() { return context_.back(); }
|
|
|
|
void push(Value&& ctx) { context_.push_back(std::move(ctx)); }
|
|
|
|
Value pop() {
|
|
auto ctx = std::move(context_.back());
|
|
context_.pop_back();
|
|
return ctx;
|
|
}
|
|
|
|
size_t size() const noexcept { return context_.size(); }
|
|
bool empty() const noexcept { return context_.empty(); }
|
|
|
|
std::vector<Result<Value>> Execute(Span<std::function<Result<Value>()>> tasks);
|
|
|
|
private:
|
|
vector<Value> context_;
|
|
TaskGraph* graph_;
|
|
};
|
|
|
|
class MMDEPLOY_API TaskGraph {
|
|
friend class Context;
|
|
|
|
public:
|
|
using TaskFunction = std::function<Result<void>(Context& ctx)>;
|
|
|
|
class Handle {
|
|
public:
|
|
const std::string& name() const noexcept { return name_; }
|
|
void set_name(const std::string& name) { name_ = name; }
|
|
|
|
private:
|
|
std::string name_;
|
|
};
|
|
|
|
~TaskGraph();
|
|
|
|
TaskGraph() = default;
|
|
TaskGraph(const TaskGraph&) = delete;
|
|
TaskGraph& operator=(const TaskGraph&) = delete;
|
|
|
|
Handle* Add(TaskFunction fn);
|
|
|
|
Result<Value> Run(Value inputs);
|
|
|
|
private:
|
|
std::vector<Result<Value>> Execute(Span<std::function<Result<Value>()>> tasks);
|
|
|
|
vector<TaskFunction> function_;
|
|
vector<unique_ptr<Handle>> handle_;
|
|
#if MMDEPLOY_USE_TASKFLOW
|
|
tf::Executor executor_;
|
|
#endif
|
|
// profiling utils
|
|
std::vector<double> time_;
|
|
int64_t count_{};
|
|
};
|
|
|
|
class MMDEPLOY_API Node {
|
|
public:
|
|
virtual ~Node() = default;
|
|
virtual void Build(TaskGraph& graph) = 0;
|
|
const std::vector<std::string>& inputs() const noexcept { return inputs_; }
|
|
const std::vector<std::string>& outputs() const noexcept { return outputs_; }
|
|
const std::string& name() const noexcept { return name_; }
|
|
|
|
protected:
|
|
std::string name_;
|
|
std::vector<std::string> inputs_;
|
|
std::vector<std::string> outputs_;
|
|
};
|
|
|
|
} // namespace graph
|
|
|
|
MMDEPLOY_DECLARE_REGISTRY(graph::Node);
|
|
|
|
} // namespace mmdeploy
|
|
|
|
#endif // MMDEPLOY_SRC_EXPERIMENTAL_PIPELINE_IR_H_
|