From 028e652150ca540259a4ec9536e0080c8566d706 Mon Sep 17 00:00:00 2001
From: Randy Dodgen <dodgen@google.com>
Date: Tue, 20 Oct 2020 11:09:33 -0700
Subject: [PATCH] Add a Bazel-level config setting for experimental changes to
registration
Bazel command-line: --//tensorflow:enable_registration_v2=true|false
C++: TF_OPTION_REGISTRATION_V2()
So far, enabling this setting disallows use of 'selective registration', which
has been unsupported for some time.
PiperOrigin-RevId: 338095316
Change-Id: I3d37006e639e98935bd3b2f5788ac060d5a4c692
---
tensorflow/BUILD | 29 ++++++++
tensorflow/core/BUILD | 1 +
tensorflow/core/framework/BUILD | 18 ++++-
.../core/framework/registration_options.h.tpl | 25 +++++++
.../core/framework/selective_registration.h | 14 +++-
tensorflow/tensorflow.bzl | 70 +++++++++++++++++++
6 files changed, 155 insertions(+), 2 deletions(-)
create mode 100644 tensorflow/core/framework/registration_options.h.tpl
diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index 455a0483486..274a829f575 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -3,6 +3,7 @@
# learning applications.
load("@bazel_skylib//lib:selects.bzl", "selects")
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
load("//tensorflow:tensorflow.bzl", "VERSION", "tf_cc_shared_object", "tf_custom_op_library_additional_deps_impl", "tf_native_cc_binary")
load(
"//tensorflow/core/platform:build_config.bzl",
@@ -569,6 +570,33 @@ selects.config_setting_group(
],
)
+# 'enable_registration_v2' opts-in to a different implementation of op and
+# kernel registration - REGISTER_OP, REGISTER_KERNEL_BUILDER, etc.
+#
+# This setting is currently experimental. The 'v2' implementation does _not_
+# correspond to a particular, finalized design; rather, it relates to
+# developing one.
+#
+# The current aim of the 'v2' implementation is to allow 'unused' ops and
+# kernels to be discarded by the linker (to the benefit of binary size).
+bool_flag(
+ name = "enable_registration_v2",
+ build_setting_default = False,
+ visibility = ["//visibility:public"],
+)
+
+config_setting(
+ name = "registration_v1",
+ flag_values = {":enable_registration_v2": "False"},
+ visibility = ["//visibility:public"],
+)
+
+config_setting(
+ name = "registration_v2",
+ flag_values = {":enable_registration_v2": "True"},
+ visibility = ["//visibility:public"],
+)
+
# DO NOT ADD ANY NEW EXCEPTIONS TO THIS LIST!
# Instead, please use public APIs or public build rules TF provides.
# If you need functionality that is not exposed, we will work with you to expand our public APIs.
@@ -612,6 +640,7 @@ bzl_library(
"//third_party/mkl:build_defs_bzl",
"//third_party/mkl_dnn:build_defs_bzl",
"//third_party/ngraph:build_defs_bzl",
+ "@bazel_skylib//rules:common_settings",
"@local_config_cuda//cuda:build_defs_bzl",
"@local_config_rocm//rocm:build_defs_bzl",
"@local_config_tensorrt//:build_defs_bzl",
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 943ead9d451..8613d979d6d 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -448,6 +448,7 @@ tf_cuda_library(
"//tensorflow/core/framework:reader_op_kernel.h",
"//tensorflow/core/framework:register_types.h",
"//tensorflow/core/framework:register_types_traits.h",
+ "//tensorflow/core/framework:registration_options.h",
"//tensorflow/core/framework:resource_mgr.h",
"//tensorflow/core/framework:resource_op_kernel.h",
"//tensorflow/core/framework:rng_alg.h",
diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index e1b6ccb3680..de196f20da9 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -10,6 +10,7 @@ load(
"tf_cc_tests",
"tf_copts",
"tf_cuda_library",
+ "tf_gen_options_header",
)
# buildifier: disable=same-origin-load
@@ -155,6 +156,7 @@ exports_files(
"op.h",
"op_def_builder.h",
"op_def_util.h",
+ "registration_options",
"selective_registration.h",
"shape_inference.h",
],
@@ -216,6 +218,7 @@ filegroup(
"reader_op_kernel.h",
"register_types.h",
"register_types_traits.h",
+ "registration_options.h",
"rendezvous.h",
"resource_handle.h",
"resource_mgr.h",
@@ -400,6 +403,7 @@ filegroup(
"queue_interface.h",
"reader_interface.h",
"register_types_traits.h",
+ "registration_options.h",
"rendezvous.cc",
"rendezvous.h",
"resource_mgr.cc",
@@ -967,9 +971,21 @@ cc_library(
],
)
+tf_gen_options_header(
+ name = "gen_registration_options",
+ build_settings = {
+ "//tensorflow:enable_registration_v2": "REGISTRATION_V2",
+ },
+ output_header = "registration_options.h",
+ template = "registration_options.h.tpl",
+)
+
cc_library(
name = "selective_registration",
- hdrs = ["selective_registration.h"],
+ hdrs = [
+ "registration_options.h",
+ "selective_registration.h",
+ ],
deps = tf_selective_registration_deps(),
)
diff --git a/tensorflow/core/framework/registration_options.h.tpl b/tensorflow/core/framework/registration_options.h.tpl
new file mode 100644
index 00000000000..375a1088b51
--- /dev/null
+++ b/tensorflow/core/framework/registration_options.h.tpl
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FRAMEWORK_REGISTRATION_OPTIONS_TMPL_H_
+#define TENSORFLOW_CORE_FRAMEWORK_REGISTRATION_OPTIONS_TMPL_H_
+
+// This header is generated from a template; see the tf_gen_options_header()
+// build rule. Template placeholders of the form '#define_option X' result in
+// macros of the form 'TF_OPTION_X()'.
+
+#define_option REGISTRATION_V2
+
+#endif // TENSORFLOW_CORE_FRAMEWORK_REGISTRATION_OPTIONS_TMPL_H_
diff --git a/tensorflow/core/framework/selective_registration.h b/tensorflow/core/framework/selective_registration.h
index c9bbcb8bfe8..06ea0e00004 100644
--- a/tensorflow/core/framework/selective_registration.h
+++ b/tensorflow/core/framework/selective_registration.h
@@ -35,6 +35,10 @@ limitations under the License.
#include <type_traits>
#include <utility>
+#include "tensorflow/core/framework/registration_options.h"
+
+#if !TF_OPTION_REGISTRATION_V2()
+
#ifdef SELECTIVE_REGISTRATION
// Experimental selective registration support to reduce binary size.
@@ -66,12 +70,20 @@ limitations under the License.
!defined(SHOULD_REGISTER_OP_KERNEL))
static_assert(false, "ops_to_register.h must define SHOULD_REGISTER macros");
#endif
-#else
+#else // SELECTIVE_REGISTRATION
#define SHOULD_REGISTER_OP(op) true
#define SHOULD_REGISTER_OP_GRADIENT true
#define SHOULD_REGISTER_OP_KERNEL(clz) true
+#endif // SELECTIVE_REGISTRATION
+
+#else // ! TF_OPTION_REGISTRATION_V2()
+
+#ifdef SELECTIVE_REGISTRATION
+#error TF_OPTION_REGISTRATION_V2(): Compile-time selective registration is not supported
#endif
+#endif // ! TF_OPTION_REGISTRATION_V2()
+
namespace tensorflow {
// An InitOnStartupMarker is 'initialized' on program startup, purely for the
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index 39a822065da..c5b756a3398 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -51,6 +51,7 @@ load(
"//third_party/ngraph:build_defs.bzl",
"if_ngraph",
)
+load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
# version for the shared libraries, can
# not contain rc or alpha, only numbers.
@@ -262,6 +263,12 @@ def if_libtpu(if_true, if_false = []):
"//conditions:default": if_false,
})
+def if_registration_v2(if_true, if_false = []):
+ return select({
+ "//tensorflow:registration_v2": if_true,
+ "//conditions:default": if_false,
+ })
+
# Linux systems may required -lrt linker flag for e.g. clock_gettime
# see https://github.com/tensorflow/tensorflow/issues/15129
def lrt_if_needed():
@@ -2822,3 +2829,66 @@ def internal_tfrt_deps():
def internal_cuda_deps():
return []
+
+def _tf_gen_options_header_impl(ctx):
+ header_depset = depset([ctx.outputs.output_header])
+
+ define_vals = {True: "true", False: "false"}
+ substitutions = {}
+ for target, identifier in ctx.attr.build_settings.items():
+ setting_val = target[BuildSettingInfo].value
+ lines = [
+ "// %s" % target.label,
+ "#define TF_OPTION_%s() %s" % (identifier, define_vals[setting_val]),
+ ]
+ substitutions["#define_option %s" % identifier] = "\n".join(lines)
+
+ ctx.actions.expand_template(
+ template = ctx.file.template,
+ output = ctx.outputs.output_header,
+ substitutions = substitutions,
+ )
+
+ return [
+ DefaultInfo(files = header_depset),
+ ]
+
+tf_gen_options_header = rule(
+ attrs = {
+ "output_header": attr.output(
+ doc = "File path for the generated header (output)",
+ mandatory = True,
+ ),
+ "template": attr.label(
+ doc = """Template for the header.
+ For each option name 'X' (see build_settings attribute),
+ '#define_option X' results in a macro 'TF_OPTION_X()'
+ """,
+ allow_single_file = True,
+ mandatory = True,
+ ),
+ "build_settings": attr.label_keyed_string_dict(
+ doc = """Dictionary from build-setting labels to option names. Example:
+ {"//tensorflow:x_setting" : "X"}
+ """,
+ providers = [BuildSettingInfo],
+ ),
+ },
+ implementation = _tf_gen_options_header_impl,
+ doc = """
+ Generates a header file for Bazel build settings.
+
+ This is an alternative to setting preprocessor defines on the compiler
+ command line. It has a few advantages:
+ - Usage of the options requires #include-ing the header, and thus a
+ Bazel-level dependency.
+ - Each option has a definition site in source code, which mentions the
+ corresponding Bazel setting. This is particularly useful when
+ navigating code with the assistance of static analysis (e.g.
+ https://cs.opensource.google/tensorflow).
+ - Each option is represented as a FUNCTION()-style macro, which is always
+ defined (i.e. one uses #if instead of #ifdef). This allows forms like
+ 'if constexpr (TF_OPTION_FOO()) { ... }', and helps catch missing
+ dependencies (if 'F' is undefined, '#if F()' results in an error).
+ """,
+)