| 1058 | } |
| 1059 | |
| 1060 | bool MemoryUsageTracker::Check() const { |
| 1061 | auto elements_are_unique = [](const BufferIdList& vec) { |
| 1062 | return vec.size() == std::set<BufferId>(vec.begin(), vec.end()).size(); |
| 1063 | }; |
| 1064 | |
| 1065 | // Verify buffers_defined per instruction. |
| 1066 | for (auto* instruction : computation_->instructions()) { |
| 1067 | const BufferIdList& defined_buffers = |
| 1068 | instruction_list_.GetItem(instruction)->buffers_defined; |
| 1069 | CHECK(elements_are_unique(defined_buffers)) |
| 1070 | << "Instruction " << instruction->name() |
| 1071 | << " does not have unique defined buffers: " |
| 1072 | << absl::StrJoin( |
| 1073 | defined_buffers, ", ", [this](string* out, BufferId buffer_id) { |
| 1074 | absl::StrAppend(out, buffers_.at(buffer_id).ToString()); |
| 1075 | }); |
| 1076 | |
| 1077 | for (const Buffer& buffer : buffers_) { |
| 1078 | if (buffer.defining_instruction->instruction == instruction) { |
| 1079 | CHECK(absl::c_linear_search(defined_buffers, buffer.id)) |
| 1080 | << "Instruction " << instruction->name() |
| 1081 | << " defined buffers is missing: " << buffer.ToString(); |
| 1082 | } |
| 1083 | } |
| 1084 | } |
| 1085 | |
| 1086 | // Verify buffers_used per instruction. |
| 1087 | for (auto* instruction : computation_->instructions()) { |
| 1088 | const BufferIdList& used_buffers = |
| 1089 | instruction_list_.GetItem(instruction)->buffers_used; |
| 1090 | CHECK(elements_are_unique(used_buffers)) |
| 1091 | << "Instruction " << instruction->name() |
| 1092 | << " does not have unique used buffers: " |
| 1093 | << absl::StrJoin( |
| 1094 | used_buffers, ", ", [this](string* out, BufferId buffer_id) { |
| 1095 | absl::StrAppend(out, buffers_.at(buffer_id).ToString()); |
| 1096 | }); |
| 1097 | } |
| 1098 | for (const Buffer& buffer : buffers_) { |
| 1099 | int64 unfinished_uses = 0; |
| 1100 | for (Item* user : buffer.users) { |
| 1101 | const BufferIdList& used_buffers = user->buffers_used; |
| 1102 | CHECK(absl::c_linear_search(used_buffers, buffer.id)) |
| 1103 | << "Instruction " << user->instruction->name() |
| 1104 | << " used buffers is missing " << buffer.ToString(); |
| 1105 | if (!IsFinished(user)) { |
| 1106 | unfinished_uses++; |
| 1107 | } |
| 1108 | } |
| 1109 | CHECK_EQ(buffer.unfinished_user_count, unfinished_uses) |
| 1110 | << "Incorrect unplaced use count for " << buffer.ToString(); |
| 1111 | } |
| 1112 | return true; |
| 1113 | } |
| 1114 | |
| 1115 | // Computes and returns the cost of rematerializing the given instruction. |
| 1116 | // Cost per rematerialized instruction is defined as: |