From 7630bade64b5f21b341bcf4a5b5e7acbc0c7fdf6 Mon Sep 17 00:00:00 2001
From: scxfjiang <xuefei.jiang@amd.com>
Date: Fri, 22 Nov 2024 11:43:05 +0000
Subject: [PATCH] use xla flags to enable hipblaslt instead of env vars

---
 tensorflow/compiler/xla/stream_executor/gpu/BUILD          | 1 +
 tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc | 6 ++----
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/xla/stream_executor/gpu/BUILD b/tensorflow/compiler/xla/stream_executor/gpu/BUILD
index 96f560cc1181be..5b07fe0723f9e0 100644
--- a/tensorflow/compiler/xla/stream_executor/gpu/BUILD
+++ b/tensorflow/compiler/xla/stream_executor/gpu/BUILD
@@ -78,6 +78,7 @@ cc_library(
         #"//tensorflow/core/platform:env",
         "//tensorflow/tsl/util:env_var",
         "@com_google_absl//absl/types:any",
+        "//tensorflow/compiler/xla:debug_options_flags",
     ]),
 )
 
diff --git a/tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc b/tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc
index c9fe6109ef47fd..a47a0b01a6e895 100644
--- a/tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc
+++ b/tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/stream_executor/stream_executor.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/tsl/util/env_var.h"
+#include "tensorflow/compiler/xla/debug_options_flags.h"
 
 namespace stream_executor {
 
@@ -33,10 +34,7 @@ using xla::PrimitiveType;
 
 bool GpuBlasLtEnabled() {
   static std::atomic_bool result{[] {
-    bool value = false;
-    tsl::ReadBoolFromEnvVar("TF_ENABLE_GPU_BLASLT",
-                     /*default_value=*/false, &value);
-    return value;
+    return xla::GetDebugOptionsFromFlags().xla_gpu_enable_cublaslt();
   }()};
   return result;
 }