1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // This file provides some common support for 'registration' of e.g. ops and 17 // kernels. In particular, it relates to the REGISTER_OP (op registration) and 18 // REGISTER_KERNEL_BUILDER (kernel registration) macros. 19 // 20 // Note that there are two sides to 'registration': 21 // - Definition (compile-time): making op and kernel definitions _available_. 22 // - Usage (run-time): adding particular (available) definitions of ops and 23 // kernels to the global OpRegistry / KernelRegistry, to be found when 24 // constructing and executing graphs. 25 // 26 // Currently, definition and usage happen to be coupled together: all 27 // 'available' definitions (from the REGISTER_*' macros) are added to the global 28 // registries on startup / library load. 29 30 #ifndef TENSORFLOW_CORE_FRAMEWORK_REGISTRATION_REGISTRATION_H_ 31 #define TENSORFLOW_CORE_FRAMEWORK_REGISTRATION_REGISTRATION_H_ 32 33 #include <string.h> 34 35 #include <type_traits> 36 #include <utility> 37 38 #include "tensorflow/core/framework/registration/options.h" 39 40 #if !TF_OPTION_REGISTRATION_V2() 41 42 #ifdef SELECTIVE_REGISTRATION 43 44 // Experimental selective registration support to reduce binary size. 45 // 46 // To use selective registration, when building: 47 // 1. define SELECTIVE_REGISTRATION, e.g. in gcc by passing 48 // -DSELECTIVE_REGISTRATION to compilation. 49 // 2. Provide ops_to_register.h. This file is not included in the repo and must 50 // be placed by the user or a tool where the compiler can find it. It must 51 // define the constants and functions used in the macros below. The 52 // functions should be defined as valid constexpr functions, so that they are 53 // evaluated at compile time: this is needed to make symbols referenced by 54 // un-registered objects unused, and therefore allow the linker to strip them 55 // out. See python/tools/print_selective_registration_header.py for a tool 56 // that can be used to generate ops_to_register.h. 57 // 58 // ops_to_register.h should define macros for: 59 // // Ops for which this is false will not be registered. 60 // SHOULD_REGISTER_OP(op) 61 // // If this is false, then no gradient ops are registered. 62 // SHOULD_REGISTER_OP_GRADIENT 63 // // Op kernel classes where this is false won't be registered. 64 // SHOULD_REGISTER_OP_KERNEL(clz) 65 // The macros should be defined using constexprs. 66 67 #include "ops_to_register.h" 68 69 #if (!defined(SHOULD_REGISTER_OP) || !defined(SHOULD_REGISTER_OP_GRADIENT) || \ 70 !defined(SHOULD_REGISTER_OP_KERNEL)) 71 static_assert(false, "ops_to_register.h must define SHOULD_REGISTER macros"); 72 #endif 73 #else // SELECTIVE_REGISTRATION 74 #define SHOULD_REGISTER_OP(op) true 75 #define SHOULD_REGISTER_OP_GRADIENT true 76 #define SHOULD_REGISTER_OP_KERNEL(clz) true 77 #endif // SELECTIVE_REGISTRATION 78 79 #else // ! TF_OPTION_REGISTRATION_V2() 80 81 #ifdef SELECTIVE_REGISTRATION 82 #error TF_OPTION_REGISTRATION_V2(): Compile-time selective registration is not supported 83 #endif 84 85 #endif // ! TF_OPTION_REGISTRATION_V2() 86 87 namespace tensorflow { 88 89 // An InitOnStartupMarker is 'initialized' on program startup, purely for the 90 // side-effects of that initialization - the struct itself is empty. (The type 91 // is expected to be used to define globals.) 92 // 93 // The '<<' operator should be used in initializer expressions to specify what 94 // to run on startup. The following values are accepted: 95 // - An InitOnStartupMarker. Example: 96 // InitOnStartupMarker F(); 97 // InitOnStartupMarker const kInitF = 98 // InitOnStartupMarker{} << F(); 99 // - Something to call, which returns an InitOnStartupMarker. Example: 100 // InitOnStartupMarker const kInit = 101 // InitOnStartupMarker{} << []() { G(); return 102 // 103 // See also: TF_INIT_ON_STARTUP_IF 104 struct InitOnStartupMarker { 105 constexpr InitOnStartupMarker operator<<(InitOnStartupMarker) const { 106 return *this; 107 } 108 109 template <typename T> 110 constexpr InitOnStartupMarker operator<<(T&& v) const { 111 return std::forward<T>(v)(); 112 } 113 }; 114 115 // Conditional initializer expressions for InitOnStartupMarker: 116 // TF_INIT_ON_STARTUP_IF(cond) << f 117 // If 'cond' is true, 'f' is evaluated (and called, if applicable) on startup. 118 // Otherwise, 'f' is *not evaluated*. Note that 'cond' is required to be a 119 // constant-expression, and so this approximates #ifdef. 120 // 121 // The implementation uses the ?: operator (!cond prevents evaluation of 'f'). 122 // The relative precedence of ?: and << is significant; this effectively expands 123 // to (see extra parens): 124 // !cond ? InitOnStartupMarker{} : (InitOnStartupMarker{} << f) 125 // 126 // Note that although forcing 'cond' to be a constant-expression should not 127 // affect binary size (i.e. the same optimizations should apply if it 'happens' 128 // to be one), it was found to be necessary (for a recent version of clang; 129 // perhaps an optimizer bug). 130 // 131 // The parens are necessary to hide the ',' from the preprocessor; it could 132 // otherwise act as a macro argument separator. 133 #define TF_INIT_ON_STARTUP_IF(cond) \ 134 (::std::integral_constant<bool, !(cond)>::value) \ 135 ? ::tensorflow::InitOnStartupMarker{} \ 136 : ::tensorflow::InitOnStartupMarker {} 137 138 // Wrapper for generating unique IDs (for 'anonymous' InitOnStartup definitions) 139 // using __COUNTER__. The new ID (__COUNTER__ already expanded) is provided as a 140 // macro argument. 141 // 142 // Usage: 143 // #define M_IMPL(id, a, b) ... 144 // #define M(a, b) TF_NEW_ID_FOR_INIT(M_IMPL, a, b) 145 #define TF_NEW_ID_FOR_INIT_2(m, c, ...) m(c, __VA_ARGS__) 146 #define TF_NEW_ID_FOR_INIT_1(m, c, ...) TF_NEW_ID_FOR_INIT_2(m, c, __VA_ARGS__) 147 #define TF_NEW_ID_FOR_INIT(m, ...) \ 148 TF_NEW_ID_FOR_INIT_1(m, __COUNTER__, __VA_ARGS__) 149 150 } // namespace tensorflow 151 152 #endif // TENSORFLOW_CORE_FRAMEWORK_REGISTRATION_REGISTRATION_H_ 153