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 #include +#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). + """, +)