In cpu compiler's CompileAheadOfTime, pass ordering when compiling entry computation.

PiperOrigin-RevId: 158140349
This commit is contained in:
A. Unique TensorFlower 2017-06-06 07:57:37 -07:00 committed by TensorFlower Gardener
parent f3f53e8b39
commit 2a61c1652e
2 changed files with 18 additions and 10 deletions

View File

@ -397,7 +397,8 @@ StatusOr<std::unique_ptr<Executable>> CpuCompiler::Compile(
llvm::Function * ir_function, llvm::Function * ir_function,
ir_emitter.EmitComputation( ir_emitter.EmitComputation(
embedded_computation, embedded_computation->name(), embedded_computation, embedded_computation->name(),
/*is_entry_computation=*/computation_is_parallel)); /*is_entry_computation=*/computation_is_parallel,
/*instruction_order=*/nullptr));
// If this computation is parallel, remember it in the function name map. // If this computation is parallel, remember it in the function name map.
// This way we know what function to execute when we try to run code for // This way we know what function to execute when we try to run code for
// the Call instruction. // the Call instruction.
@ -625,7 +626,8 @@ CpuCompiler::CompileAheadOfTime(std::vector<std::unique_ptr<HloModule>> modules,
TF_ASSIGN_OR_RETURN( TF_ASSIGN_OR_RETURN(
llvm::Function * entry_function, llvm::Function * entry_function,
ir_emitter.EmitComputation(computation, entry_point_name, ir_emitter.EmitComputation(computation, entry_point_name,
/*is_entry_computation=*/true)); /*is_entry_computation=*/true,
&module_sequence.at(computation)));
entry_function->setName(llvm_ir::AsStringRef(entry_point_name)); entry_function->setName(llvm_ir::AsStringRef(entry_point_name));

View File

@ -67,17 +67,23 @@ class IrEmitter : public DfsHloVisitorWithDefault {
~IrEmitter() override; ~IrEmitter() override;
// Emit and return the given HLO computation as an LLVM IR // Emit and return the given HLO computation as an LLVM IR
// function. function_name_prefix is the desired name of the function. If the // function.
// name is not unique among already emitted functions then a suffix is //
// appended to make the name unique. is_entry_computation indicates that this // function_name_prefix is the desired name of the function. If the name is
// is the entry computation of the HLO module. If 'instruction_order' is given // not unique among already emitted functions then a suffix is appended to
// then the HLO instructions are emitted in the given order. In this case, // make the name unique.
// 'instruction_order' must be a topological sort of the set of nodes //
// accessible from the root of the computation. // is_entry_computation indicates that this is the entry computation of the
// HLO module.
//
// If 'instruction_order' is not NULL, then the HLO instructions are emitted
// in the given order. In this case, 'instruction_order' must be a
// topological sort of the set of nodes accessible from the root of the
// computation.
StatusOr<llvm::Function*> EmitComputation( StatusOr<llvm::Function*> EmitComputation(
HloComputation* computation, const string& function_name_prefix, HloComputation* computation, const string& function_name_prefix,
bool is_entry_computation, bool is_entry_computation,
std::vector<const HloInstruction*>* instruction_order = nullptr); std::vector<const HloInstruction*>* instruction_order);
protected: protected:
// //