| 559 | } // namespace |
| 560 | |
| 561 | StatusOr<std::unique_ptr<Executable>> CpuCompiler::RunBackend( |
| 562 | std::unique_ptr<HloModule> module, se::StreamExecutor* stream_exec, |
| 563 | se::DeviceMemoryAllocator* /*device_allocator*/) { |
| 564 | VLOG(1) << "Compiling: " << module->name(); |
| 565 | XLA_SCOPED_LOGGING_TIMER( |
| 566 | absl::StrFormat("Compiling [%s] for CPU using JIT", module->name())); |
| 567 | auto slow_compile_alarm = SlowCompilationAlarm(); |
| 568 | |
| 569 | TF_RET_CHECK(stream_exec != nullptr); |
| 570 | absl::call_once(llvm_command_line_options_initialized, |
| 571 | &llvm_ir::InitializeLLVMCommandLineOptions, module->config()); |
| 572 | |
| 573 | ModuleHook pre_optimization_ir_hook; |
| 574 | ModuleHook post_optimization_ir_hook; |
| 575 | std::tie(pre_optimization_ir_hook, post_optimization_ir_hook) = |
| 576 | GetIRModuleHooks(*module, user_pre_optimization_hook_, |
| 577 | user_post_optimization_hook_); |
| 578 | |
| 579 | // Compile must be thread-safe so create a new LLVM context for the module. |
| 580 | auto llvm_context = absl::make_unique<llvm::LLVMContext>(); |
| 581 | auto llvm_module = |
| 582 | absl::make_unique<llvm::Module>("__compute_module", *llvm_context); |
| 583 | |
| 584 | auto jit = absl::make_unique<SimpleOrcJIT>( |
| 585 | CompilerTargetOptions(module->config()), |
| 586 | CodeGenOptLevel(module->config()), |
| 587 | options::OptimizeForSizeRequested(module->config()), |
| 588 | module->config().debug_options().xla_llvm_disable_expensive_passes(), |
| 589 | llvm_ir::GetCpuFastMathFlags(module->config()), pre_optimization_ir_hook, |
| 590 | post_optimization_ir_hook, |
| 591 | OrcJITPostCompilationHook::Create(module.get())); |
| 592 | llvm_module->setDataLayout(jit->data_layout()); |
| 593 | llvm_module->setTargetTriple(jit->target_triple().getTriple()); |
| 594 | |
| 595 | HloComputation* entry_computation = module->entry_computation(); |
| 596 | std::unordered_map<const HloInstruction*, int64> instruction_to_profile_idx; |
| 597 | std::unordered_map<const HloComputation*, int64> computation_to_profile_idx; |
| 598 | std::unique_ptr<HloProfileIndexMap> hlo_profile_index_map; |
| 599 | std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data; |
| 600 | if (module->config().hlo_profiling_enabled()) { |
| 601 | TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts( |
| 602 | *module, &instruction_to_profile_idx, &computation_to_profile_idx, |
| 603 | &hlo_profile_index_map, &hlo_profile_printer_data)); |
| 604 | } |
| 605 | |
| 606 | std::unique_ptr<Executable> cpu_executable; |
| 607 | |
| 608 | // Cache these flags here since we'll want to access them after the module's |
| 609 | // ownership is std::moved. |
| 610 | const bool embed_ir_in_executable = |
| 611 | module->config().debug_options().xla_embed_ir_in_executable(); |
| 612 | |
| 613 | // Select an order for emitting the HLO instructions for each |
| 614 | // computation. Using this sequence enables tighter buffer liveness analysis |
| 615 | // and reduced memory usage (as compared to using DependencyHloOrdering). |
| 616 | TF_ASSIGN_OR_RETURN(HloSchedule schedule, |
| 617 | ScheduleModule(module.get(), BufferSizeBytesFunction(), |
| 618 | ComputationSchedulerToModuleScheduler( |
nothing calls this directly
no test coverage detected