From 84c850527c0382f1435be956f1f6b1d34dd7291f Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Fri, 31 Jan 2020 16:52:07 -0800
Subject: [PATCH] Restrict deprecation warning to XLA_CPU and XLA_GPU devices

PiperOrigin-RevId: 292642710
Change-Id: I2c12ba0093017de0231d6341794ac0e0909e56c7
---
 tensorflow/compiler/jit/BUILD         |  1 +
 tensorflow/compiler/jit/xla_device.cc | 26 ++++++++++++++++----------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index feaf4c0cd57..c283328403b 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -161,6 +161,7 @@ XLA_DEVICE_DEPS = [
     ":xla_tensor",
     "@com_google_absl//absl/base",
     "@com_google_absl//absl/memory",
+    "@com_google_absl//absl/strings",
     "@com_google_absl//absl/synchronization",
     "@com_google_absl//absl/types:optional",
     "//tensorflow/compiler/jit/ops:xla_ops",
diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc
index 6d81a3c1673..830aaf74186 100644
--- a/tensorflow/compiler/jit/xla_device.cc
+++ b/tensorflow/compiler/jit/xla_device.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/base/call_once.h"
 #include "absl/memory/memory.h"
+#include "absl/strings/match.h"
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/xla_compile_on_demand_op.h"
 #include "tensorflow/compiler/jit/xla_device_context.h"
@@ -388,27 +389,32 @@ Status XlaDevice::TryGetDeviceContext(DeviceContext** out_context) {
 }
 
 // Warn about XLA_CPU/XLA_GPU exactly once.
-static void ShowXlaDeviceDeprecationWarning() {
+static void ShowXlaDeviceDeprecationWarning(
+    absl::string_view compilation_device_name) {
   static absl::once_flag once;
-  absl::call_once(once, [] {
-    LOG(WARNING) << "XLA_GPU and XLA_CPU devices are deprecated and will be "
-                    "removed in subsequent releases. Instead, use either "
-                    "@tf.function(experimental_compile=True) for must-compile "
-                    "semantics, or run with TF_XLA_FLAGS=--tf_xla_auto_jit=2 "
-                    "for auto-clustering best-effort compilation.";
-  });
+  if (absl::StrContains(compilation_device_name, "CPU") ||
+      absl::StrContains(compilation_device_name, "GPU")) {
+    absl::call_once(once, [] {
+      LOG(WARNING)
+          << "XLA_GPU and XLA_CPU devices are deprecated and will be "
+             "removed in subsequent releases. Instead, use either "
+             "@tf.function(experimental_compile=True) for must-compile "
+             "semantics, or run with TF_XLA_FLAGS=--tf_xla_auto_jit=2 "
+             "for auto-clustering best-effort compilation.";
+    });
+  }
 }
 
 void XlaDevice::Compute(OpKernel* op_kernel, OpKernelContext* context) {
   VLOG(2) << "XlaDevice::Compute " << op_kernel->name() << ":"
           << op_kernel->type_string();
-  ShowXlaDeviceDeprecationWarning();
+  ShowXlaDeviceDeprecationWarning(jit_device_name_.type_string());
   op_kernel->Compute(context);
 }
 
 void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context,
                              AsyncOpKernel::DoneCallback done) {
-  ShowXlaDeviceDeprecationWarning();
+  ShowXlaDeviceDeprecationWarning(jit_device_name_.type_string());
   VLOG(2) << "XlaDevice::ComputeAsync " << op_kernel->name() << ":"
           << op_kernel->type_string();
   op_kernel->ComputeAsync(context, done);