in src/core/scheduler/scheduler.cc [406:518]
void Graph::AddOperation(OpFunc &&op, const BlockVec &read_blocks,
const BlockVec &write_blocks, string op_name) {
dirty_ = true;
// if the size of both read_blocks and write_blocks is zero,
// this operation is used for synchronization
if (read_blocks.size() == 0 && write_blocks.size() == 0) {
AddSyncOp(std::move(op), op_name);
return;
}
// create new node
Node *node = new Node(nodes_.size(), std::move(op), op_name);
// create edges for read_blocks
for (size_t i = 0; i < read_blocks.size(); ++i) {
Block *blk = read_blocks[i];
Node *src_node = nullptr;
BlkInfo *blkInfo = nullptr;
// update leaf blocks
auto iter = leaf_blocks_.find(blk);
if (iter != leaf_blocks_.end()) {
leaf_blocks_.erase(iter);
}
// check if the block is already in the computational graph
auto it = blocks_.find(blk);
if (it == blocks_.end()) {
blkInfo = new BlkInfo(blocks_.size(), blk, BlockType::kInput);
blocks_[blk] = blkInfo;
} else {
blkInfo = it->second;
if (blkInfo->type_ == BlockType::kEnd) {
blkInfo->type_ = BlockType::kInter;
}
// update the existing edge, update dst node and create new edge
Edge *write_edge = blkInfo->write_edge_;
if (write_edge) {
if (!write_edge->dst_node_) {
// change the dst node of the write_edge
write_edge->dst_node_ = node;
node->AddInEdge(write_edge);
blkInfo->graph_ref_ += 1;
continue;
} else {
src_node = write_edge->src_node_;
}
}
}
// create new edge for new block
Edge *edge = new Edge(edges_.size(), blk, src_node, node);
blkInfo->graph_ref_ += 1;
if (src_node) {
src_node->AddOutEdge(edge);
}
node->AddInEdge(edge);
edges_.push_back(edge);
}
// update last node for write_blocks
for (size_t i = 0; i < write_blocks.size(); ++i) {
Block *blk = write_blocks[i];
BlkInfo *blkInfo = nullptr;
// update leaf blocks
leaf_blocks_.insert(blk);
auto it = blocks_.find(blk);
if (it == blocks_.end()) {
blkInfo = new BlkInfo(blocks_.size(), blk, BlockType::kEnd);
blocks_[blk] = blkInfo;
} else {
blkInfo = it->second;
if (blkInfo->type_ == BlockType::kInput) {
blkInfo->type_ = BlockType::kParam;
}
Edge *write_edge = blkInfo->write_edge_;
if (write_edge) {
if (!write_edge->dst_node_) {
write_edge->dst_node_ = node;
node->AddInEdge(write_edge);
} else {
Node *lastNode = write_edge->src_node_;
auto outEdges = lastNode->out_edges();
for (auto outEdge : outEdges) {
if (outEdge->blk_ == blk && outEdge->dst_node_ != node) {
Edge *edge =
new Edge(edges_.size(), blk, outEdge->dst_node_, node);
outEdge->dst_node_->AddOutEdge(edge);
node->AddInEdge(edge);
}
}
}
}
}
// create new edge for new block
Edge *edge = new Edge(edges_.size(), blk, node, nullptr);
blkInfo->write_edge_ = edge;
blkInfo->graph_ref_ += 1;
node->AddOutEdge(edge);
edges_.push_back(edge);
}
// add node into nodes
nodes_.push_back(node);
}