In cpu compiler's CompileAheadOfTime, pass ordering when compiling entry computation.
PiperOrigin-RevId: 158140349
This commit is contained in:
parent
f3f53e8b39
commit
2a61c1652e
@ -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));
|
||||||
|
|
||||||
|
@ -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:
|
||||||
//
|
//
|
||||||
|
Loading…
Reference in New Issue
Block a user