| 1510 | namespace { |
| 1511 | |
| 1512 | Maybe<void> InferOpOutSbpParallel( |
| 1513 | Operator* op, const OpNodeSignature& upstream_signature, |
| 1514 | const std::function<const BlobDesc&(const std::string&)>& ConstBlobDesc4Ibn, |
| 1515 | const SbpSignature& sbp_sig_conf, const ParallelDesc& parallel_desc) { |
| 1516 | const auto& SbpParallel4Ibn = [&](const std::string& ibn) -> const SbpParallel* { |
| 1517 | const auto& map = upstream_signature.sbp_signature().bn_in_op2sbp_parallel(); |
| 1518 | return &map.at(ibn); |
| 1519 | }; |
| 1520 | HashMap<std::string, SbpInferHint> ibn2sbp_infer_hint; |
| 1521 | for (const std::string& ibn : op->input_bns()) { |
| 1522 | const ParallelDesc* pd = ¶llel_desc; |
| 1523 | const BlobDesc* logical_blob_desc = &ConstBlobDesc4Ibn(ibn); |
| 1524 | const SbpParallel* sbp_parallel = SbpParallel4Ibn(ibn); |
| 1525 | ibn2sbp_infer_hint.emplace(ibn, SbpInferHint(pd, logical_blob_desc, sbp_parallel)); |
| 1526 | } |
| 1527 | SbpSignature sbp_signature; |
| 1528 | JUST(op->InferSbpSignature(&sbp_signature, sbp_sig_conf, ibn2sbp_infer_hint)); |
| 1529 | JUST(op->FillSbpSignature(sbp_signature)); |
| 1530 | return Maybe<void>::Ok(); |
| 1531 | } |
| 1532 | |
| 1533 | Maybe<void> InferLocalSignature(Operator* op, const OpNodeSignature& upstream_signature, |
| 1534 | bool is_local, const ParallelDesc& parallel_desc) { |
no test coverage detected