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; }