OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Save compile options with the compiled IFRT IR program to be used later for serialization

Browse Source 
PiperOrigin-RevId: 826301016
This commit is contained in:
A. Unique TensorFlower 2025-10-30 20:53:48 -07:00 committed by TensorFlower Gardener
parent e61bac51b1
commit a95c558dc4
2 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc vendored
View File

@ -305,11 +305,16 @@ absl::Status PopulateLayouts(mlir::ModuleOp mlir_module,
absl::StatusOr<CompiledIfrtIrProgram> CompiledIfrtIrProgram::Create( absl::StatusOr<CompiledIfrtIrProgram> CompiledIfrtIrProgram::Create(
std::unique_ptr<xla::ifrt::IfrtIRProgram> ifrt_ir_program, std::unique_ptr<xla::ifrt::IfrtIRProgram> ifrt_ir_program,
std::unique_ptr<xla::ifrt::IfrtIRCompileOptions> compile_options, std::unique_ptr<xla::ifrt::IfrtIRCompileOptions> ifrt_ir_compile_options,
xla::ifrt::Client* client, xla::ifrt::Client* client,
std::shared_ptr<xla::ifrt::AtomProgramCompiler> atom_program_compiler) { std::shared_ptr<xla::ifrt::AtomProgramCompiler> atom_program_compiler) {
TraceMe traceme([]() { return "ProgramCompiler::CompileForInterpreter"; }); 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<xla::ifrt::IfrtIRCompileOptions> compile_options =
std::move(ifrt_ir_compile_options);
std::vector<xla::ifrt::Device*> devices; std::vector<xla::ifrt::Device*> devices;
devices.reserve(compile_options->device_assignments.size()); devices.reserve(compile_options->device_assignments.size());
for (const auto& device_id : compile_options->device_assignments) { for (const auto& device_id : compile_options->device_assignments) {
@ -394,8 +399,7 @@ absl::StatusOr<CompiledIfrtIrProgram> CompiledIfrtIrProgram::Create(
} }
TF_RETURN_IF_ERROR(xla::ifrt::createOutlinedAtomProgramsToCompiledPipeline( TF_RETURN_IF_ERROR(xla::ifrt::createOutlinedAtomProgramsToCompiledPipeline(
pm, std::move(atom_program_compiler), compile_pipeline_options, pm, std::move(atom_program_compiler), compile_pipeline_options,
std::move(compile_options), atom_executable_map, compile_options, atom_executable_map, std::move(bound_executable_map)));
std::move(bound_executable_map)));
{ {
TraceMe traceme( TraceMe traceme(
@ -444,6 +448,7 @@ absl::StatusOr<CompiledIfrtIrProgram> CompiledIfrtIrProgram::Create(
/*donatable_input_indices=*/std::move(donatable_input_indices), /*donatable_input_indices=*/std::move(donatable_input_indices),
/*program=*/std::move(ifrt_ir_program), /*program=*/std::move(ifrt_ir_program),
/*device_assignments=*/std::move(device_assignments), /*device_assignments=*/std::move(device_assignments),
/*compile_options=*/compile_options,
}; };
} }

3
third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.h vendored
View File

@ -60,6 +60,9 @@ struct CompiledIfrtIrProgram {
// ids obtained from IFRT client. // ids obtained from IFRT client.
std::vector<xla::ifrt::DeviceId> device_assignments; std::vector<xla::ifrt::DeviceId> device_assignments;
// The compile options used to compile the program.
std::shared_ptr<xla::ifrt::IfrtIRCompileOptions> compile_options;
// Compiles an IFRT IR program. // Compiles an IFRT IR program.
static absl::StatusOr<CompiledIfrtIrProgram> Create( static absl::StatusOr<CompiledIfrtIrProgram> Create(
std::unique_ptr<xla::ifrt::IfrtIRProgram> ifrt_ir_program, std::unique_ptr<xla::ifrt::IfrtIRProgram> ifrt_ir_program,