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

Add some annotations for op / kernel registrations - visible in the Clang AST /

Browse Source 
LLVM IR

This change adds a macro for [[clang::annotate]], and some initial uses in the
op and kernel registration macros. The intent is to make it easier to identify
and work with these registrations in AST / IR-level tooling (in particular,
with less sensitivity to particular implementation).

PiperOrigin-RevId: 327529042
Change-Id: I50b16e7ac18250573a2dcc0c48262fdfba7752a1
This commit is contained in:
Randy Dodgen 2020-08-19 16:24:04 -07:00 committed by TensorFlower Gardener
parent add3c56b2e
commit 1087d48004
3 changed files with 25 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/framework/op.h
View File

@ -313,6 +313,7 @@ struct OpDefBuilderReceiver {
#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)
#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)
#define REGISTER_OP_UNIQ(ctr, name) \
TF_ATTRIBUTE_ANNOTATE("tf:op") \
static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr \
TF_ATTRIBUTE_UNUSED = \
::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( \
@ -326,6 +327,8 @@ struct OpDefBuilderReceiver {
#define REGISTER_SYSTEM_OP_UNIQ_HELPER(ctr, name) \
REGISTER_SYSTEM_OP_UNIQ(ctr, name)
#define REGISTER_SYSTEM_OP_UNIQ(ctr, name) \
TF_ATTRIBUTE_ANNOTATE("tf:op") \
TF_ATTRIBUTE_ANNOTATE("tf:op:system") \
static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr \
TF_ATTRIBUTE_UNUSED = \
::tensorflow::register_op::OpDefBuilderWrapper<true>(name)

3
tensorflow/core/framework/op_kernel.h
View File

@ -1457,6 +1457,7 @@ class Name : public KernelDefBuilder {
#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...) \
constexpr bool should_register_##ctr##__flag = \
SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__); \
TF_ATTRIBUTE_ANNOTATE("tf:kernel") \
static ::tensorflow::kernel_factory::OpKernelRegistrar \
registrar__body__##ctr##__object( \
should_register_##ctr##__flag \
@ -1479,6 +1480,8 @@ class Name : public KernelDefBuilder {
REGISTER_SYSTEM_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__)
#define REGISTER_SYSTEM_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...) \
TF_ATTRIBUTE_ANNOTATE("tf:kernel") \
TF_ATTRIBUTE_ANNOTATE("tf:kernel:system") \
static ::tensorflow::kernel_factory::OpKernelRegistrar \
registrar__body__##ctr##__object( \
::tensorflow::register_kernel::system::kernel_builder.Build(), \

19
tensorflow/core/platform/macros.h
View File

@ -74,6 +74,25 @@ limitations under the License.
#define TF_HAS_BUILTIN(x) 0
#endif
// C++11-style attributes (N2761)
#if defined(__has_cpp_attribute)
// Safely checks if an attribute is supported. Equivalent to
// ABSL_HAVE_CPP_ATTRIBUTE.
#define TF_HAS_CPP_ATTRIBUTE(n) __has_cpp_attribute(n)
#else
#define TF_HAS_CPP_ATTRIBUTE(n) 0
#endif
// [[clang::annotate("x")]] allows attaching custom strings (e.g. "x") to
// declarations (variables, functions, fields, etc.) for use by tools. They are
// represented in the Clang AST (as AnnotateAttr nodes) and in LLVM IR, but not
// in final output.
#if TF_HAS_CPP_ATTRIBUTE(clang::annotate)
#define TF_ATTRIBUTE_ANNOTATE(str) [[clang::annotate(str)]]
#else
#define TF_ATTRIBUTE_ANNOTATE(str)
#endif
// Compilers can be told that a certain branch is not likely to be taken
// (for instance, a CHECK failure), and use that information in static
// analysis. Giving it this information can help it optimize for the