| 1230 | } |
| 1231 | |
| 1232 | StatusOr<bool> CopyInsertion::Run(HloModule* module) { |
| 1233 | // Copy insertion is performed in three steps: |
| 1234 | // |
| 1235 | // (1) Add copies conservatively to guarantee that there is no live-range |
| 1236 | // interference. This is done simplistically and usually results in more |
| 1237 | // copies than is strictly necessary. |
| 1238 | // |
| 1239 | // (2) Using a more fine-grained analysis, remove as many copies that were |
| 1240 | // added in (1) as possible while ensuring no live-range interference. |
| 1241 | // |
| 1242 | // (3) Add copies to resolve issues not related to live range interference |
| 1243 | // such as parameters and constants live out of the entry computation. |
| 1244 | // |
| 1245 | // We add copies then remove them (step (1) then (2)) rather than simply |
| 1246 | // adding only the copies that are necessary because, in general, it is |
| 1247 | // difficult to figure out the minimal set of copies to add once there is |
| 1248 | // interference. On the other hand, it is easy to determine if removing a copy |
| 1249 | // will introduce interference. |
| 1250 | // |
| 1251 | // The final copy insertion in (3) is done separately to simplify the |
| 1252 | // implementation of copy removal in (2) which is the most complicated part of |
| 1253 | // the pass. As is, copy removal only has to reason about live range |
| 1254 | // interference. If all copies were added in step (1) then copy removal would |
| 1255 | // also have to reason about things like constants and parameters live out of |
| 1256 | // the computation. |
| 1257 | std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module); |
| 1258 | if (!call_graph->IsFlattened()) { |
| 1259 | return FailedPrecondition( |
| 1260 | "Call graph must be flattened before copy insertion."); |
| 1261 | } |
| 1262 | |
| 1263 | TF_RETURN_IF_ERROR(AddCopiesToResolveInterference(module)); |
| 1264 | |
| 1265 | // Simplify the tuple structures introduced by the deep copies. This should be |
| 1266 | // done before removing copies (RemoveUnnecessaryCopies) because tuple |
| 1267 | // simplification changes dependencies in the graph which changes live range |
| 1268 | // interference in the graph. Also run DCE to remove the dead Tuple/GTE |
| 1269 | // instructions introduced by tuple simplification. |
| 1270 | TupleSimplifier tuple_simplifier; |
| 1271 | HloDCE dce; |
| 1272 | TF_RETURN_IF_ERROR(tuple_simplifier.Run(module).status()); |
| 1273 | TF_RETURN_IF_ERROR(dce.Run(module).status()); |
| 1274 | DumpHloModuleDuringPassIfEnabled( |
| 1275 | name(), "after adding copies to resolve interference", *module); |
| 1276 | |
| 1277 | TF_RETURN_IF_ERROR( |
| 1278 | RemoveUnnecessaryCopies(DependencyHloOrdering(module), module)); |
| 1279 | DumpHloModuleDuringPassIfEnabled(name(), "after removing unnecessary copies", |
| 1280 | *module); |
| 1281 | TF_RETURN_IF_ERROR(AddSpecialCaseCopies(*call_graph, module)); |
| 1282 | DumpHloModuleDuringPassIfEnabled(name(), "after adding special-case copies", |
| 1283 | *module); |
| 1284 | |
| 1285 | TF_RETURN_IF_ERROR(tuple_simplifier.Run(module).status()); |
| 1286 | TF_RETURN_IF_ERROR(dce.Run(module).status()); |
| 1287 | |
| 1288 | if (VLOG_IS_ON(1)) { |
| 1289 | int64 num_total_copies = 0; |
nothing calls this directly
no test coverage detected