From a090692ef8c9280981387e34ae2248b1cae0f1ee Mon Sep 17 00:00:00 2001
From: Georgiy Manuilov <george.manuilov@gmail.com>
Date: Thu, 19 Nov 2020 12:27:44 +0300
Subject: [PATCH] PR review fixes

---
 tensorflow/stream_executor/gpu/asm_compiler.cc | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/stream_executor/gpu/asm_compiler.cc b/tensorflow/stream_executor/gpu/asm_compiler.cc
index f3b789d28d9..0e2bd303f02 100644
--- a/tensorflow/stream_executor/gpu/asm_compiler.cc
+++ b/tensorflow/stream_executor/gpu/asm_compiler.cc
@@ -108,7 +108,7 @@ static void WarnIfBadPtxasVersion(const std::string& ptxas_path) {
 port::StatusOr<absl::Span<const uint8>> CompileGpuAsmOrGetCached(
     int device_ordinal, const char* ptx, GpuAsmOpts compilation_options) {
   using PtxCacheKey = std::tuple<int, std::string, GpuAsmOpts::PtxOptionsTuple>;
-  using PtxCompilerResult = port::StatusOr<std::vector<tensorflow::uint8>>;
+  using PtxCompilerResult = port::StatusOr<std::vector<uint8>>;
   static tensorflow::mutex ptx_cache_mutex(tensorflow::LINKER_INITIALIZED);
   static auto& ptx_cache TF_GUARDED_BY(ptx_cache_mutex) =
       *new absl::flat_hash_map<PtxCacheKey, PtxCompilerResult>();
@@ -122,13 +122,14 @@ port::StatusOr<absl::Span<const uint8>> CompileGpuAsmOrGetCached(
     it = ptx_cache.emplace(cache_key, std::move(compiled)).first;
   }
 
+  CHECK(it != ptx_cache.end());
+ 
   // Failed compilation attempts are cached.
   if (TF_PREDICT_FALSE(!it->second.ok())) {
     return it->second.status();
   }
 
-  CHECK(it != ptx_cache.end());
-  const std::vector<uint8>& compiled = it->second.ValueOrDie();
+  TF_ASSIGN_OR_RETURN(const std::vector<uint8>& compiled, it->second);
   return absl::MakeSpan(compiled);
 }