From 5b5aab7f63ad9667edee3bb837adc4033ccbeac3 Mon Sep 17 00:00:00 2001 From: Russell Power Date: Fri, 2 Oct 2020 20:20:11 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 335147548 Change-Id: Ib445cfbcb28421b4eb522d4d9524e4a64fe631df --- tensorflow/c/eager/BUILD | 6 +- tensorflow/c/eager/c_api.cc | 4 +- tensorflow/compiler/jit/BUILD | 10 +- .../compiler/jit/xla_compilation_cache.cc | 4 +- tensorflow/compiler/tf2xla/BUILD | 10 +- tensorflow/compiler/tf2xla/xla_compiler.cc | 4 +- tensorflow/compiler/xla/rpc/BUILD | 4 +- tensorflow/core/BUILD | 7 +- tensorflow/core/common_runtime/BUILD | 4 +- .../core/platform/default/build_config.bzl | 4 +- tensorflow/core/tpu/kernels/BUILD | 120 +++++++----------- .../tpu/kernels/tpu_compilation_cache_grpc.cc | 6 +- .../tpu/kernels/tpu_compilation_cache_grpc.h | 6 +- .../tpu_compilation_cache_rpc_lookup.cc | 2 +- .../tpu_compilation_cache_rpc_support.cc | 6 +- .../tpu/kernels/tpu_compilation_metrics.cc | 4 +- .../core/tpu/kernels/tpu_compile_op_impl.cc | 4 +- tensorflow/core/tpu/kernels/tpu_pod_state.cc | 4 +- tensorflow/stream_executor/gpu/BUILD | 4 +- tensorflow/tensorflow.bzl | 6 +- 20 files changed, 94 insertions(+), 125 deletions(-) diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index 54771ffa840..b90b2644269 100644 --- a/tensorflow/c/eager/BUILD +++ b/tensorflow/c/eager/BUILD @@ -3,7 +3,7 @@ load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") load( "//tensorflow:tensorflow.bzl", - "if_tpu", + "if_libtpu", "tf_cc_test", "tf_copts", "tf_cuda_cc_test", @@ -289,7 +289,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/lib/llvm_rtti", - ] + if_tpu( + ] + if_libtpu( if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"], if_true = [], ), @@ -354,7 +354,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/lib/llvm_rtti", - ] + if_tpu( + ] + if_libtpu( if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"], if_true = [], ), diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index 91c5a72ad64..5f388bfe0cd 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/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" -#if defined(PLATFORM_GOOGLE) && !defined(LIBTFTPU) +#if defined(PLATFORM_GOOGLE) && !defined(LIBTPU_ON_GCE) #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) { -#if defined(PLATFORM_GOOGLE) && !defined(LIBTFTPU) +#if defined(PLATFORM_GOOGLE) && !defined(LIBTPU_ON_GCE) return tensorflow::wrap(new tfrt::tf::ContextInterface(opts->async)); #else status->status = tensorflow::errors::Unimplemented("TFRT is not supported"); diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index b3655dcba63..da3db1789b5 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -4,7 +4,7 @@ load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") 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:tensorflow.bzl", "if_libtpu", "tf_copts") load("//tensorflow/stream_executor:build_defs.bzl", "if_cuda_or_rocm") # buildifier: disable=same-origin-load @@ -77,7 +77,7 @@ cc_library( "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", - ] + if_tpu( + ] + if_libtpu( if_false = ["//tensorflow/compiler/xla/service:cpu_plugin"], if_true = [], ), @@ -114,7 +114,7 @@ cc_library( "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", - ] + if_tpu( + ] + if_libtpu( if_false = [ "//tensorflow/compiler/xla/service:cpu_plugin", # buildcleaner: keep ], @@ -141,7 +141,7 @@ cc_library( "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "//tensorflow/core/common_runtime/gpu:gpu_init", - ] + if_tpu( + ] + if_libtpu( if_false = [ "//tensorflow/compiler/xla/service:gpu_plugin", # buildcleaner: keep ], @@ -375,7 +375,7 @@ cc_library( "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:logging", - ] + if_tpu( + ] + if_libtpu( if_false = [ "//tensorflow/compiler/mlir:array_container_utils", "//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes", diff --git a/tensorflow/compiler/jit/xla_compilation_cache.cc b/tensorflow/compiler/jit/xla_compilation_cache.cc index 3ab3c19e439..91c2e0f9fdb 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -47,7 +47,7 @@ limitations under the License. #include "tensorflow/core/public/version.h" #include "tensorflow/core/util/dump_graph.h" -#if !defined(LIBTFTPU) +#if !defined(LIBTPU_ON_GCE) #include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h" #include "tensorflow/compiler/mlir/utils/array_container_utils.h" #endif @@ -289,7 +289,7 @@ Status XlaCompilationCache::CompileSingleOp( }); const ConfigProto* config = ctx->function_library()->config_proto(); bool use_mlir = config && config->experimental().enable_mlir_bridge(); -#ifdef LIBTFTPU +#ifdef LIBTPU_ON_GCE if (use_mlir && has_tensor_list_arg) { LOG(WARNING) << "MLIR is not supported in this environment."; } diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 7b289f00bd0..03510701af4 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -1,5 +1,5 @@ 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:tensorflow.bzl", "if_libtpu", "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", @@ -298,7 +298,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/stream_executor:platform", - ] + if_tpu( + ] + if_libtpu( if_false = [ "//tensorflow/compiler/xla/service:cpu_plugin", "//tensorflow/compiler/xla/service/cpu:buffer_info_util", @@ -369,7 +369,7 @@ cc_library( "//tensorflow/core:lib_internal", "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", - ] + if_tpu( + ] + if_libtpu( if_false = [ "//tensorflow/compiler/mlir:array_container_utils", "//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes", @@ -877,13 +877,13 @@ cc_library( cc_library( name = "mlir_bridge_pass_registration", - srcs = if_tpu( + srcs = if_libtpu( if_false = [ "mlir_bridge_pass_registration.cc", ], if_true = [], ), - deps = if_tpu( + deps = if_libtpu( if_false = [ ":mlir_bridge_pass", "//tensorflow/compiler/mlir:mlir_graph_optimization_pass_registration", diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index e67e183aef9..c62b8286bbe 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -56,7 +56,7 @@ limitations under the License. #include "tensorflow/core/protobuf/graph_debug_info.pb.h" #include "tensorflow/core/util/dump_graph.h" -#ifndef LIBTFTPU +#ifndef LIBTPU_ON_GCE #include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h" #include "tensorflow/compiler/mlir/utils/array_container_utils.h" #endif @@ -733,7 +733,7 @@ Status XlaCompiler::CompileFunction( } VLOG(1) << "===================================================="; -#ifdef LIBTFTPU +#ifdef LIBTPU_ON_GCE if (GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge) { VLOG(1) << "MLIR is not supported in this environment."; } diff --git a/tensorflow/compiler/xla/rpc/BUILD b/tensorflow/compiler/xla/rpc/BUILD index cb2c194858c..15022d1a879 100644 --- a/tensorflow/compiler/xla/rpc/BUILD +++ b/tensorflow/compiler/xla/rpc/BUILD @@ -2,7 +2,7 @@ load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency") load( "//tensorflow:tensorflow.bzl", - "if_tpu", + "if_libtpu", "tf_cc_binary", "tf_cc_test", ) @@ -57,7 +57,7 @@ cc_library( "//tensorflow/core:framework_internal", "//tensorflow/core:lib", tf_grpc_cc_dependency(), - ] + if_tpu( + ] + if_libtpu( if_false = ["//tensorflow/compiler/xla/service:cpu_plugin"], if_true = [], ), diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index a664c96d15b..a1149ec1187 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -68,9 +68,9 @@ load( "if_chromiumos", "if_cuda_or_rocm", "if_ios", + "if_libtpu", "if_mobile", "if_not_windows", - "if_tpu", "tf_android_core_proto_headers", "tf_cc_test", "tf_cc_test_mkl", @@ -894,8 +894,7 @@ cc_library( "//tensorflow/c/kernels:summary_op_lib", ] + if_chromiumos( [], - # Non-tpu platforms don't need tpu dependency. It would be best to guard - # them by if_tpu. But there is no such flag yet. + # Non-tpu platforms don't need tpu dependency. [ ":tpu_configuration_ops_op_lib", ":tpu_cross_replica_ops_op_lib", @@ -916,7 +915,7 @@ cc_library( ]) + if_tensorrt([ "//tensorflow/compiler/tf2tensorrt:trt_engine_resource_ops_op_lib", "//tensorflow/compiler/tf2tensorrt:trt_op_libs", - ]) + if_tpu( + ]) + if_libtpu( if_false = ["//tensorflow/compiler/mlir/tensorflow:mlir_passthrough_op"], if_true = [], ), diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD index f0378ad7538..66dfdffaccf 100644 --- a/tensorflow/core/common_runtime/BUILD +++ b/tensorflow/core/common_runtime/BUILD @@ -1,6 +1,6 @@ load( "//tensorflow:tensorflow.bzl", - "if_tpu", + "if_libtpu", "tf_cc_test", "tf_cc_test_mkl", "tf_cc_tests", @@ -93,7 +93,7 @@ cc_library( deps = [ ":core_cpu", "//tensorflow/core/common_runtime/gpu:gpu_runtime", - ] + if_tpu(["//tensorflow/core/tpu:tpu_runtime"]), + ] + if_libtpu(["//tensorflow/core/tpu:tpu_runtime"]), ) filegroup( diff --git a/tensorflow/core/platform/default/build_config.bzl b/tensorflow/core/platform/default/build_config.bzl index 9a1068d5a3a..d049ecec70a 100644 --- a/tensorflow/core/platform/default/build_config.bzl +++ b/tensorflow/core/platform/default/build_config.bzl @@ -1,7 +1,7 @@ # Platform-specific build configurations. load("@com_google_protobuf//:protobuf.bzl", "proto_gen") -load("//tensorflow:tensorflow.bzl", "clean_dep", "if_not_windows", "if_tpu") +load("//tensorflow:tensorflow.bzl", "clean_dep", "if_libtpu", "if_not_windows") load("//tensorflow/core/platform:build_config_root.bzl", "if_static") load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm") @@ -814,4 +814,4 @@ def if_llvm_system_z_available(then, otherwise = []): }) def tf_tpu_dependencies(): - return if_tpu(["//tensorflow/core/tpu/kernels"]) + return if_libtpu(["//tensorflow/core/tpu/kernels"]) diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD index 90107abf4b6..8de50acfd6c 100644 --- a/tensorflow/core/tpu/kernels/BUILD +++ b/tensorflow/core/tpu/kernels/BUILD @@ -5,13 +5,11 @@ load( "//tensorflow/core/platform:build_config.bzl", "tf_proto_library", ) +load("//tensorflow:tensorflow.bzl", "if_libtpu", "tf_copts") load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency") # buildifier: disable=same-origin-load load("//tensorflow:tensorflow.bzl", "tf_kernel_library") # buildifier: disable=same-origin-load # Config setting to enable go/libtpu support. -WITH_TPU_SUPPORT = "//tensorflow:with_tpu_support" - -DEFAULT = "//conditions:default" package( default_visibility = [ @@ -44,10 +42,10 @@ cc_library( name = "tpu_compile_op_common", srcs = ["tpu_compile_op_common.cc"], hdrs = ["tpu_compile_op_common.h"], - deps = select({ - WITH_TPU_SUPPORT: [":tpu_compilation_metrics"], - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_compilation_metrics"], - }) + [ + deps = if_libtpu( + [":tpu_compilation_metrics"], + ["//tensorflow/core/tpu/kernels:tpu_compilation_metrics"], + ) + [ ":tpu_compilation_cache_entry_unloader", ":tpu_compilation_cache_interface", ":tpu_compilation_metrics_hdrs", @@ -97,14 +95,10 @@ tf_kernel_library( name = "tpu_configuration_ops", srcs = ["tpu_configuration_ops.cc"], hdrs = ["tpu_configuration_ops.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), - deps = select({ - WITH_TPU_SUPPORT: [":tpu_util"], - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_util"], - }) + [ + deps = if_libtpu( + [":tpu_util"], + ["//tensorflow/core/tpu/kernels:tpu_util"], + ) + [ ":tpu_compilation_cache_factory", ":tpu_compilation_cache_interface", ":tpu_compilation_cache_local_lookup", @@ -346,10 +340,10 @@ cc_library( name = "tpu_compilation_cache_interface", srcs = ["tpu_compilation_cache_interface.cc"], hdrs = ["tpu_compilation_cache_interface.h"], - deps = select({ - WITH_TPU_SUPPORT: [":tpu_compilation_metrics"], - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_compilation_metrics"], - }) + [ + deps = if_libtpu( + [":tpu_compilation_metrics"], + ["//tensorflow/core/tpu/kernels:tpu_compilation_metrics"], + ) + [ ":compiled_subgraph", ":tpu_compilation_cache_common_proto_cc", ":tpu_compilation_cache_entry", @@ -424,10 +418,7 @@ cc_library( cc_library( name = "tpu_compilation_metrics", srcs = ["tpu_compilation_metrics.cc"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), + copts = tf_copts(), deps = [ ":tpu_compilation_metrics_hdrs", ], @@ -529,14 +520,11 @@ cc_library( cc_library( name = "tpu_compilation_cache_rpc_support_hdrs", hdrs = ["tpu_compilation_cache_rpc_support.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), - deps = select({ - WITH_TPU_SUPPORT: [":tpu_compilation_cache_proto_cc"], # build_cleaner: keep - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_compilation_cache_cc_proto"], # build_cleaner: keep - }) + [ + copts = tf_copts(), + deps = if_libtpu( + [":tpu_compilation_cache_proto_cc"], + ["//tensorflow/core/tpu/kernels:tpu_compilation_cache_cc_proto"], + ) + [ ":tpu_compilation_cache_entry", ":tpu_compilation_cache_interface", ":tpu_compilation_cache_lookup", @@ -550,10 +538,7 @@ cc_library( cc_library( name = "tpu_compilation_cache_rpc_support", srcs = ["tpu_compilation_cache_rpc_support.cc"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), + copts = tf_copts(), deps = [ ":tpu_compilation_cache_common_proto_cc", ":tpu_compilation_cache_proto_cc", @@ -572,14 +557,11 @@ cc_library( name = "tpu_compilation_cache_rpc_lookup", srcs = ["tpu_compilation_cache_rpc_lookup.cc"], hdrs = ["tpu_compilation_cache_rpc_lookup.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), - deps = select({ - WITH_TPU_SUPPORT: [":tpu_compilation_cache_rpc_support"], - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_compilation_cache_rpc_support"], - }) + [ + copts = tf_copts(), + deps = if_libtpu( + [":tpu_compilation_cache_rpc_support"], + ["//tensorflow/core/tpu/kernels:tpu_compilation_cache_rpc_support"], + ) + [ ":tpu_compilation_cache_grpc", ":tpu_compilation_cache_interface", ":tpu_compilation_cache_lookup", @@ -617,14 +599,11 @@ cc_library( name = "tpu_compilation_cache_grpc", srcs = ["tpu_compilation_cache_grpc.cc"], hdrs = ["tpu_compilation_cache_grpc.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), - deps = select({ - WITH_TPU_SUPPORT: [":tpu_compilation_cache_proto_cc"], - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_compilation_cache_cc_proto"], - }) + [ + copts = tf_copts(), + deps = if_libtpu( + [":tpu_compilation_cache_proto_cc"], + ["//tensorflow/core/tpu/kernels:tpu_compilation_cache_cc_proto"], + ) + [ ":tpu_compilation_cache_common_proto_cc", tf_grpc_cc_dependency(), ], @@ -634,20 +613,17 @@ cc_library( name = "tpu_compilation_cache_service", srcs = ["tpu_compilation_cache_service.cc"], hdrs = ["tpu_compilation_cache_service.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), - deps = select({ - WITH_TPU_SUPPORT: [ - ":tpu_compilation_cache_rpc_support", # build_cleaner: keep - ":tpu_compilation_cache_proto_cc", # build_cleaner: keep + copts = tf_copts(), + deps = if_libtpu( + [ + ":tpu_compilation_cache_rpc_support", + ":tpu_compilation_cache_proto_cc", ], - DEFAULT: [ - "//tensorflow/core/tpu/kernels:tpu_compilation_cache_rpc_support", # build_cleaner: keep - "//tensorflow/core/tpu/kernels:tpu_compilation_cache_cc_proto", # build_cleaner: keep + [ + "//tensorflow/core/tpu/kernels:tpu_compilation_cache_rpc_support", + "//tensorflow/core/tpu/kernels:tpu_compilation_cache_cc_proto", ], - }) + [ + ) + [ ":tpu_compilation_cache_common_proto_cc", ":tpu_compilation_cache_grpc", ":tpu_compilation_cache_interface", @@ -704,10 +680,7 @@ cc_library( name = "tpu_compile_op_impl", srcs = ["tpu_compile_op_impl.cc"], hdrs = ["tpu_compile_op_impl.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), + copts = tf_copts(), deps = [ ":tpu_compilation_cache_key", ":tpu_compile_c_api_hdrs", @@ -952,14 +925,11 @@ cc_library( name = "tpu_pod_state", srcs = ["tpu_pod_state.cc"], hdrs = ["tpu_pod_state.h"], - copts = select({ - WITH_TPU_SUPPORT: ["-DLIBTFTPU"], - DEFAULT: [], - }), - deps = select({ - WITH_TPU_SUPPORT: [":tpu_util"], - DEFAULT: ["//tensorflow/core/tpu/kernels:tpu_util"], - }) + [ + copts = tf_copts(), + deps = if_libtpu( + [":tpu_util"], + ["//tensorflow/core/tpu/kernels:tpu_util"], + ) + [ ":tpu_compilation_cache_service", "//tensorflow/c:tf_status", "//tensorflow/c:tf_status_helper", diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.cc b/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.cc index 207a60e7b48..c3aa62805c0 100644 --- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.cc +++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.cc @@ -30,11 +30,11 @@ namespace tensorflow { namespace tpu { static const char* grpcTpuCompilationCacheService_method_names[] = { -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) "/tensorflow.tpu.TpuCompilationCacheServiceExternal/GetTpuProgram", -#else // LIBTFTPU +#else // LIBTPU_ON_GCE "/tensorflow.tpu.TpuCompilationCacheService/GetTpuProgram", -#endif // LIBTFTPU +#endif // LIBTPU_ON_GCE }; std::unique_ptr diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.h b/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.h index 324fc9e6f08..55877d15df2 100644 --- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.h +++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_grpc.h @@ -35,7 +35,7 @@ limitations under the License. #include -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) #include "tensorflow/core/tpu/kernels/tpu_compilation_cache.pb.h" #else #include "tensorflow/core/tpu/kernels/tpu_compilation_cache.pb.h" // copybara" @@ -48,7 +48,7 @@ namespace grpc { class TpuCompilationCacheService final { public: using RequestType = ::tensorflow::tpu::GetTpuProgramRequest; -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) using ResponseType = ::tensorflow::tpu::GetTpuProgramResponseExternal; #else using ResponseType = ::tensorflow::tpu::GetTpuProgramResponse; @@ -59,7 +59,7 @@ class TpuCompilationCacheService final { enum class MethodId { kGetTpuProgram = 0 }; static constexpr char const* service_full_name() { -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) return "tensorflow.tpu.TpuCompilationCacheServiceExternal"; #else return "tensorflow.tpu.TpuCompilationCacheService"; diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_lookup.cc b/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_lookup.cc index 8b0fb674682..7846cc7bbb3 100644 --- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_lookup.cc +++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_lookup.cc @@ -25,7 +25,7 @@ namespace tensorflow { namespace tpu { namespace { -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) using ResponseType = GetTpuProgramResponseExternal; #else using ResponseType = GetTpuProgramResponse; diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_support.cc b/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_support.cc index 9a6ca6be7e4..29ec8701a37 100644 --- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_support.cc +++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_rpc_support.cc @@ -17,7 +17,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" #include "tensorflow/core/platform/casts.h" -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) #include "tensorflow/core/tpu/kernels/tpu_compilation_cache.pb.h" #endif #include "tensorflow/core/tpu/kernels/tpu_compilation_cache_common.pb.h" @@ -30,7 +30,7 @@ std::shared_ptr<::grpc::ChannelCredentials> CreateChannelCredentials() { return ::grpc::InsecureChannelCredentials(); // NOLINT } -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) template <> Status DeserializeRpcResponseToCacheEntry( absl::string_view local_proto_key, GetTpuProgramResponseExternal* response, @@ -156,6 +156,6 @@ xla::StatusOr> SerializeCacheEntryToBufferSlices( return std::vector<::grpc::Slice>{::grpc::Slice(encoded_header)}; } -#endif // LIBTFTPU +#endif // LIBTPU_ON_GCE } // namespace tpu } // namespace tensorflow diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_metrics.cc b/tensorflow/core/tpu/kernels/tpu_compilation_metrics.cc index e1a65ad0f32..ce982a1bd9a 100644 --- a/tensorflow/core/tpu/kernels/tpu_compilation_metrics.cc +++ b/tensorflow/core/tpu/kernels/tpu_compilation_metrics.cc @@ -19,7 +19,7 @@ namespace tpu { // TODO(henrytan): remove this once `TpuCompilationCache` migration to OSS is // completed. -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) /* static */ void TpuCompilationMetrics::IncrementCacheLookupCount( bool is_cache_hit, absl::string_view session_name) { @@ -36,7 +36,7 @@ void TpuCompilationMetrics::IncrementCompilationCount( absl::string_view session_name) { // A placeholder for tracking metrics. } -#endif // LIBTFTPU +#endif // LIBTPU_ON_GCE } // namespace tpu } // namespace tensorflow diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_impl.cc b/tensorflow/core/tpu/kernels/tpu_compile_op_impl.cc index 8703dd818f5..270c2c53d7a 100644 --- a/tensorflow/core/tpu/kernels/tpu_compile_op_impl.cc +++ b/tensorflow/core/tpu/kernels/tpu_compile_op_impl.cc @@ -68,11 +68,11 @@ class TpuCompileOpImplFactory : public CompileOpImplFactory { } }; -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) REGISTER_MODULE_INITIALIZER(tpu_compile_op_impl_factory, { VLOG(1) << "register TpuCompileOpImplFactory()"; CompileOpImplFactory::Register(new TpuCompileOpImplFactory()); }); -#endif // LIBTFTPU +#endif // LIBTPU_ON_GCE } // namespace tpu } // namespace tensorflow diff --git a/tensorflow/core/tpu/kernels/tpu_pod_state.cc b/tensorflow/core/tpu/kernels/tpu_pod_state.cc index 7b02998b343..898f02b28e9 100644 --- a/tensorflow/core/tpu/kernels/tpu_pod_state.cc +++ b/tensorflow/core/tpu/kernels/tpu_pod_state.cc @@ -18,7 +18,7 @@ limitations under the License. #include "tensorflow/c/tf_status_helper.h" #include "tensorflow/core/tpu/tpu_api.h" -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) #include "tensorflow/core/tpu/kernels/tpu_util.h" #else #include "tensorflow/core/tpu/kernels/tpu_util.h" // copybara" @@ -54,7 +54,7 @@ xla::StatusOr> ConstructCacheService(ResourceMgr* rmgr, int serving_port, tpu::TpuCompilationCacheInterface* compilation_cache) { xla::StatusOr> server_builder; -#if defined(LIBTFTPU) +#if defined(LIBTPU_ON_GCE) server_builder = tpu::CreateServerBuilder(serving_port); #else server_builder = tpu::CreateServerBuilderGoogle(serving_port); diff --git a/tensorflow/stream_executor/gpu/BUILD b/tensorflow/stream_executor/gpu/BUILD index 7fbb40e0ae9..8626c34d383 100644 --- a/tensorflow/stream_executor/gpu/BUILD +++ b/tensorflow/stream_executor/gpu/BUILD @@ -10,7 +10,7 @@ load( "//tensorflow/core/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", ) -load("//tensorflow:tensorflow.bzl", "if_tpu", "tf_copts") +load("//tensorflow:tensorflow.bzl", "if_libtpu", "tf_copts") load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load( "//tensorflow/core/platform:rules_cc.bzl", @@ -70,7 +70,7 @@ cc_library( "//tensorflow/stream_executor:device_options", "//tensorflow/stream_executor/lib", "//tensorflow/stream_executor/platform", - ] + if_tpu( + ] + if_libtpu( if_false = ["@local_config_cuda//cuda:cuda_headers"], if_true = [], ), diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 8ed12136c55..990130e1702 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -261,8 +261,8 @@ def if_nccl(if_true, if_false = []): "//conditions:default": if_true, }) -def if_tpu(if_true, if_false = []): - """Shorthand for select()ing whether to build for TPUs.""" +def if_libtpu(if_true, if_false = []): + """Shorthand for select()ing whether to build support for using TPUs via libtpu.so""" return select({ str(Label("//tensorflow:with_tpu_support")): if_true, "//conditions:default": if_false, @@ -328,7 +328,7 @@ def tf_copts( (if_not_windows(["-fno-exceptions"]) if not allow_exceptions else []) + if_cuda(["-DGOOGLE_CUDA=1"]) + if_nvcc(["-DTENSORFLOW_USE_NVCC=1"]) + - if_tpu(["-DLIBTFTPU"]) + + if_libtpu(["-DLIBTPU_ON_GCE"], []) + if_xla_available(["-DTENSORFLOW_USE_XLA=1"]) + if_tensorrt(["-DGOOGLE_TENSORRT=1"]) + if_mkl(["-DINTEL_MKL=1", "-DENABLE_MKLDNN_V1", "-DENABLE_INTEL_MKL_BFLOAT16"]) +