#ifndef GRAPH_UTILS_PM_DAG_CHECK_PASS_HPP
#define GRAPH_UTILS_PM_DAG_CHECK_PASS_HPP
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <unordered_map>
#include "graph/utils/pm/pass_base.hpp"
namespace dnnl {
namespace impl {
namespace graph {
namespace pass {
class dag_check_pass_t : public graph::pass::pass_base_t {
public:
explicit dag_check_pass_t(std::string pbackend, std::string pname)
: graph::pass::pass_base_t(std::move(pbackend), std::move(pname)) {}
static graph::pass::pass_base_ptr create(
std::string pbackend, std::string pname) {
return std::make_shared<dag_check_pass_t>(
std::move(pbackend), std::move(pname));
}
bool is_cycle_detected(const op_t *cur_op,
std::unordered_map<size_t, bool> &visited,
std::unordered_map<size_t, bool> &recursion_stack) {
if (!visited.at(cur_op->get_id())) {
visited[cur_op->get_id()] = true;
recursion_stack[cur_op->get_id()] = true;
const auto &inputs = cur_op->get_input_values();
for (auto it = inputs.begin(); it != inputs.end(); ++it) {
if (!((*it)->has_producer())) continue;
op_t &next_op = (*it)->get_producer();
if (is_cycle_detected(&next_op, visited, recursion_stack))
return true;
}
recursion_stack[cur_op->get_id()] = false;
} else if (recursion_stack.at(cur_op->get_id())) {
return true;
}
return false;
}
status_t run(graph_t &agraph) override {
std::unordered_map<size_t, bool> visited;
std::unordered_map<size_t, bool> recursion_stack;
for (const std::shared_ptr<op_t> &aop : agraph.get_ops()) {
visited[aop->get_id()] = false;
recursion_stack[aop->get_id()] = false;
}
for (const std::shared_ptr<op_t> &aop : agraph.get_ops()) {
if (is_cycle_detected(aop.get(), visited, recursion_stack))
return status::invalid_graph;
}
return status::success;
}
};
} } } }
#endif