experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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

Browse Source 
PiperOrigin-RevId: 332884872
Change-Id: I65691fa2021c065e6c2ab57815d5a2b342d30ee2
This commit is contained in:
Wenhao Jia 2020-09-21 10:48:19 -07:00 committed by TensorFlower Gardener
parent 0451edab7d
commit 8f1362de18
171 changed files with 639 additions and 129 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/c/BUILD
View File

@ -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(

34
tensorflow/c/eager/BUILD
View File

@ -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",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
] + if_tpu(
if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
if_true = [],
),
)
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",
"@com_google_absl//absl/types:span",
],
] + if_tpu(
if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
if_true = [],
),
)
tf_cuda_cc_test(

4
tensorflow/c/eager/c_api.cc
View File

@ -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");

2
tensorflow/c/eager/parallel_device/BUILD
View File

@ -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",

2
tensorflow/c/experimental/filesystem/BUILD
View File

@ -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.

2
tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
View File

@ -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")

3
tensorflow/c/experimental/filesystem/plugins/hadoop/BUILD
View File

@ -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(),

2
tensorflow/c/experimental/filesystem/plugins/posix/BUILD
View File

@ -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")

2
tensorflow/c/experimental/filesystem/plugins/windows/BUILD
View File

@ -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")

4
tensorflow/c/experimental/gradients/BUILD
View File

@ -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",
],
)

3
tensorflow/c/experimental/ops/BUILD
View File

@ -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

2
tensorflow/c/experimental/saved_model/core/BUILD
View File

@ -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

2
tensorflow/c/experimental/saved_model/core/ops/BUILD
View File

@ -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(

2
tensorflow/c/experimental/saved_model/core/revived_types/BUILD
View File

@ -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(

2
tensorflow/c/experimental/saved_model/internal/BUILD
View File

@ -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.

2
tensorflow/c/experimental/saved_model/public/BUILD
View File

@ -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 = [

1
tensorflow/c/experimental/stream_executor/BUILD
View File

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

16
tensorflow/c/kernels/BUILD
View File

@ -1,9 +1,13 @@
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
"tf_gen_op_libs",
"tf_kernel_library",
)
# buildifier: disable=same-origin-load
load("//tensorflow:tensorflow.bzl", "filegroup")
# buildifier: disable=same-origin-load
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"],

8
tensorflow/cc/BUILD
View File

@ -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

2
tensorflow/cc/experimental/base/public/BUILD
View File

@ -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 = [

2
tensorflow/cc/experimental/base/tests/BUILD
View File

@ -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")

1
tensorflow/cc/profiler/BUILD
View File

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

2
tensorflow/cc/saved_model/BUILD
View File

@ -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",

2
tensorflow/cc/saved_model/experimental/public/BUILD
View File

@ -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 = [

1
tensorflow/cc/tools/BUILD
View File

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

2
tensorflow/compiler/aot/BUILD
View File

@ -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__",
],

62
tensorflow/compiler/jit/BUILD
View File

@ -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",
"@com_google_absl//absl/strings",
],
] + if_tpu(
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",
"//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
"@com_google_absl//absl/base",
"@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",
"@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",
],
] + if_tpu(
if_false = [
"//tensorflow/compiler/mlir:array_container_utils",
"//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
],
if_true = [],
),
)
tf_cc_test(

1
tensorflow/compiler/jit/graphcycles/BUILD
View File

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

2
tensorflow/compiler/jit/kernels/BUILD
View File

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

1
tensorflow/compiler/jit/ops/BUILD
View File

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

1
tensorflow/compiler/jit/tests/BUILD
View File

@ -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

15
tensorflow/compiler/jit/xla_compilation_cache.cc
View File

@ -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,

1
tensorflow/compiler/mlir/tensorflow/c/BUILD
View File

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

2
tensorflow/compiler/plugin/BUILD
View File

@ -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

7
tensorflow/compiler/tf2tensorrt/BUILD
View File

@ -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",

64
tensorflow/compiler/tf2xla/BUILD
View File

@ -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",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/types:span",
"@com_google_absl//absl/types:variant",
],
] + if_tpu(
if_false = [
"//tensorflow/compiler/mlir:array_container_utils",
"//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 = [
"mlir_bridge_pass_registration.cc",
],
deps = [
":mlir_bridge_pass",
"//tensorflow/compiler/mlir:mlir_graph_optimization_pass_registration",
"//tensorflow/core:core_cpu",
],
srcs = if_tpu(
if_false = [
"mlir_bridge_pass_registration.cc",
],
if_true = [],
),
deps = if_tpu(
if_false = [
":mlir_bridge_pass",
"//tensorflow/compiler/mlir:mlir_graph_optimization_pass_registration",
"//tensorflow/core:core_cpu",
],
if_true = [],
),
alwayslink = 1,
)

1
tensorflow/compiler/tf2xla/cc/BUILD
View File

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

1
tensorflow/compiler/tf2xla/kernels/BUILD
View File

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

3
tensorflow/compiler/tf2xla/lib/BUILD
View File

@ -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

