Runs a heap simulation for the given 'computation', assuming the given 'instruction_sequence'.
| 147 | // Runs a heap simulation for the given 'computation', assuming the given |
| 148 | // 'instruction_sequence'. |
| 149 | Status HeapSimulator::RunComputation( |
| 150 | const HloComputation& computation, |
| 151 | const HloInstructionSequence& instruction_sequence, |
| 152 | const HloAliasAnalysis& alias_analysis, HloLiveRange* hlo_live_range) { |
| 153 | XLA_VLOG_LINES(1, computation.parent()->ToString()); |
| 154 | XLA_VLOG_LINES(2, computation.ToString()); |
| 155 | |
| 156 | VLOG(1) << hlo_live_range->ToString(); |
| 157 | |
| 158 | HloDataflowAnalysis& dataflow_analysis = alias_analysis.dataflow_analysis(); |
| 159 | |
| 160 | // Record the buffer define/free event for each time step. We free all |
| 161 | // remaining buffers (entry parameter, etc) after the program has finished |
| 162 | // running, so we set the size of to program_end_time + 1. |
| 163 | std::vector<std::vector<const HloValue*>> buffers_defined( |
| 164 | hlo_live_range->schedule_end_time() + 1); |
| 165 | std::vector<std::vector<const HloValue*>> buffers_freed( |
| 166 | hlo_live_range->schedule_end_time() + 1); |
| 167 | |
| 168 | // values_to_assign tracks the HloValues that we need to assign a buffer to. |
| 169 | // Note that we only need to assign a buffer to a value when both of the |
| 170 | // following conditions are met: |
| 171 | // |
| 172 | // - The user specifically asks us to assign a buffer to a set of HloValues, |
| 173 | // and the value is in the set. If the user don't provide such a set, by |
| 174 | // default we assign buffer to all HloValues. |
| 175 | // |
| 176 | // - If the instruction is in a nested call of the current computation, only |
| 177 | // assign a buffer if we are doing global heap simulation. |
| 178 | std::vector<const HloValue*> values_to_assign; |
| 179 | values_to_assign.reserve(dataflow_analysis.values().size()); |
| 180 | |
| 181 | for (const HloValue* value : dataflow_analysis.values()) { |
| 182 | // Ignore buffers that are not tracked. |
| 183 | if (hlo_live_range->instruction_schedule().count( |
| 184 | value->defining_instruction()) == 0) { |
| 185 | continue; |
| 186 | } |
| 187 | if (IgnoreBuffer(value)) { |
| 188 | continue; |
| 189 | } |
| 190 | values_to_assign.push_back(value); |
| 191 | } |
| 192 | |
| 193 | auto& buffer_live_ranges = hlo_live_range->buffer_live_ranges(); |
| 194 | |
| 195 | absl::c_sort(values_to_assign, |
| 196 | [&](const HloValue* value1, const HloValue* value2) { |
| 197 | const auto& live_range1 = buffer_live_ranges.at(value1); |
| 198 | const auto& live_range2 = buffer_live_ranges.at(value2); |
| 199 | return std::forward_as_tuple(live_range1.start, |
| 200 | live_range1.end, value1->id()) < |
| 201 | std::forward_as_tuple(live_range2.start, |
| 202 | live_range2.end, value2->id()); |
| 203 | }); |
| 204 | |
| 205 | // For each value that we need to assign a buffer to, add the define and free |
| 206 | // events. |
no test coverage detected