From 2ed84d32c46dffba6523f23a4d2dd9d7a54095cf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 11 Sep 2023 10:03:44 -0700
Subject: [PATCH] change triton autotuning AutotunerUtil::AddResult non-ok
 status to a non-blocking LOG(WARNING)

PiperOrigin-RevId: 564420754
---
 third_party/xla/xla/service/gpu/autotuner_util.cc   | 10 +++-------
 third_party/xla/xla/service/gpu/autotuner_util.h    |  4 ++--
 third_party/xla/xla/service/gpu/triton_autotuner.cc |  9 +++++++--
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/autotuner_util.cc b/third_party/xla/xla/service/gpu/autotuner_util.cc
index bbed2a63e37..1eb16148c91 100644
--- a/third_party/xla/xla/service/gpu/autotuner_util.cc
+++ b/third_party/xla/xla/service/gpu/autotuner_util.cc
@@ -134,15 +134,11 @@ static AutotuneResult* TryFindInCache(const AutotuneCacheKey& key) {
   return TryFindInCache(key) != nullptr;
 }
 
-/*static*/ Status AutotunerUtil::AddResult(const AutotuneCacheKey& key,
-                                           AutotuneResult result) {
+/*static*/ bool AutotunerUtil::AddResult(const AutotuneCacheKey& key,
+                                         AutotuneResult result) {
   absl::MutexLock lock(&autotune_cache_mu);
   auto [_, inserted] = autotune_cache.emplace(key, std::move(result));
-  if (!inserted) {
-    return absl::AlreadyExistsError(absl::StrCat(
-        "The key is already in the autotune cache: ", key.ToString()));
-  }
-  return OkStatus();
+  return inserted;
 }
 
 /*static*/ StatusOr<AutotuneResult> AutotunerUtil::Autotune(
diff --git a/third_party/xla/xla/service/gpu/autotuner_util.h b/third_party/xla/xla/service/gpu/autotuner_util.h
index 9ac0e2540f2..0633fbd2a2a 100644
--- a/third_party/xla/xla/service/gpu/autotuner_util.h
+++ b/third_party/xla/xla/service/gpu/autotuner_util.h
@@ -178,10 +178,10 @@ struct AutotunerUtil {
 
   // Adds the result to the autotune cache.
   //
-  // It is an error to call this, if the key is already present in the cache.
+  // Returns true if the entry is inserted.
   //
   // Normally, we don't have to use this low level method.
-  static Status AddResult(const AutotuneCacheKey& key, AutotuneResult result);
+  static bool AddResult(const AutotuneCacheKey& key, AutotuneResult result);
 
   // Creates a RedzoneAllocator from a given config. If `force_stream` is
   // provided, than it is used for checking redzones.
diff --git a/third_party/xla/xla/service/gpu/triton_autotuner.cc b/third_party/xla/xla/service/gpu/triton_autotuner.cc
index 8b0ed2d42da..c78c4c7fbb3 100644
--- a/third_party/xla/xla/service/gpu/triton_autotuner.cc
+++ b/third_party/xla/xla/service/gpu/triton_autotuner.cc
@@ -712,8 +712,13 @@ Status Autotune(const AutotuneConfig& config, AutotunerCompileUtil& util,
           DumpAutotunedFusions(config, result, fusion, fusion_id_for_dump));
     }
 
-    TF_RETURN_IF_ERROR(AutotunerUtil::AddResult(
-        AutotunerUtil::GetKey(fusion, config), std::move(result)));
+    const AutotuneCacheKey key = AutotunerUtil::GetKey(fusion, config);
+    if (!AutotunerUtil::AddResult(key, std::move(result))) {
+      // In the context of model server, concurrent autotuning is expected and
+      // insertion of identical autotuning keys is accepted.
+      LOG(WARNING) << "AutotunerUtil::AddResult already existed: "
+                   << key.ToString();
+    }
 
     fusion_id_for_dump += 1;
   }
-- 
GitLab