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

Fixed unwanted PTX compilation result moving in CompileGpuAsmOrGetCached

Browse Source
This commit is contained in:
Georgiy Manuilov 2020-12-12 15:08:47 +03:00 committed by GitHub
parent a090692ef8
commit 8b35d402a8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/stream_executor/gpu/asm_compiler.cc
View File

@ -118,18 +118,22 @@ port::StatusOr<absl::Span<const uint8>> CompileGpuAsmOrGetCached(
compilation_options.ToTuple()};
auto it = ptx_cache.find(cache_key);
if (it == ptx_cache.end()) {
auto compiled = CompileGpuAsm(device_ordinal, ptx, compilation_options);
PtxCompilerResult compiled = CompileGpuAsm(device_ordinal, ptx,
compilation_options);
it = ptx_cache.emplace(cache_key, std::move(compiled)).first;
}
CHECK(it != ptx_cache.end());
// Failed compilation attempts are cached.
// Use separate status check and ValueOrDie invocation on ptx_cache
// entry to avoid value moving introduced by TF_ASSIGN_OR_RETURN.
if (TF_PREDICT_FALSE(!it->second.ok())) {
return it->second.status();
}
TF_ASSIGN_OR_RETURN(const std::vector<uint8>& compiled, it->second);
const std::vector<uint8>& compiled = it->second.ValueOrDie();
return absl::MakeSpan(compiled);
}