-
Notifications
You must be signed in to change notification settings - Fork 74k
/
register_types.h
172 lines (138 loc) · 6.14 KB
/
register_types.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
/* Copyright 2015 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_FRAMEWORK_REGISTER_TYPES_H_
#define TENSORFLOW_FRAMEWORK_REGISTER_TYPES_H_
// This file is used by cuda code and must remain compilable by nvcc.
#include "tensorflow/core/framework/numeric_types.h"
#include "tensorflow/core/platform/types.h"
// Two sets of macros:
// - TF_CALL_float, TF_CALL_double, etc. which call the given macro with
// the type name as the only parameter - except on platforms for which
// the type should not be included.
// - Macros to apply another macro to lists of supported types. These also call
// into TF_CALL_float, TF_CALL_double, etc. so they filter by target platform
// as well.
// If you change the lists of types, please also update the list in types.cc.
//
// See example uses of these macros in core/ops.
//
//
// Each of these TF_CALL_XXX_TYPES(m) macros invokes the macro "m" multiple
// times by passing each invocation a data type supported by TensorFlow.
//
// The different variations pass different subsets of the types.
// TF_CALL_ALL_TYPES(m) applied "m" to all types supported by TensorFlow.
// The set of types depends on the compilation platform.
//.
// This can be used to register a different template instantiation of
// an OpKernel for different signatures, e.g.:
/*
#define REGISTER_PARTITION(type) \
REGISTER_TF_OP_KERNEL("partition", DEVICE_CPU, #type ", int32", \
PartitionOp<type>);
TF_CALL_ALL_TYPES(REGISTER_PARTITION)
#undef REGISTER_PARTITION
*/
#if !defined(IS_MOBILE_PLATFORM) || defined(SUPPORT_SELECTIVE_REGISTRATION)
// All types are supported, so all macros are invoked.
//
// Note: macros are defined in same order as types in types.proto, for
// readability.
#define TF_CALL_float(m) m(float)
#define TF_CALL_double(m) m(double)
#define TF_CALL_int32(m) m(::tensorflow::int32)
#define TF_CALL_uint8(m) m(::tensorflow::uint8)
#define TF_CALL_int16(m) m(::tensorflow::int16)
#define TF_CALL_int8(m) m(::tensorflow::int8)
#define TF_CALL_string(m) m(string)
#define TF_CALL_complex64(m) m(::tensorflow::complex64)
#define TF_CALL_int64(m) m(::tensorflow::int64)
#define TF_CALL_bool(m) m(bool)
#define TF_CALL_qint8(m) m(::tensorflow::qint8)
#define TF_CALL_quint8(m) m(::tensorflow::quint8)
#define TF_CALL_qint32(m) m(::tensorflow::qint32)
#define TF_CALL_bfloat16(m) m(::tensorflow::bfloat16)
#define TF_CALL_qint16(m) m(::tensorflow::qint16)
#define TF_CALL_quint16(m) m(::tensorflow::quint16)
#define TF_CALL_uint16(m) m(::tensorflow::uint16)
#define TF_CALL_complex128(m) m(::tensorflow::complex128)
#define TF_CALL_half(m) m(Eigen::half)
#elif defined(__ANDROID_TYPES_FULL__)
// Only half, float, int32, int64, and quantized types are supported.
#define TF_CALL_float(m) m(float)
#define TF_CALL_double(m)
#define TF_CALL_int32(m) m(::tensorflow::int32)
#define TF_CALL_uint8(m)
#define TF_CALL_int16(m)
#define TF_CALL_int8(m)
#define TF_CALL_string(m)
#define TF_CALL_complex64(m)
#define TF_CALL_int64(m) m(::tensorflow::int64)
#define TF_CALL_bool(m)
#define TF_CALL_qint8(m) m(::tensorflow::qint8)
#define TF_CALL_quint8(m) m(::tensorflow::quint8)
#define TF_CALL_qint32(m) m(::tensorflow::qint32)
#define TF_CALL_bfloat16(m)
#define TF_CALL_qint16(m) m(::tensorflow::qint16)
#define TF_CALL_quint16(m) m(::tensorflow::quint16)
#define TF_CALL_uint16(m)
#define TF_CALL_complex128(m)
#define TF_CALL_half(m) m(Eigen::half)
#else // defined(IS_MOBILE_PLATFORM) && !defined(__ANDROID_TYPES_FULL__)
// Only float and int32 are supported.
#define TF_CALL_float(m) m(float)
#define TF_CALL_double(m)
#define TF_CALL_int32(m) m(::tensorflow::int32)
#define TF_CALL_uint8(m)
#define TF_CALL_int16(m)
#define TF_CALL_int8(m)
#define TF_CALL_string(m)
#define TF_CALL_complex64(m)
#define TF_CALL_int64(m)
#define TF_CALL_bool(m)
#define TF_CALL_qint8(m)
#define TF_CALL_quint8(m)
#define TF_CALL_qint32(m)
#define TF_CALL_bfloat16(m)
#define TF_CALL_qint16(m)
#define TF_CALL_quint16(m)
#define TF_CALL_uint16(m)
#define TF_CALL_complex128(m)
#define TF_CALL_half(m)
#endif // defined(IS_MOBILE_PLATFORM) - end of TF_CALL_type defines
// Defines for sets of types.
#define TF_CALL_INTEGRAL_TYPES(m) \
TF_CALL_int64(m) TF_CALL_int32(m) TF_CALL_uint16(m) TF_CALL_int16(m) \
TF_CALL_uint8(m) TF_CALL_int8(m)
#define TF_CALL_REAL_NUMBER_TYPES(m) \
TF_CALL_INTEGRAL_TYPES(m) TF_CALL_half(m) TF_CALL_float(m) TF_CALL_double(m)
#define TF_CALL_REAL_NUMBER_TYPES_NO_INT32(m) \
TF_CALL_half(m) TF_CALL_float(m) TF_CALL_double(m) TF_CALL_int64(m) \
TF_CALL_uint16(m) TF_CALL_int16(m) TF_CALL_uint8(m) TF_CALL_int8(m)
// Call "m" for all number types, including complex64 and complex128.
#define TF_CALL_NUMBER_TYPES(m) \
TF_CALL_REAL_NUMBER_TYPES(m) TF_CALL_complex64(m) TF_CALL_complex128(m)
#define TF_CALL_NUMBER_TYPES_NO_INT32(m) \
TF_CALL_REAL_NUMBER_TYPES_NO_INT32(m) \
TF_CALL_complex64(m) TF_CALL_complex128(m)
#define TF_CALL_POD_TYPES(m) TF_CALL_NUMBER_TYPES(m) TF_CALL_bool(m)
// Call "m" on all types.
#define TF_CALL_ALL_TYPES(m) TF_CALL_POD_TYPES(m) TF_CALL_string(m)
// Call "m" on all types supported on GPU.
#define TF_CALL_GPU_NUMBER_TYPES(m) \
TF_CALL_half(m) TF_CALL_float(m) TF_CALL_double(m)
#define TF_CALL_GPU_NUMBER_TYPES_NO_HALF(m) TF_CALL_float(m) TF_CALL_double(m)
// Call "m" on all quantized types.
// TODO(cwhipkey): include TF_CALL_qint16(m) TF_CALL_quint16(m)
#define TF_CALL_QUANTIZED_TYPES(m) \
TF_CALL_qint8(m) TF_CALL_quint8(m) TF_CALL_qint32(m)
#endif // TENSORFLOW_FRAMEWORK_REGISTER_TYPES_H_