#ifndef GRAPH_UTILS_PM_OP_DEPTH_CHECK_PASS_HPP
#define GRAPH_UTILS_PM_OP_DEPTH_CHECK_PASS_HPP
#include <memory>
#include <queue>
#include <string>
#include <utility>
#include <unordered_set>
#include "graph/utils/pm/pass_base.hpp"
namespace dnnl {
namespace impl {
namespace graph {
namespace utils {
namespace pm {
class graph_op_depth_check_pass_t : public graph::pass::pass_base_t {
public:
explicit graph_op_depth_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<graph_op_depth_check_pass_t>(
std::move(pbackend), std::move(pname));
}
impl::status_t run(graph_t &agraph) override {
std::queue<op_t *> cur_layer_ops;
std::queue<op_t *> next_layer_ops;
std::unordered_set<op_t *> visited;
int64_t cur_op_depth = 0;
for (auto &root_op : agraph.get_output_ops()) {
cur_layer_ops.push(root_op);
root_op->set_attr<int64_t>(op_attr::op_depth, cur_op_depth);
}
while (!cur_layer_ops.empty()) {
cur_op_depth++;
while (!cur_layer_ops.empty()) {
op_t *cur_op = cur_layer_ops.front();
cur_layer_ops.pop();
auto &input_values = cur_op->get_input_values();
for (auto &input_value : input_values) {
if (input_value->has_producer()) {
auto next_layer_op = &(input_value->get_producer());
next_layer_op->set_attr<int64_t>(
op_attr::op_depth, cur_op_depth);
if (visited.find(next_layer_op) != visited.end())
continue;
next_layer_ops.push(next_layer_op);
visited.insert(next_layer_op);
}
}
}
swap(cur_layer_ops, next_layer_ops);
visited.clear();
}
return impl::status::success;
}
};
} } } } }
#endif