Function GetHostInputLbis4OpNode

oneflow/core/graph/task_graph.cpp:485–512

}

void GetHostInputLbis4OpNode(const OpNode* op_node,
                             std::vector<LogicalBlobId>* host_mem_input_lbis) {
  host_mem_input_lbis->clear();
  if (op_node->op().op_conf().has_user_conf()) {
    const auto& user_conf = op_node->op().op_conf().user_conf();
    const auto& op_type_name = user_conf.op_type_name();
    if (user_op::UserOpHostMemoryInputRegistry::Get().HasHostMemoryInput(op_type_name)) {
      const auto& inputs = [&]() -> std::vector<std::pair<std::string, int32_t>> {
        const auto& arg_map = op_node->op().op_conf().user_conf().input();
        std::vector<std::pair<std::string, int32_t>> arg_vec;
        for (auto it = arg_map.begin(); it != arg_map.end(); ++it) {
          for (int32_t i = 0; i < it->second.s_size(); ++i) {
            arg_vec.emplace_back(std::make_pair(it->first, i));
          }
        }
        return arg_vec;
      }();
      for (const auto& pair : inputs) {
        if (user_op::UserOpHostMemoryInputRegistry::Get().IsHostMemoryInput4Op(
                op_type_name, pair.first, pair.second)) {
          const LogicalBlobId& host_input_lbi =
              GenLogicalBlobId(user_conf.input().at(pair.first).s(pair.second));
          host_mem_input_lbis->emplace_back(host_input_lbi);
        }
      }
    }
  }
}

HashMap<DeviceType, CreateSubTskGphBuilderFn>* GlobalDeviceType2CreateSubTskGphBuilderFn() {
  static HashMap<DeviceType, CreateSubTskGphBuilderFn>