flow.h

Go to the documentation of this file.

#ifndef FLOW_INCLUDE_FLOW_H_
#define FLOW_INCLUDE_FLOW_H_

#include <string>
#include <vector>
#include <map>
#include <queue>
#include "flow.pb.h"

#include "node.h"
#include "task.h"
#include "schedule.h"
#include "tensor.hpp"

class Flow {
public:
    Flow();

    ~Flow();

    void init(std::vector<std::string> graphPaths,
        DataType precision,
        AffinityPolicy affinityPolicy = AFFINITY_CPU_HIGH_PERFORMANCE,
        int cpuThreads = 1,
        bool useGPU = true);

    void enqueue(Task task);

    std::vector<Task> dequeue(bool block = false);

    unsigned int size();

private:
    Schedule<flow::GraphParameter, Node, Tensor> schedule;
    std::queue<std::shared_ptr<Task>> tasks;
};
#endif  // FLOW_INCLUDE_FLOW_H_