Add target environment constraints to a subset of TensorFlow packages and targets.

PiperOrigin-RevId: 332884872
Change-Id: I65691fa2021c065e6c2ab57815d5a2b342d30ee2

tensorflow/c/BUILD

@@ -1,6 +1,7 @@
 # Description:
 # C API for TensorFlow, for use by client language bindings.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",
@@ -9,6 +10,11 @@ load(
     "tf_custom_op_library",
     "tf_kernel_library",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 
 package(

tensorflow/c/eager/BUILD

@@ -1,12 +1,20 @@
 # Experimental extensions to the C API for eager execution of kernels.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
+    "if_tpu",
     "tf_cc_test",
     "tf_copts",
     "tf_cuda_cc_test",
     "tf_cuda_library",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "internal_tfrt_deps")
 load(
     "//tensorflow/core/platform:build_config.bzl",
     "tf_kernel_tests_linkstatic",
@@ -87,7 +95,7 @@ tf_cuda_library(
         "//tensorflow/core/distributed_runtime:server_lib",
         "//tensorflow/core/distributed_runtime:worker_env",
         "//tensorflow/core:gpu_runtime",
-    ],
+    ] + internal_tfrt_deps(),
     alwayslink = 1,
 )
 
@@ -248,22 +256,24 @@ cc_library(
         ":c_api_unified_internal",
         ":gradients_internal",
         ":tape",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
         "//tensorflow/c:c_api",
         "//tensorflow/c:tf_status_helper",
         "//tensorflow/c/experimental/ops:array_ops",
         "//tensorflow/c/experimental/ops:math_ops",
         "//tensorflow/c/experimental/ops:nn_ops",
         "//tensorflow/cc/profiler",
-        "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/lib/llvm_rtti",
-        "@com_google_absl//absl/container:flat_hash_map",
+    ] + if_tpu(
-        "@com_google_absl//absl/container:flat_hash_set",
+        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
-        "@com_google_absl//absl/strings",
+        if_true = [],
-        "@com_google_absl//absl/types:span",
+    ),
-    ],
 )
 
 cc_library(
@@ -311,6 +321,8 @@ cc_library(
         ":c_api_unified_internal",
         ":gradients_internal",
         ":gradients_util",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
         "//tensorflow/c:c_api",
         "//tensorflow/c:tf_status_helper",
         "//tensorflow/c/experimental/gradients:math_grad",
@@ -319,13 +331,13 @@ cc_library(
         "//tensorflow/c/experimental/ops:math_ops",
         "//tensorflow/c/experimental/ops:nn_ops",
         "//tensorflow/cc/profiler",
-        "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/lib/llvm_rtti",
-        "@com_google_absl//absl/strings",
+    ] + if_tpu(
-        "@com_google_absl//absl/types:span",
+        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
-    ],
+        if_true = [],
+    ),
 )
 
 tf_cuda_cc_test(

tensorflow/c/eager/c_api.cc

@@ -39,7 +39,7 @@ limitations under the License.
 #include "tensorflow/c/eager/tfe_op_internal.h"
 #include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
 #include "tensorflow/c/tf_tensor_internal.h"
-#ifdef PLATFORM_GOOGLE
+#if defined(PLATFORM_GOOGLE) && !defined(LIBTFTPU)
 #include "tensorflow/core/tfrt/eager/c_api_tfrt.h"
 #endif
 #include "tensorflow/core/common_runtime/device.h"
@@ -729,7 +729,7 @@ void TFE_DeleteContextOptions(TFE_ContextOptions* options) { delete options; }
 
 TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) {
   if (opts->use_tfrt) {
-#ifdef PLATFORM_GOOGLE
+#if defined(PLATFORM_GOOGLE) && !defined(LIBTFTPU)
     return tensorflow::wrap(new tfrt::tf::ContextInterface(opts->async));
 #else
     status->status = tensorflow::errors::Unimplemented("TFRT is not supported");

tensorflow/c/eager/parallel_device/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",

tensorflow/c/experimental/filesystem/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental filesystem C APIs for TensorFlow.
 # Will be moved in proper place once all filesystems are converted to the
 # modular framework.

tensorflow/c/experimental/filesystem/plugins/gcs/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental gcs filesystem plugin.
 load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object", "tf_cc_test")
 

tensorflow/c/experimental/filesystem/plugins/hadoop/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental hadoop filesystem plugin.
 load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object", "tf_cc_test")
 
@@ -20,6 +22,7 @@ cc_library(
     name = "hadoop_filesystem_impl",
     srcs = ["hadoop_filesystem.cc"],
     hdrs = ["hadoop_filesystem.h"],
+    compatible_with = [],
     copts = select({
         "//conditions:default": [],
         "//tensorflow:windows": get_win_copts(),

tensorflow/c/experimental/filesystem/plugins/posix/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental posix filesystem plugin.
 load("//tensorflow:tensorflow.bzl", "tf_cc_shared_object")
 

tensorflow/c/experimental/filesystem/plugins/windows/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental windows filesystem plugin.
 load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object")
 

tensorflow/c/experimental/gradients/BUILD

@@ -1,3 +1,6 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Library of gradient functions.
 package(
     licenses = ["notice"],  # Apache 2.0
@@ -76,6 +79,7 @@ cc_library(
         ":array_grad",
         ":math_grad",
         ":nn_grad",
+        "//tensorflow/c/eager:gradients_internal",
     ],
 )
 

tensorflow/c/experimental/ops/BUILD

@@ -1,3 +1,6 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental ops. These will eventually be replaced by machine-generated versions.
 package(
     licenses = ["notice"],  # Apache 2.0

tensorflow/c/experimental/saved_model/core/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Experimental SavedModel C APIs for TensorFlow. See RFC
 # https://github.com/tensorflow/community/pull/207
 # Targets in this directory are pure C++ "Classes" underlying the C API types

tensorflow/c/experimental/saved_model/core/ops/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # This package contains written convenience helpers for Eager Operations
 # used by SavedModel. Once we autogenerate C++ Eager Op wrappers, we can remove these.
 load(

tensorflow/c/experimental/saved_model/core/revived_types/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # This package contains classes corresponding to Revived SavedObjectGraph types
 # used by SavedModel. See https://cs.opensource.google/tensorflow/tensorflow/+/c575e2ba93c442121d98d3f125d83fed1339924d:tensorflow/core/protobuf/saved_object_graph.proto;l=56-62
 package(

tensorflow/c/experimental/saved_model/internal/BUILD

@@ -9,6 +9,8 @@
 # Note(bmzhao): The *.cc files in this directory form the direct implementation of the
 # C API functions exposed in tf/c/experimental/saved_model/public/.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Note(bmzhao): All *type.h files in this directory are the internal definitions of
 # the opaque C types. These headers should only be visible to internal tensorflow
 # implementors.

tensorflow/c/experimental/saved_model/public/BUILD

@@ -8,6 +8,8 @@
 # programmatic checks that all "public" headers only include other "public"
 # headers.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     # This is intentionally public
     default_visibility = [

tensorflow/c/experimental/stream_executor/BUILD

@@ -1,6 +1,7 @@
 # Description:
 # StreamExecutor C API.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",

tensorflow/c/kernels/BUILD

@@ -1,9 +1,13 @@
-load(
+# buildifier: disable=same-origin-load
-    "//tensorflow:tensorflow.bzl",
+load("//tensorflow:tensorflow.bzl", "filegroup")
-    "tf_cc_test",
+
-    "tf_gen_op_libs",
+# buildifier: disable=same-origin-load
-    "tf_kernel_library",
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs")
-)
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 
 package(
     default_visibility = ["//visibility:public"],

tensorflow/cc/BUILD

@@ -2,16 +2,22 @@
 # TensorFlow is a computational framework, primarily for use in machine
 # learning applications.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "cc_library_with_android_deps",
     "tf_cc_binary",
     "tf_cc_test",
     "tf_copts",
-    "tf_gen_op_wrappers_cc",
     "transitive_hdrs",
 )
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrappers_cc")
+
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],  # Apache 2.0

tensorflow/cc/experimental/base/public/BUILD

@@ -8,6 +8,8 @@
 # 2. Are std:: types
 # 3. Wrap an opaque C type
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     # This is intentionally public
     default_visibility = [

tensorflow/cc/experimental/base/tests/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 # Tests for the C++ header-only base types.
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 

tensorflow/cc/profiler/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 
 package(

tensorflow/cc/saved_model/BUILD

@@ -1,6 +1,8 @@
 # Description:
 # TensorFlow SavedModel.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "if_android",

tensorflow/cc/saved_model/experimental/public/BUILD

@@ -1,6 +1,8 @@
 # Experimental C++ SavedModel Header Only APIs. See RFC
 # https://github.com/tensorflow/community/pull/207
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     # This is intentionally public
     default_visibility = [

tensorflow/cc/tools/BUILD

@@ -1,6 +1,7 @@
 # Description:
 # TensorFlow cc tools.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",

tensorflow/compiler/aot/BUILD

@@ -1,4 +1,5 @@
 load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
+load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
 load("//tensorflow/core/platform:build_config.bzl", "if_llvm_aarch64_available", "if_llvm_system_z_available")
 
@@ -303,6 +304,7 @@ cc_library(
     name = "aot_only_var_handle_op",
     srcs = ["aot_only_var_handle_op.cc"],
     hdrs = ["aot_only_var_handle_op.h"],
+    compatible_with = get_compatible_with_cloud(),
     visibility = [
         "//tensorflow/compiler/tf2xla:__pkg__",
     ],

tensorflow/compiler/jit/BUILD

@@ -1,5 +1,16 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "tf_cc_test")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "if_tpu", "tf_copts")
 load("//tensorflow/stream_executor:build_defs.bzl", "if_cuda_or_rocm")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library", "tf_jit_compilation_passes_extra_deps")
 load("//tensorflow/core/platform:build_config.bzl", "tf_additional_all_protos", "tf_proto_library")
 load("//tensorflow/core/platform:build_config_root.bzl", "tf_cuda_tests_tags")
@@ -66,8 +77,10 @@ cc_library(
         "//tensorflow/compiler/jit/kernels:xla_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
-        "//tensorflow/compiler/xla/service:cpu_plugin",
+    ] + if_tpu(
-    ],
+        if_false = ["//tensorflow/compiler/xla/service:cpu_plugin"],
+        if_true = [],
+    ),
     alwayslink = 1,
 )
 
@@ -94,15 +107,19 @@ cc_library(
         ":jit_compilation_passes",
         ":xla_device",
         ":xla_kernel_creator",  # buildcleaner: keep
+        "@com_google_absl//absl/memory",
         "//tensorflow/compiler/jit/kernels:xla_ops",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_op_registry",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
-        "//tensorflow/compiler/xla/service:cpu_plugin",  # buildcleaner: keep
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:lib",
-        "@com_google_absl//absl/memory",
+    ] + if_tpu(
-    ],
+        if_false = [
+            "//tensorflow/compiler/xla/service:cpu_plugin",  # buildcleaner: keep
+        ],
+        if_true = [],
+    ),
     alwayslink = 1,
 )
 
@@ -115,17 +132,21 @@ cc_library(
         ":jit_compilation_passes",
         ":xla_device",
         ":xla_kernel_creator",  # buildcleaner: keep
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
         "//tensorflow/compiler/jit/kernels:xla_ops",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_op_registry",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
-        "//tensorflow/compiler/xla/service:gpu_plugin",  # buildcleaner: keep
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core/common_runtime/gpu:gpu_init",
-        "@com_google_absl//absl/memory",
+    ] + if_tpu(
-        "@com_google_absl//absl/strings",
+        if_false = [
-    ],
+            "//tensorflow/compiler/xla/service:gpu_plugin",  # buildcleaner: keep
+        ],
+        if_true = [],
+    ),
     alwayslink = 1,
 )
 
@@ -329,12 +350,17 @@ cc_library(
     name = "xla_compilation_cache",
     srcs = ["xla_compilation_cache.cc"],
     hdrs = ["xla_compilation_cache.h"],
+    copts = tf_copts(),
     deps = [
         ":flags",
         ":xla_activity_listener",
         ":xla_activity_proto_cc",
-        "//tensorflow/compiler/mlir:array_container_utils",
+        "@com_google_absl//absl/base",
-        "//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
         "//tensorflow/compiler/tf2xla:common",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_context",
@@ -349,13 +375,13 @@ cc_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/platform:logging",
-        "@com_google_absl//absl/base",
+    ] + if_tpu(
-        "@com_google_absl//absl/container:flat_hash_map",
+        if_false = [
-        "@com_google_absl//absl/container:inlined_vector",
+            "//tensorflow/compiler/mlir:array_container_utils",
-        "@com_google_absl//absl/strings",
+            "//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
-        "@com_google_absl//absl/types:optional",
+        ],
-        "@com_google_absl//absl/types:span",
+        if_true = [],
-    ],
+    ),
 )
 
 tf_cc_test(

tensorflow/compiler/jit/graphcycles/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/compiler/jit/kernels/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     default_visibility = [
         "//tensorflow/compiler/tf2xla:internal",

tensorflow/compiler/jit/ops/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
 
 package(

tensorflow/compiler/jit/tests/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 licenses(["notice"])  # Apache 2.0

tensorflow/compiler/jit/xla_compilation_cache.cc

@@ -23,8 +23,6 @@ limitations under the License.
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/xla_activity.pb.h"
 #include "tensorflow/compiler/jit/xla_activity_listener.h"
-#include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h"
-#include "tensorflow/compiler/mlir/utils/array_container_utils.h"
 #include "tensorflow/compiler/tf2xla/shape_util.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
@@ -49,6 +47,11 @@ limitations under the License.
 #include "tensorflow/core/public/version.h"
 #include "tensorflow/core/util/dump_graph.h"
 
+#if !defined(LIBTFTPU)
+#include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h"
+#include "tensorflow/compiler/mlir/utils/array_container_utils.h"
+#endif
+
 namespace tensorflow {
 
 constexpr int64 XlaCompilationCache::kDefaultCompilationThreshold;
@@ -286,6 +289,13 @@ Status XlaCompilationCache::CompileSingleOp(
         });
     const ConfigProto* config = ctx->function_library()->config_proto();
     bool use_mlir = config && config->experimental().enable_mlir_bridge();
+#ifdef LIBTFTPU
+    if (use_mlir && has_tensor_list_arg) {
+      LOG(WARNING) << "MLIR is not supported in this environment.";
+    }
+    return compiler->CompileGraph(compile_options, node_def.name(),
+                                  std::move(graph), args, result);
+#else
     // TODO(b/155596779): Support TensorList args.
     if (!use_mlir || !has_tensor_list_arg) {
       return compiler->CompileGraph(compile_options, node_def.name(),
@@ -297,6 +307,7 @@ Status XlaCompilationCache::CompileSingleOp(
         *graph, mlir::SpanToArrayRef<XlaCompiler::Argument>(args),
         options.device_type.type_string(), compile_options.use_tuple_arg,
         *options.flib_def, debug_info, options.shape_representation_fn, result);
+#endif
   };
   return CompileImpl(options, name, args, compile_op,
                      /*compile_threshold=*/absl::nullopt,

tensorflow/compiler/mlir/tensorflow/c/BUILD

@@ -19,6 +19,7 @@ tf_cuda_library(
     srcs = [
         "c_api_unified_experimental_mlir.cc",
     ],
+    compatible_with = [],
     copts = tf_copts(),
     deps = [
         "//tensorflow/c:c_api",

tensorflow/compiler/plugin/BUILD

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 """Configuration file for an XLA plugin.
 
   please don't check in changes to this file. to prevent changes appearing

tensorflow/compiler/tf2tensorrt/BUILD

@@ -234,6 +234,7 @@ tf_cuda_library(
     name = "trt_engine_utils",
     srcs = ["utils/trt_engine_utils.cc"],
     hdrs = ["utils/trt_engine_utils.h"],
+    compatible_with = [],
     deps = [
         ":trt_logging",
         ":utils",
@@ -248,6 +249,7 @@ tf_cuda_library(
     name = "trt_logging",
     srcs = ["utils/trt_logger.cc"],
     hdrs = ["utils/trt_logger.h"],
+    compatible_with = [],
     visibility = ["//visibility:public"],
     deps = [
         ":common_utils",
@@ -289,6 +291,7 @@ tf_cuda_library(
         "utils/trt_lru_cache.h",
         "utils/trt_shape_optimization_profiles.h",
     ],
+    compatible_with = [],
     deps = [
         ":trt_allocator",
         ":trt_logging",
@@ -306,6 +309,7 @@ tf_cuda_library(
     name = "trt_allocator",
     srcs = ["utils/trt_allocator.cc"],
     hdrs = ["utils/trt_allocator.h"],
+    compatible_with = [],
     deps = [
         "//tensorflow/core:framework_headers_lib",
         "//tensorflow/core:framework_lite",
@@ -365,6 +369,7 @@ tf_cuda_library(
     hdrs = [
         "convert/logger_registry.h",
     ],
+    compatible_with = [],
     copts = tf_copts(),
     deps = [
         "@com_google_absl//absl/strings",
@@ -385,6 +390,7 @@ tf_cuda_library(
         "convert/convert_nodes.h",
         "convert/trt_optimization_pass.h",
     ],
+    compatible_with = [],
     deps = [
         ":common_utils",
         ":logger_registry",
@@ -543,6 +549,7 @@ tf_cuda_library(
     name = "trt_plugins",
     srcs = ["plugin/trt_plugin.cc"],
     hdrs = ["plugin/trt_plugin.h"],
+    compatible_with = [],
     deps = [
         "//tensorflow/core:framework_lite",
         "//tensorflow/core:lib_proto_parsing",

tensorflow/compiler/tf2xla/BUILD

@@ -1,4 +1,5 @@
-load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test", "tf_cuda_cc_test", "tf_openmp_copts")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "if_tpu", "tf_cc_binary", "tf_cc_test", "tf_copts", "tf_cuda_cc_test", "tf_openmp_copts")
 load(
     "//tensorflow/core/platform/default:cuda_build_defs.bzl",
     "if_cuda_is_configured",
@@ -10,6 +11,11 @@ load(
     "tf_proto_library_cc",
 )
 load("//tensorflow/compiler/xla:xla.bzl", "xla_py_proto_library")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_portable_proto_library")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured")
 load("//tensorflow/compiler/xla/service/cpu:build_defs.bzl", "runtime_copts")
@@ -165,6 +171,7 @@ cc_library(
     name = "mlir_tf2xla",
     srcs = ["mlir_tf2xla.cc"],
     hdrs = ["tf2xla.h"],
+    compatible_with = [],
     visibility = ["//visibility:public"],
     deps = [
         ":tf2xla_proto_cc",
@@ -258,6 +265,7 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "//third_party/eigen3",
         "//tensorflow/core/framework:numeric_types",
+        "//tensorflow/core/platform:bfloat16",
     ] + tf_additional_tensor_coding_deps(),
     alwayslink = 1,
 )
@@ -303,14 +311,18 @@ cc_library(
         "//tensorflow/compiler/xla/client:client_library",
         "//tensorflow/compiler/xla/client:local_client",
         "//tensorflow/compiler/xla/client:xla_computation",
-        "//tensorflow/compiler/xla/service:cpu_plugin",
         "//tensorflow/compiler/xla/service:platform_util",
-        "//tensorflow/compiler/xla/service/cpu:buffer_info_util",
-        "//tensorflow/compiler/xla/service/cpu:cpu_executable",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/stream_executor:platform",
-    ],
+    ] + if_tpu(
+        if_false = [
+            "//tensorflow/compiler/xla/service:cpu_plugin",
+            "//tensorflow/compiler/xla/service/cpu:buffer_info_util",
+            "//tensorflow/compiler/xla/service/cpu:cpu_executable",
+        ],
+        if_true = [],
+    ),
 )
 
 cc_library(
@@ -334,6 +346,7 @@ cc_library(
         "xla_op_kernel.h",
         "xla_op_registry.h",
     ],
+    copts = tf_copts(),
     visibility = [":friends"],
     deps = [
         ":common",
@@ -349,11 +362,13 @@ cc_library(
         ":xla_helpers",
         ":xla_op_registry",
         ":xla_resource",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/types:variant",
         "//tensorflow/compiler/jit:common",
         "//tensorflow/compiler/jit:flags",
         "//tensorflow/compiler/jit:shape_inference",
-        "//tensorflow/compiler/mlir:array_container_utils",
-        "//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
         "//tensorflow/compiler/xla:protobuf_util",
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:status_macros",
@@ -371,11 +386,13 @@ cc_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:ops",
         "//tensorflow/core:protos_all_cc",
-        "@com_google_absl//absl/algorithm:container",
+    ] + if_tpu(
-        "@com_google_absl//absl/memory",
+        if_false = [
-        "@com_google_absl//absl/types:span",
+            "//tensorflow/compiler/mlir:array_container_utils",
-        "@com_google_absl//absl/types:variant",
+            "//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
-    ],
+        ],
+        if_true = [],
+    ),
     alwayslink = 1,
 )
 
@@ -864,6 +881,7 @@ cc_library(
     name = "mlir_bridge_pass",
     srcs = ["mlir_bridge_pass.cc"],
     hdrs = ["mlir_bridge_pass.h"],
+    compatible_with = [],
     deps = [
         "//tensorflow/compiler/jit:flags",
         "//tensorflow/compiler/mlir:mlir_graph_optimization_pass",
@@ -877,14 +895,20 @@ cc_library(
 
 cc_library(
     name = "mlir_bridge_pass_registration",
-    srcs = [
+    srcs = if_tpu(
-        "mlir_bridge_pass_registration.cc",
+        if_false = [
-    ],
+            "mlir_bridge_pass_registration.cc",
-    deps = [
+        ],
-        ":mlir_bridge_pass",
+        if_true = [],
-        "//tensorflow/compiler/mlir:mlir_graph_optimization_pass_registration",
+    ),
-        "//tensorflow/core:core_cpu",
+    deps = if_tpu(
-    ],
+        if_false = [
+            ":mlir_bridge_pass",
+            "//tensorflow/compiler/mlir:mlir_graph_optimization_pass_registration",
+            "//tensorflow/core:core_cpu",
+        ],
+        if_true = [],
+    ),
     alwayslink = 1,
 )
 

tensorflow/compiler/tf2xla/cc/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_cc")
 
 package(

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
 
 package(

tensorflow/compiler/tf2xla/lib/BUILD

@@ -1,5 +1,8 @@
 # Utilities for building XLA computations.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     default_visibility = ["//tensorflow/compiler/tf2xla:friends"],
     licenses = ["notice"],  # Apache 2.0

tensorflow/compiler/tf2xla/ops/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_custom_op_library",

tensorflow/compiler/tf2xla/xla_compiler.cc

@@ -23,8 +23,6 @@ limitations under the License.
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/shape_inference.h"
-#include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h"
-#include "tensorflow/compiler/mlir/utils/array_container_utils.h"
 #include "tensorflow/compiler/tf2xla/graph_compiler.h"
 #include "tensorflow/compiler/tf2xla/rearrange_function_argument.h"
 #include "tensorflow/compiler/tf2xla/shape_util.h"
@@ -58,6 +56,11 @@ limitations under the License.
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/util/dump_graph.h"
 
+#ifndef LIBTFTPU
+#include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h"
+#include "tensorflow/compiler/mlir/utils/array_container_utils.h"
+#endif
+
 namespace tensorflow {
 namespace {
 
@@ -730,6 +733,13 @@ Status XlaCompiler::CompileFunction(
   }
 
   VLOG(1) << "====================================================";
+#ifdef LIBTFTPU
+  if (GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge) {
+    VLOG(1) << "MLIR is not supported in this environment.";
+  }
+  TF_RETURN_IF_ERROR(
+      CompileGraph(options, function_id, std::move(graph), args, result));
+#else
   if (GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge) {
     VLOG(1) << "Using MLIR bridge";
     GraphDebugInfo debug_info;
@@ -742,6 +752,7 @@ Status XlaCompiler::CompileFunction(
     TF_RETURN_IF_ERROR(
         CompileGraph(options, function_id, std::move(graph), args, result));
   }
+#endif
   VLOG(1) << "====================================================";
 
   cache_[{function_id, arg_vector}] = *result;

tensorflow/compiler/xla/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "tf_cc_test")
 load(
     "//tensorflow/core/platform:build_config.bzl",

tensorflow/compiler/xla/client/BUILD

@@ -1,6 +1,8 @@
 # Description:
 #   XLA client libraries.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/compiler/xla/client/lib/BUILD

@@ -1,5 +1,7 @@
 # Common computation builders for XLA.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites", "xla_test")
 
 package(

tensorflow/compiler/xla/pjrt/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 

tensorflow/compiler/xla/pjrt/distributed/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library_cc")
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")

tensorflow/compiler/xla/python/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow/core/platform:build_config.bzl", "pyx_library")
 load("//tensorflow/compiler/xla:xla.bzl", "xla_py_test_deps")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")

tensorflow/compiler/xla/python/tpu_driver/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library_cc")
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load(
@@ -77,6 +78,7 @@ cc_library(
 cc_library(
     name = "direct_tpu_driver",
     srcs = ["direct_tpu_driver.cc"],
+    compatible_with = [],
     deps = [
         ":tpu_driver",
         "@com_google_absl//absl/strings:str_format",

tensorflow/compiler/xla/python/tpu_driver/client/BUILD

@@ -1,3 +1,9 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "pybind_extension")
 
 package(
@@ -11,6 +17,7 @@ cc_library(
     hdrs = [
         "tpu_client.h",
     ],
+    compatible_with = [],
     deps = [
         "//tensorflow/compiler/xla:literal",
         "//tensorflow/compiler/xla:shape_util",

tensorflow/compiler/xla/rpc/BUILD

@@ -1,6 +1,8 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load(
     "//tensorflow:tensorflow.bzl",
+    "if_tpu",
     "tf_cc_binary",
     "tf_cc_test",
 )
@@ -50,13 +52,15 @@ cc_library(
     srcs = ["grpc_service_main.cc"],
     deps = [
         ":grpc_service",
-        "//tensorflow/compiler/xla/service:cpu_plugin",
+        "@com_google_absl//absl/strings:str_format",
         "//tensorflow/compiler/xla/service:platform_util",
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
-        "@com_google_absl//absl/strings:str_format",
         tf_grpc_cc_dependency(),
-    ],
+    ] + if_tpu(
+        if_false = ["//tensorflow/compiler/xla/service:cpu_plugin"],
+        if_true = [],
+    ),
 )
 
 tf_cc_binary(

tensorflow/compiler/xla/service/BUILD

@@ -1,6 +1,7 @@
 # Description:
 #   XLA service implementation.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test")
 load(
     "//tensorflow/core/platform:build_config.bzl",
@@ -8,6 +9,12 @@ load(
     "tf_proto_library_py",
 )
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "internal_hlo_deps")
 load(
     "//tensorflow/core/platform/default:cuda_build_defs.bzl",
     "if_cuda_is_configured",
@@ -1018,6 +1025,7 @@ cc_library(
 
 cc_library(
     name = "cpu_plugin",
+    compatible_with = [],
     deps = [
         ":service",
         "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
@@ -1037,6 +1045,7 @@ config_setting(
 # GPU plugin should be used or not.
 cc_library(
     name = "gpu_plugin",
+    compatible_with = [],
     deps = select(
         {
             ":with_mlir_gpu_support": [
@@ -1051,6 +1060,7 @@ cc_library(
 
 cc_library(
     name = "gpu_plugin_no_mlir",
+    compatible_with = [],
     deps = [
         ":service",
         "//tensorflow/compiler/xla/service/gpu:gpu_compiler",
@@ -1067,6 +1077,7 @@ cc_library(
 
 cc_library(
     name = "gpu_plugin_mlir",
+    compatible_with = [],
     deps = [
         ":service",
         "//tensorflow/compiler/xla/service/gpu:gpu_transfer_manager",
@@ -1147,6 +1158,10 @@ cc_library(
         ":maybe_owning_device_memory",
         ":shaped_buffer",
         ":stream_pool",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/types:variant",
         "//tensorflow/compiler/xla:debug_options_flags",
         "//tensorflow/compiler/xla:executable_run_options",
         "//tensorflow/compiler/xla:shape_tree",
@@ -1161,11 +1176,7 @@ cc_library(
         "//tensorflow/stream_executor",
         "//tensorflow/stream_executor:device_description",
         "//tensorflow/stream_executor:device_memory_allocator",
-        "@com_google_absl//absl/memory",
+    ] + internal_hlo_deps(),
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/types:span",
-        "@com_google_absl//absl/types:variant",
-    ],
 )
 
 cc_library(
@@ -1194,6 +1205,7 @@ cc_library(
     name = "llvm_compiler",
     srcs = ["llvm_compiler.cc"],
     hdrs = ["llvm_compiler.h"],
+    compatible_with = [],
     deps = [
         ":compiler",
         "//tensorflow/core:lib_internal",
@@ -4042,6 +4054,7 @@ cc_library(
     name = "elemental_ir_emitter",
     srcs = ["elemental_ir_emitter.cc"],
     hdrs = ["elemental_ir_emitter.h"],
+    compatible_with = [],
     deps = [
         ":hlo",
         ":hlo_casting_utils",

tensorflow/compiler/xla/service/cpu/BUILD

@@ -6,6 +6,12 @@ load(
     "//third_party/mkl:build_defs.bzl",
     "mkl_deps",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test", "tf_openmp_copts")
 load(":build_defs.bzl", "runtime_copts")
 load("//tensorflow/core/platform:build_config.bzl", "if_llvm_system_z_available")
@@ -742,6 +748,7 @@ cc_library(
     name = "runtime_single_threaded_matmul",
     srcs = ["runtime_single_threaded_matmul.cc"],
     hdrs = ["runtime_single_threaded_matmul.h"],
+    compatible_with = get_compatible_with_cloud(),
     copts = runtime_copts(),
     visibility = ["//visibility:public"],
     deps = [

tensorflow/compiler/xla/service/gpu/BUILD

@@ -26,6 +26,11 @@ load(
     "//tensorflow/core/platform/default:cuda_build_defs.bzl",
     "if_cuda_is_configured",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "if_nccl")
 load("//third_party/mlir:tblgen.bzl", "gentbl")
 
@@ -61,6 +66,7 @@ cc_library(
     name = "gpu_executable_run_options",
     srcs = ["gpu_executable_run_options.cc"],
     hdrs = ["gpu_executable_run_options.h"],
+    compatible_with = get_compatible_with_cloud(),
     visibility = ["//visibility:public"],
     deps = [
         "//tensorflow/compiler/xla:statusor",
@@ -434,6 +440,7 @@ tf_cuda_library(
         ["dummy_all_reduce_thunk.cc"],
     ),
     hdrs = ["nccl_all_reduce_thunk.h"],
+    compatible_with = [],
     deps = [
         ":buffer_allocations",
         ":hlo_execution_profiler",
@@ -470,6 +477,7 @@ cc_library(
     hdrs = [
         "gpu_debug_info_manager.h",
     ],
+    compatible_with = get_compatible_with_cloud(),
     deps = [
         "//tensorflow/compiler/xla/service:buffer_assignment",
         "//tensorflow/compiler/xla/service:hlo",
@@ -1374,6 +1382,7 @@ cc_library(
 cc_library(
     name = "xfeed_queue",
     hdrs = ["xfeed_queue.h"],
+    compatible_with = get_compatible_with_cloud(),
     deps = [
         "//tensorflow/core:lib",
         "@com_google_absl//absl/base:core_headers",
@@ -1384,6 +1393,7 @@ cc_library(
     name = "infeed_manager",
     srcs = ["infeed_manager.cc"],
     hdrs = ["infeed_manager.h"],
+    compatible_with = get_compatible_with_cloud(),
     deps = [
         ":xfeed_queue",
         "//tensorflow/compiler/xla:shape_tree",
@@ -1398,6 +1408,7 @@ cc_library(
     name = "outfeed_manager",
     srcs = ["outfeed_manager.cc"],
     hdrs = ["outfeed_manager.h"],
+    compatible_with = get_compatible_with_cloud(),
     deps = [
         ":xfeed_queue",
         "//tensorflow/compiler/xla:literal",

tensorflow/compiler/xla/service/interpreter/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow/core/platform:build_config_root.bzl",
     "if_static",

tensorflow/compiler/xla/service/spmd/BUILD

@@ -1,5 +1,6 @@
 # Description: SPMD partitioning pass.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/compiler/xla/tests/BUILD

@@ -6,7 +6,14 @@ load(
     "//tensorflow/core/platform:build_config_root.bzl",
     "tf_cuda_tests_tags",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "genrule")
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 
 package(
     default_visibility = [":friends"],
@@ -177,6 +184,7 @@ cc_library(
 tf_cc_binary(
     name = "local_client_aot_test_helper",
     srcs = ["local_client_aot_test_helper.cc"],
+    compatible_with = [],
     deps = [
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:util",
@@ -239,6 +247,7 @@ cc_library(
     testonly = True,
     srcs = ["llvm_irgen_test_base.cc"],
     hdrs = ["llvm_irgen_test_base.h"],
+    compatible_with = [],
     deps = [
         ":codegen_test_base",
         ":filecheck",
@@ -266,6 +275,7 @@ cc_library(
     testonly = True,
     srcs = ["filecheck.cc"],
     hdrs = ["filecheck.h"],
+    compatible_with = [],
     data = [
         "@llvm-project//llvm:FileCheck",
     ],

tensorflow/compiler/xla/tools/BUILD

@@ -1,5 +1,7 @@
 # Tools and utilities that aid in XLA development and usage.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "if_cuda_or_rocm",

tensorflow/compiler/xrt/BUILD

@@ -1,11 +1,12 @@
 # Description: Operations defined for XRT
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_custom_op_py_library",
-    "tf_gen_op_libs",
     "tf_gen_op_wrapper_py",
 )
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs")
 load(
     "//tensorflow/core/platform:build_config.bzl",
     "tf_proto_library_cc",

tensorflow/compiler/xrt/cc/BUILD

@@ -1,7 +1,4 @@
-load(
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrappers_cc")
-    "//tensorflow:tensorflow.bzl",
-    "tf_gen_op_wrappers_cc",
-)
 
 package(
     default_visibility = ["//visibility:public"],

tensorflow/compiler/xrt/kernels/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     default_visibility = [
         "//learning/brain:__subpackages__",

tensorflow/compiler/xrt/tests/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_cc_test")
 load(
     "//tensorflow/core/platform:build_config_root.bzl",

tensorflow/core/BUILD

@@ -65,13 +65,13 @@
 
 load(
     "//tensorflow:tensorflow.bzl",
-    "cc_header_only_library",
     "if_android",
     "if_chromiumos",
     "if_cuda_or_rocm",
     "if_ios",
     "if_mobile",
     "if_not_windows",
+    "if_tpu",
     "tf_android_core_proto_headers",
     "tf_cc_test",
     "tf_cc_test_mkl",
@@ -80,13 +80,18 @@ load(
     "tf_cuda_library",
     "tf_defines_nortti_if_lite_protos",
     "tf_features_nomodules_if_mobile",
-    "tf_gen_op_libs",
     "tf_genrule_cmd_append_to_srcs",
     "tf_opts_nortti_if_lite_protos",
     "tf_portable_full_lite_protos",
     "transitive_hdrs",
 )
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs")
+
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "if_nccl")
 
@@ -102,6 +107,12 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_tests_gpu")
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_monitoring_framework_deps")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "genrule")
+
 # For platform specific build config
 load(
     "//tensorflow/core/platform:build_config.bzl",
@@ -882,7 +893,6 @@ cc_library(
         "//tensorflow/c/kernels:histogram_summary_op_lib",
         "//tensorflow/c/kernels:merge_summary_op_lib",
         "//tensorflow/c/kernels:summary_op_lib",
-        "//tensorflow/compiler/mlir/tensorflow:mlir_passthrough_op",
     ] + if_chromiumos(
         [],
         # Non-tpu platforms don't need tpu dependency. It would be best to guard
@@ -907,7 +917,10 @@ cc_library(
     ]) + if_tensorrt([
         "//tensorflow/compiler/tf2tensorrt:trt_engine_resource_ops_op_lib",
         "//tensorflow/compiler/tf2tensorrt:trt_op_libs",
-    ]),
+    ]) + if_tpu(
+        if_false = ["//tensorflow/compiler/mlir/tensorflow:mlir_passthrough_op"],
+        if_true = [],
+    ),
     alwayslink = 1,
 )
 

tensorflow/core/api_def/BUILD

@@ -6,6 +6,8 @@
 #   :python_api_def
 #   :java_api_def
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_binary",

tensorflow/core/api_def/java_api/BUILD

@@ -1,6 +1,8 @@
 # Description:
 #   Provides Java ApiDef access and ApiDef validation for TensorFlow.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
 package(
     default_visibility = ["//visibility:private"],
     licenses = ["notice"],  # Apache 2.0

tensorflow/core/common_runtime/BUILD

@@ -1,6 +1,5 @@
 load(
     "//tensorflow:tensorflow.bzl",
-    "cc_header_only_library",
     "if_tpu",
     "tf_cc_test",
     "tf_cc_test_mkl",
@@ -10,12 +9,18 @@ load(
     "tf_openmp_copts",
 )
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library")
+
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cc_test_gpu")
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
 # For platform specific build config
 load(
     "//tensorflow/core/platform:build_config.bzl",

tensorflow/core/common_runtime/eager/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",

tensorflow/core/common_runtime/gpu/BUILD

@@ -14,6 +14,9 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_tests_gpu")
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
 # For platform specific build config
 load(
     "//tensorflow/core/platform:build_config.bzl",

tensorflow/core/data/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load(
     "//tensorflow/core/platform:build_config.bzl",

tensorflow/core/data/service/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("@com_github_grpc_grpc//bazel:cc_grpc_library.bzl", "cc_grpc_library")
 load(
     "//tensorflow/core/platform:build_config.bzl",
@@ -5,10 +6,17 @@ load(
     "tf_proto_library",
     "tf_protos_profiler_service",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load(
     "//tensorflow:tensorflow.bzl",
-    "cc_header_only_library",
     "tf_cc_test",
 )
 
@@ -152,6 +160,7 @@ tf_cc_test(
 cc_grpc_library(
     name = "dispatcher_cc_grpc_proto",
     srcs = [":dispatcher_proto"],
+    compatible_with = get_compatible_with_cloud(),
     generate_mocks = True,
     grpc_only = True,
     deps = [":dispatcher_proto_cc"],
@@ -422,6 +431,7 @@ tf_cc_test(
 cc_grpc_library(
     name = "worker_cc_grpc_proto",
     srcs = [":worker_proto"],
+    compatible_with = get_compatible_with_cloud(),
     generate_mocks = True,
     grpc_only = True,
     deps = [":worker_proto_cc"],

tensorflow/core/debug/BUILD

@@ -12,6 +12,7 @@
 #    a watch state.
 # ":debug_node_key" - Defines a struct used for tracking tensors.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load(
     "//tensorflow:tensorflow.bzl",

tensorflow/core/distributed_runtime/BUILD

@@ -3,9 +3,11 @@
 # to be distributed and performed in parallel across multiple
 # processes.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_copts", "tf_cuda_library")
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")  # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")  # buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")  # buildifier: disable=same-origin-load
 
 # For platform specific build config
 load(

tensorflow/core/distributed_runtime/eager/BUILD

@@ -1,3 +1,9 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load(
     "//tensorflow:tensorflow.bzl",

tensorflow/core/distributed_runtime/rpc/BUILD

@@ -1,6 +1,7 @@
 # Description:
 #   RPC communication interfaces and implementations for TensorFlow.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "if_windows",
@@ -12,6 +13,7 @@ load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")  # buildifier: disable=sa
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_tests")  # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_grpc_dependency")  # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")  # buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")  # buildifier: disable=same-origin-load
 
 # For platform specific build config
 load(

tensorflow/core/distributed_runtime/rpc/eager/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 

tensorflow/core/example/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "@io_bazel_rules_closure//closure:defs.bzl",
     "closure_proto_library",

tensorflow/core/framework/BUILD

@@ -9,8 +9,16 @@ load(
     "tf_cc_tests",
     "tf_copts",
     "tf_cuda_library",
-    "tf_generate_proto_text_sources",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_generate_proto_text_sources")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_selective_registration_deps")
 load(
     "//tensorflow/core/platform:build_config_root.bzl",
     "if_static",
@@ -20,9 +28,6 @@ load(
     "cc_library",
 )
 
-# buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "tf_selective_registration_deps")
-
 package(
     default_visibility = [
         "//tensorflow/core:__subpackages__",

tensorflow/core/graph/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",

tensorflow/core/grappler/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_library")
 
 package(

tensorflow/core/grappler/clusters/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_library")
 load(

tensorflow/core/grappler/costs/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",

tensorflow/core/grappler/graph_analyzer/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/core/grappler/inputs/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/core/grappler/inputs/testdata/BUILD

@@ -1,6 +1,8 @@
 # Description:
 # grappler testdata packages.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
 package(
     licenses = ["notice"],  # Apache 2.0
 )

tensorflow/core/grappler/optimizers/BUILD

@@ -1,4 +1,11 @@
-load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cc_test_mkl", "tf_copts", "tf_cuda_cc_test", "tf_kernel_library")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cc_test_mkl", "tf_copts", "tf_cuda_cc_test")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
 
 # Platform specific build config
 load(

tensorflow/core/grappler/optimizers/data/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load("//tensorflow/core/platform:build_config.bzl", "tf_protos_all")
 

tensorflow/core/grappler/optimizers/data/vectorization/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load("//tensorflow/core/platform:build_config.bzl", "tf_protos_all")
 

tensorflow/core/grappler/utils/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load(
     "//tensorflow/core/platform:build_config.bzl",

tensorflow/core/grappler/verifiers/BUILD

@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/core/kernels/BUILD

@@ -1,7 +1,7 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load(
     "//tensorflow:tensorflow.bzl",
-    "cc_header_only_library",
     "if_android",
     "if_cuda_or_rocm",
     "if_mobile",
@@ -12,11 +12,16 @@ load(
     "tf_cc_tests",
     "tf_copts",
     "tf_cuda_library",
-    "tf_kernel_library",
     "tf_opts_nortti_if_lite_protos",
 )
 load("//tensorflow/core/kernels/mlir_generated:build_defs.bzl", "if_mlir_generated_gpu_kernels_enabled")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "if_nccl")
 
@@ -26,6 +31,9 @@ load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_tests")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
 # For platform specific build config
 load(
     "//tensorflow/core/platform:build_config.bzl",
@@ -501,6 +509,7 @@ tf_cuda_library(
     name = "gpu_utils",
     srcs = if_cuda_or_rocm(["gpu_utils.cc"]),
     hdrs = ["gpu_utils.h"],
+    compatible_with = [],
     deps = [
         ":gpu_util_hdrs",
         "//tensorflow/core:autotuning_proto_cc",
@@ -1476,6 +1485,7 @@ tf_kernel_library(
 tf_kernel_library(
     name = "cudnn_rnn_kernels",
     srcs = ["cudnn_rnn_ops.cc"],
+    compatible_with = [],
     visibility = ["//visibility:public"],
     deps = [
         ":bounds_check_lib",
@@ -3832,13 +3842,13 @@ tf_kernel_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core/util:image_resizer_state",
         "//tensorflow/core/util/proto:proto_utils",
-        "//tensorflow/stream_executor/gpu:gpu_asm_opts",
     ] + select({
         ":xsmm_convolutions": [
             "@libxsmm_archive//:xsmm_avx",
         ],
         "//conditions:default": [],
     }) + if_cuda([
+        "//tensorflow/stream_executor/gpu:gpu_asm_opts",
         "//tensorflow/core/platform/default/build_config:cublas_plugin",
         "//tensorflow/core/platform/default/build_config:cudnn_plugin",
         "//tensorflow/stream_executor:tf_allocator_adapter",

tensorflow/core/kernels/batching_util/BUILD

@@ -1,5 +1,6 @@
 # Description: Utilities.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/core/kernels/boosted_trees/BUILD

@@ -1,11 +1,12 @@
 # Description:
 #   OpKernels for boosted trees ops.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",
-    "tf_kernel_library",
 )
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
 load(
     "//tensorflow/core/platform:build_config.bzl",
     "tf_proto_library",

tensorflow/core/kernels/boosted_trees/quantiles/BUILD

@@ -1,6 +1,8 @@
 # Description:
 #   This directory contains common quantile utilities used in boosted_trees.
 
+load("//tensorflow:tensorflow.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 
 package(

tensorflow/core/kernels/data/BUILD

@@ -1,13 +1,19 @@
 # Description:
 #   OpKernels for tf.data
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "if_not_mobile",
     "tf_cc_test",
-    "tf_kernel_library",
 )
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],  # Apache 2.0

tensorflow/core/kernels/data/experimental/BUILD

@@ -1,11 +1,9 @@
 # Description:
 #   Contains experimental kernels for datasets and iterators.
 
-load(
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
-    "//tensorflow:tensorflow.bzl",
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
-    "tf_cc_test",
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
-    "tf_kernel_library",
-)
 
 package(
     default_visibility = ["//tensorflow:internal"],

tensorflow/core/kernels/data/experimental/sql/BUILD

@@ -1,6 +1,8 @@
 # Description:
 #   SQL library.
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],  # Apache 2.0

tensorflow/core/kernels/fuzzing/BUILD

@@ -1,5 +1,12 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow/core/kernels/fuzzing:tf_ops_fuzz_target_lib.bzl", "tf_ops_fuzz_target_lib")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow/core/kernels/fuzzing:tf_ops_fuzz_target_lib.bzl", "tf_oss_fuzz_corpus")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow/core/kernels/fuzzing:tf_ops_fuzz_target_lib.bzl", "tf_oss_fuzz_dict")
 
 package(

tensorflow/core/kernels/hexagon/BUILD

@@ -1,11 +1,12 @@
 # Description:
 #   quantization-specific OpKernels for hexagon
 
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_cc_test",
-    "tf_kernel_library",
 )
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
 
 package(
     default_visibility = ["//visibility:public"],

tensorflow/core/kernels/image/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load(
     "//tensorflow:tensorflow.bzl",
@@ -6,14 +7,21 @@ load(
     "tf_cc_test",
     "tf_cc_tests",
     "tf_copts",
-    "tf_kernel_library",
 )
 load(
     "//tensorflow/core/platform:build_config_root.bzl",
     "tf_cuda_tests_tags",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+
 # TODO(rmlarsen): Remove ASAP.
 package_group(
     name = "friends",

tensorflow/core/kernels/linalg/BUILD

@@ -1,15 +1,23 @@
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow:tensorflow.bzl",
     "if_cuda_or_rocm",
-    "tf_kernel_library",
 )
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load(
     "@local_config_rocm//rocm:build_defs.bzl",
     "if_rocm",
 )
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "filegroup")
+
+# buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+
 package(
     default_visibility = [
         "//tensorflow:__subpackages__",

tensorflow/core/kernels/neon/BUILD

@@ -1,7 +1,4 @@
-load(
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
-    "//tensorflow:tensorflow.bzl",
-    "tf_kernel_library",
-)
 
 # Description:
 # Kernel implementations using Neon intrinsics.