Method GetFastestInputOutputTimeShape

oneflow/core/graph/task_node.cpp:128–141

}

std::shared_ptr<Shape> TaskNode::GetFastestInputOutputTimeShape() const {
  std::shared_ptr<Shape> shape;
  auto UpdateRetShape = [&](TaskEdge* edge) {
    for (const auto& regst : edge->GetRegsts()) {
      if (!shape || shape->elem_cnt() < regst->data_regst_time_shape()->elem_cnt()) {
        shape = regst->data_regst_time_shape();
      }
    }
  };
  ForEachOutDataEdge(UpdateRetShape);
  if (shape) { return shape; }
  ForEachInDataEdge(UpdateRetShape);
  return shape;
}

void TaskNode::ForEachConsumedDataRegst(
    const std::function<void(const std::string&, const RegstDesc*)>& Handler) const {