Method ParseComputations

tensorflow/compiler/xla/service/hlo_parser.cc:533–567

computations ::= (computation)+

// computations ::= (computation)+
bool HloParserImpl::ParseComputations(HloModule* module) {
  HloComputation* entry_computation = nullptr;
  do {
    if (!ParseComputation(&entry_computation)) {
      return false;
    }
  } while (lexer_.GetKind() != TokKind::kEof);

  for (int i = 0; i < computations_.size(); i++) {
    // If entry_computation is not nullptr, it means the computation it pointed
    // to is marked with "ENTRY"; otherwise, no computation is marked with
    // "ENTRY", and we use the last computation as the entry computation. We
    // add the non-entry computations as embedded computations to the module.
    if ((entry_computation != nullptr &&
         computations_[i].get() != entry_computation) ||
        (entry_computation == nullptr && i != computations_.size() - 1)) {
      module->AddEmbeddedComputation(std::move(computations_[i]));
      continue;
    }
    auto computation = module->AddEntryComputation(std::move(computations_[i]));
    // The parameters and result layouts were set to default layout. Here we
    // set the layouts to what the hlo text says.
    for (int p = 0; p < computation->num_parameters(); p++) {
      const Shape& param_shape = computation->parameter_instruction(p)->shape();
      TF_CHECK_OK(module->mutable_entry_computation_layout()
                      ->mutable_parameter_layout(p)
                      ->CopyLayoutFromShape(param_shape));
    }
    const Shape& result_shape = computation->root_instruction()->shape();
    TF_CHECK_OK(module->mutable_entry_computation_layout()
                    ->mutable_result_layout()
                    ->CopyLayoutFromShape(result_shape));
  }
  return true;
}

// computation ::= ('ENTRY')? name (param_list_to_shape)? instruction_list
bool HloParserImpl::ParseComputation(HloComputation** entry_computation) {