Checks various invariants of channel instructions (send/recv and collectives).
| 1403 | // Checks various invariants of channel instructions (send/recv and |
| 1404 | // collectives). |
| 1405 | Status VerifyChannels(const HloModule& module) { |
| 1406 | absl::flat_hash_map<int64, std::vector<const HloInstruction*>> |
| 1407 | channel_instructions; |
| 1408 | |
| 1409 | // Send/Recv instruction must have a single user: the corresponding |
| 1410 | // SendDone/RecvDone. with matching channel. |
| 1411 | for (const HloComputation* computation : module.computations()) { |
| 1412 | for (const HloInstruction* instruction : computation->instructions()) { |
| 1413 | auto channel_instr = DynCast<HloChannelInstruction>(instruction); |
| 1414 | if (!channel_instr || !channel_instr->channel_id()) { |
| 1415 | continue; |
| 1416 | } |
| 1417 | channel_instructions[*channel_instr->channel_id()].push_back(instruction); |
| 1418 | |
| 1419 | switch (instruction->opcode()) { |
| 1420 | case HloOpcode::kSend: { |
| 1421 | TF_RET_CHECK(instruction->users().size() == 1); |
| 1422 | const HloInstruction* send_done = instruction->users().front(); |
| 1423 | TF_RET_CHECK(send_done->opcode() == HloOpcode::kSendDone); |
| 1424 | TF_RETURN_IF_ERROR(CheckSameChannel(instruction, send_done)); |
| 1425 | TF_RETURN_IF_ERROR(CheckSameIsHostTransfer(instruction, send_done)); |
| 1426 | break; |
| 1427 | } |
| 1428 | case HloOpcode::kRecv: { |
| 1429 | TF_RET_CHECK(instruction->users().size() == 1); |
| 1430 | const HloInstruction* recv_done = instruction->users().front(); |
| 1431 | TF_RET_CHECK(recv_done->opcode() == HloOpcode::kRecvDone); |
| 1432 | TF_RETURN_IF_ERROR(CheckSameChannel(instruction, recv_done)); |
| 1433 | TF_RETURN_IF_ERROR(CheckSameIsHostTransfer(instruction, recv_done)); |
| 1434 | break; |
| 1435 | } |
| 1436 | case HloOpcode::kSendDone: |
| 1437 | TF_RET_CHECK(instruction->operands().size() == 1); |
| 1438 | TF_RET_CHECK(instruction->operand(0)->opcode() == HloOpcode::kSend); |
| 1439 | break; |
| 1440 | case HloOpcode::kRecvDone: |
| 1441 | TF_RET_CHECK(instruction->operands().size() == 1); |
| 1442 | TF_RET_CHECK(instruction->operand(0)->opcode() == HloOpcode::kRecv); |
| 1443 | break; |
| 1444 | default: |
| 1445 | break; |
| 1446 | } |
| 1447 | } |
| 1448 | } |
| 1449 | |
| 1450 | // Iterate over each channel to check invariants. |
| 1451 | for (auto& pair : channel_instructions) { |
| 1452 | auto& instructions = pair.second; |
| 1453 | const HloInstruction* first = instructions[0]; |
| 1454 | auto sendrecv = DynCast<HloSendRecvInstruction>(first); |
| 1455 | if (sendrecv) { |
| 1456 | absl::flat_hash_set<HloOpcode> opcodes; |
| 1457 | for (const HloInstruction* instr : instructions) { |
| 1458 | opcodes.insert(instr->opcode()); |
| 1459 | auto cast = DynCast<HloSendRecvInstruction>(instr); |
| 1460 | TF_RET_CHECK(cast != nullptr) |
| 1461 | << "channel " << pair.first |
| 1462 | << " is used for different types of channel instructions"; |
no test coverage detected