experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Make ModuleHook return void rather than a Status.

Browse Source 
We were CHECK_OK'ing these return values anyway, so using Status doesn't buy us
much.

PiperOrigin-RevId: 238250643
This commit is contained in:
Justin Lebar 2019-03-13 10:19:02 -07:00 committed by TensorFlower Gardener
parent 96e08caa7d
commit 3ce012557b
4 changed files with 7 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/compiler/xla/service/cpu/compiler_functor.cc
View File

@ -94,7 +94,7 @@ std::unique_ptr<llvm::MemoryBuffer> CompilerFunctor::operator()(
XLA_VLOG_LINES(2, llvm_ir::DumpModuleToString(module));
if (pre_optimization_hook_) {
TF_CHECK_OK(pre_optimization_hook_(module));
pre_optimization_hook_(module);
}
// Add the appropriate TargetLibraryInfo and TargetTransformInfo.
@ -138,10 +138,7 @@ std::unique_ptr<llvm::MemoryBuffer> CompilerFunctor::operator()(
XLA_VLOG_LINES(2, llvm_ir::DumpModuleToString(module));
if (post_optimization_hook_) {
// TODO(jlebar): If we're going to CHECK these return values, then we might
// as well make them return void and move the CHECKs into the hooks
// themselves.
TF_CHECK_OK(post_optimization_hook_(module));
post_optimization_hook_(module);
}
// Generate code.

5
tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
View File

@ -447,10 +447,9 @@ std::pair<LLVMCompiler::ModuleHook, LLVMCompiler::ModuleHook> GetIRModuleHooks(
const auto& user_hook =
!optimized ? user_pre_optimization_hook : user_post_optimization_hook;
if (user_hook) {
TF_RETURN_IF_ERROR(user_hook(llvm_module));
user_hook(llvm_module);
}
llvm_ir::DumpIrIfEnabled(*hlo_module_ptr, llvm_module, optimized);
return Status::OK();
};
return {[hook](const llvm::Module& llvm_module) {
return hook(/*optimized=*/false, llvm_module);
@ -886,7 +885,7 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group,
{
Status verify_status = VerifyLlvmModule(llvm_module);
if (!verify_status.ok() && pre_optimization_ir_hook) {
pre_optimization_ir_hook(llvm_module).IgnoreError();
pre_optimization_ir_hook(llvm_module);
}
TF_RETURN_IF_ERROR(verify_status);
}

4
tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
View File

@ -694,7 +694,7 @@ StatusOr<std::unique_ptr<Executable>> NVPTXCompiler::RunBackend(
}
if (user_pre_optimization_hook_) {
TF_CHECK_OK(user_pre_optimization_hook_(llvm_module));
user_pre_optimization_hook_(llvm_module);
}
string ir_module_string_before_opt;
const bool embed_ir_in_executable =
@ -753,7 +753,7 @@ StatusOr<std::unique_ptr<Executable>> NVPTXCompiler::RunBackend(
llvm_ir::DumpIrIfEnabled(*module, llvm_module, /*optimized=*/true);
if (user_post_optimization_hook_) {
TF_CHECK_OK(user_post_optimization_hook_(llvm_module));
user_post_optimization_hook_(llvm_module);
}
// Write PTX to IR dump directory, if IR dumping was requested.
if (DumpingEnabledForHloModule(*module)) {

2
tensorflow/compiler/xla/service/llvm_compiler.h
View File

@ -37,7 +37,7 @@ class LLVMCompiler : public Compiler {
// A callback of this type can be run before and/or after IR-level
// optimization to e.g. dump out the generated IR to disk or gather some
// statistics.
using ModuleHook = std::function<Status(const llvm::Module&)>;
using ModuleHook = std::function<void(const llvm::Module&)>;
void SetPreOptimizationHook(ModuleHook hook) {
CHECK(!user_pre_optimization_hook_)