diff --git a/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc b/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc index c85997a33a0..49ea8fe1b29 100644 --- a/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc +++ b/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc @@ -305,11 +305,16 @@ absl::Status PopulateLayouts(mlir::ModuleOp mlir_module, absl::StatusOr CompiledIfrtIrProgram::Create( std::unique_ptr ifrt_ir_program, - std::unique_ptr compile_options, + std::unique_ptr ifrt_ir_compile_options, xla::ifrt::Client* client, std::shared_ptr atom_program_compiler) { TraceMe traceme([]() { return "ProgramCompiler::CompileForInterpreter"; }); + // Sharing the compile options with the passes and when pipeline is done add + // it to the CompiledIfrtIrProgram. + std::shared_ptr compile_options = + std::move(ifrt_ir_compile_options); + std::vector devices; devices.reserve(compile_options->device_assignments.size()); for (const auto& device_id : compile_options->device_assignments) { @@ -394,8 +399,7 @@ absl::StatusOr CompiledIfrtIrProgram::Create( } TF_RETURN_IF_ERROR(xla::ifrt::createOutlinedAtomProgramsToCompiledPipeline( pm, std::move(atom_program_compiler), compile_pipeline_options, - std::move(compile_options), atom_executable_map, - std::move(bound_executable_map))); + compile_options, atom_executable_map, std::move(bound_executable_map))); { TraceMe traceme( @@ -444,6 +448,7 @@ absl::StatusOr CompiledIfrtIrProgram::Create( /*donatable_input_indices=*/std::move(donatable_input_indices), /*program=*/std::move(ifrt_ir_program), /*device_assignments=*/std::move(device_assignments), + /*compile_options=*/compile_options, }; } diff --git a/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.h b/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.h index 5896dfc52c4..50975062748 100644 --- a/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.h +++ b/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.h @@ -60,6 +60,9 @@ struct CompiledIfrtIrProgram { // ids obtained from IFRT client. std::vector device_assignments; + // The compile options used to compile the program. + std::shared_ptr compile_options; + // Compiles an IFRT IR program. static absl::StatusOr Create( std::unique_ptr ifrt_ir_program,