| 55 | } |
| 56 | |
| 57 | StatusOr<std::vector<std::unique_ptr<Executable>>> FailoverCompiler::Compile( |
| 58 | std::unique_ptr<HloModuleGroup> module_group, |
| 59 | std::vector<std::vector<se::StreamExecutor*>> stream_execs, |
| 60 | se::DeviceMemoryAllocator* device_allocator) { |
| 61 | std::vector<std::unique_ptr<Executable>> result; |
| 62 | std::vector<std::unique_ptr<HloModule>> modules = |
| 63 | module_group->ConsumeModules(); |
| 64 | for (size_t i = 0; i < modules.size(); i++) { |
| 65 | if (stream_execs[i].size() != 1) { |
| 66 | // This is not supported by GPU compiler anyway. |
| 67 | return Unimplemented( |
| 68 | "Model partitioning not implemented for the failover compiler!"); |
| 69 | } |
| 70 | auto executable = [stream_execs, device_allocator, i, |
| 71 | this](std::unique_ptr<HloModule> module) |
| 72 | -> StatusOr<std::unique_ptr<Executable>> { |
| 73 | TF_ASSIGN_OR_RETURN( |
| 74 | auto processed_module, |
| 75 | primary_->RunHloPasses(std::move(module), stream_execs[i][0], |
| 76 | device_allocator)); |
| 77 | TF_ASSIGN_OR_RETURN( |
| 78 | auto result, |
| 79 | primary_->RunBackend(std::move(processed_module), stream_execs[i][0], |
| 80 | device_allocator)); |
| 81 | return result; |
| 82 | }(modules[i]->Clone()); |
| 83 | |
| 84 | if (IsUnimplemented(executable)) { |
| 85 | VLOG(2) << "Compile resulted in " << executable.status() |
| 86 | << ", falling back to secondary backend"; |
| 87 | TF_ASSIGN_OR_RETURN( |
| 88 | modules[i], |
| 89 | secondary_->RunHloPasses(std::move(modules[i]), stream_execs[i][0], |
| 90 | device_allocator)); |
| 91 | TF_ASSIGN_OR_RETURN( |
| 92 | executable, |
| 93 | secondary_->RunBackend(std::move(modules[i]), stream_execs[i][0], |
| 94 | device_allocator)); |
| 95 | } |
| 96 | |
| 97 | if (!executable.ok()) { |
| 98 | return executable.status(); |
| 99 | } |
| 100 | |
| 101 | result.push_back(std::move(executable.ValueOrDie())); |
| 102 | } |
| 103 | |
| 104 | return {std::move(result)}; |
| 105 | } |
| 106 | |
| 107 | StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>> |
| 108 | FailoverCompiler::CompileAheadOfTime( |
nothing calls this directly
no test coverage detected