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,
|
||||
ir_emitter.EmitComputation(
|
||||
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.
|
||||
// This way we know what function to execute when we try to run code for
|
||||
// the Call instruction.
|
||||
@ -625,7 +626,8 @@ CpuCompiler::CompileAheadOfTime(std::vector<std::unique_ptr<HloModule>> modules,
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
llvm::Function * entry_function,
|
||||
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));
|
||||
|
||||
|
@ -67,17 +67,23 @@ class IrEmitter : public DfsHloVisitorWithDefault {
|
||||
~IrEmitter() override;
|
||||
|
||||
// Emit and return the given HLO computation as an LLVM IR
|
||||
// function. function_name_prefix is the desired name of the function. If the
|
||||
// name is not unique among already emitted functions then a suffix is
|
||||
// appended to make the name unique. is_entry_computation indicates that this
|
||||
// is the entry computation of the HLO module. If 'instruction_order' is given
|
||||
// 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.
|
||||
// function.
|
||||
//
|
||||
// function_name_prefix is the desired name of the function. If the name is
|
||||
// not unique among already emitted functions then a suffix is appended to
|
||||
// make the name unique.
|
||||
//
|
||||
// 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(
|
||||
HloComputation* computation, const string& function_name_prefix,
|
||||
bool is_entry_computation,
|
||||
std::vector<const HloInstruction*>* instruction_order = nullptr);
|
||||
std::vector<const HloInstruction*>* instruction_order);
|
||||
|
||||
protected:
|
||||
//
|
||||
|
Loading…
Reference in New Issue
Block a user