1
tensorflow/compiler/tf2xla/ops/BUILD
View File

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

15
tensorflow/compiler/tf2xla/xla_compiler.cc
View File

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

2
tensorflow/compiler/xla/BUILD
View File

@ -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",

2
tensorflow/compiler/xla/client/BUILD
View File

@ -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(

2
tensorflow/compiler/xla/client/lib/BUILD
View File

@ -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(

1
tensorflow/compiler/xla/pjrt/BUILD
View File

@ -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")

1
tensorflow/compiler/xla/pjrt/distributed/BUILD
View File

@ -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")

1
tensorflow/compiler/xla/python/BUILD
View File

@ -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")

2
tensorflow/compiler/xla/python/tpu_driver/BUILD
View File

@ -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",

7
tensorflow/compiler/xla/python/tpu_driver/client/BUILD
View File

@ -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",

10
tensorflow/compiler/xla/rpc/BUILD
View File

@ -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(

23
tensorflow/compiler/xla/service/BUILD
View File

@ -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",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"@com_google_absl//absl/types:variant",
],
] + internal_hlo_deps(),
)
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",

7
tensorflow/compiler/xla/service/cpu/BUILD
View File

@ -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 = [

11
tensorflow/compiler/xla/service/gpu/BUILD
View File

@ -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",

1
tensorflow/compiler/xla/service/interpreter/BUILD
View File

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

1
tensorflow/compiler/xla/service/spmd/BUILD
View File

@ -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(

10
tensorflow/compiler/xla/tests/BUILD
View File

@ -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",
],

2
tensorflow/compiler/xla/tools/BUILD
View File

@ -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",

3
tensorflow/compiler/xrt/BUILD
View File

@ -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",

5
tensorflow/compiler/xrt/cc/BUILD
View File

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

2
tensorflow/compiler/xrt/kernels/BUILD
View File

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

1
tensorflow/compiler/xrt/tests/BUILD
View File

@ -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",

21
tensorflow/core/BUILD
View File

@ -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,
)

2
tensorflow/core/api_def/BUILD
View File

@ -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",

2
tensorflow/core/api_def/java_api/BUILD
View File

@ -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

7
tensorflow/core/common_runtime/BUILD
View File

@ -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",

2
tensorflow/core/common_runtime/eager/BUILD
View File

@ -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",

3
tensorflow/core/common_runtime/gpu/BUILD
View File

@ -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",

1
tensorflow/core/data/BUILD
View File

@ -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",

12
tensorflow/core/data/service/BUILD
View File

@ -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"],

1
tensorflow/core/debug/BUILD
View File

@ -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",

2
tensorflow/core/distributed_runtime/BUILD
View File

@ -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(

6
tensorflow/core/distributed_runtime/eager/BUILD
View File

@ -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",

2
tensorflow/core/distributed_runtime/rpc/BUILD
View File

@ -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(

1
tensorflow/core/distributed_runtime/rpc/eager/BUILD
View File

@ -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")

2
tensorflow/core/example/BUILD
View File

@ -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",

13
tensorflow/core/framework/BUILD
View File

@ -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__",

2
tensorflow/core/graph/BUILD
View File

@ -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",

2
tensorflow/core/grappler/BUILD
View File

@ -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(

2
tensorflow/core/grappler/clusters/BUILD
View File

@ -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(

2
tensorflow/core/grappler/costs/BUILD
View File

@ -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",

1
tensorflow/core/grappler/graph_analyzer/BUILD
View File

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

1
tensorflow/core/grappler/inputs/BUILD
View File

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

2
tensorflow/core/grappler/inputs/testdata/BUILD vendored
View File

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

9
tensorflow/core/grappler/optimizers/BUILD
View File

@ -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(

1
tensorflow/core/grappler/optimizers/data/BUILD
View File

@ -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")

1
tensorflow/core/grappler/optimizers/data/vectorization/BUILD
View File

@ -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")

2
tensorflow/core/grappler/utils/BUILD
View File

@ -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",

2
tensorflow/core/grappler/verifiers/BUILD
View File

@ -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(

16
tensorflow/core/kernels/BUILD
View File

@ -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",

1
tensorflow/core/kernels/batching_util/BUILD
View File

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

3
tensorflow/core/kernels/boosted_trees/BUILD
View File

@ -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",

2
tensorflow/core/kernels/boosted_trees/quantiles/BUILD
View File

@ -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(

8
tensorflow/core/kernels/data/BUILD
View File

@ -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

8
tensorflow/core/kernels/data/experimental/BUILD
View File

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

2
tensorflow/core/kernels/data/experimental/sql/BUILD
View File

@ -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

7
tensorflow/core/kernels/fuzzing/BUILD
View File

@ -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(

3
tensorflow/core/kernels/hexagon/BUILD
View File

@ -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"],

10
tensorflow/core/kernels/image/BUILD
View File

@ -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",

10
tensorflow/core/kernels/linalg/BUILD
View File

@ -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__",

5
tensorflow/core/kernels/neon/BUILD
View File

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

Some files were not shown because too many files have changed in this diff Show More