Use mlir specific gpu plugin when --define=with_mlir_gpu_support=true is specified.

PiperOrigin-RevId: 290268457
Change-Id: I35d5230d3d8f979569bae9df2576fb8a1779662e

tensorflow/compiler/jit/BUILD

@@ -77,19 +77,6 @@ cc_library(
     alwayslink = 1,
 )
 
-cc_library(
-    name = "xla_mlir_gpu_jit",
-    visibility = ["//visibility:public"],
-    deps = if_cuda_or_rocm([
-        ":jit_compilation_passes",
-        "//tensorflow/compiler/jit/kernels:xla_ops",
-        "//tensorflow/compiler/tf2xla/kernels:xla_ops",
-        "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
-        "//tensorflow/compiler/xla/service:mlir_gpu_plugin",
-    ]),
-    alwayslink = 1,
-)
-
 cc_library(
     name = "xla_cpu_device",
     srcs = ["xla_cpu_device.cc"],

tensorflow/compiler/xla/service/BUILD

@@ -931,8 +931,31 @@ cc_library(
     ],
 )
 
+# This flag enables experimental MLIR GPU support.
+config_setting(
+    name = "with_mlir_gpu_support",
+    values = {"define": "with_mlir_gpu_support=true"},
+    visibility = ["//visibility:public"],
+)
+
+# Lets us choose the right GPU plugin depending on whether the experimental MLIR
+# GPU plugin should be used or not.
 cc_library(
     name = "gpu_plugin",
+    deps = select(
+        {
+            ":with_mlir_gpu_support": [
+                ":mlir_gpu_plugin",
+            ],
+            "//conditions:default": [
+                ":gpu_plugin_no_mlir",
+            ],
+        },
+    ),
+)
+
+cc_library(
+    name = "gpu_plugin_no_mlir",
     deps = [
         ":service",
         "//tensorflow/compiler/xla/service/gpu:gpu_compiler",