OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 00:19:58 +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(
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,
std::shared_ptr<xla::ifrt::AtomProgramCompiler> 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<xla::ifrt::IfrtIRCompileOptions> compile_options =
std::move(ifrt_ir_compile_options);
std::vector<xla::ifrt::Device*> devices;
devices.reserve(compile_options->device_assignments.size());
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(
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> 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,
};
}

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.
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.
static absl::StatusOr<CompiledIfrtIrProgram> Create(
std::unique_ptr<xla::ifrt::IfrtIRProgram> ifrt_ir_program,