1 /* Copyright 2015 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 #ifndef TENSORFLOW_FRAMEWORK_KERNEL_DEF_BUILDER_H_ 17 #define TENSORFLOW_FRAMEWORK_KERNEL_DEF_BUILDER_H_ 18 19 #include "tensorflow/core/framework/types.h" 20 #include "tensorflow/core/lib/gtl/array_slice.h" 21 #include "tensorflow/core/platform/macros.h" 22 #include "tensorflow/core/platform/types.h" 23 24 namespace tensorflow { 25 26 // Forward declare proto so that kernels don't need to depend on it 27 class KernelDef; 28 29 // Builder class passed to the REGISTER_KERNEL_BUILDER() macro. 30 class KernelDefBuilder { 31 public: 32 // Starts with just the name field set. 33 // Caller MUST call Build() and take ownership of the result. 34 explicit KernelDefBuilder(const char* op_name); 35 ~KernelDefBuilder(); 36 37 // Required: specify the type of device this kernel supports. 38 // Returns *this. 39 KernelDefBuilder& Device(const char* device_type); 40 // KernelDefBuilder& Device(DeviceType device_type); 41 42 // Specify that this kernel supports a limited set of values for a 43 // particular type or list(type) attr (a further restriction than 44 // what the Op allows). 45 // Returns *this. 46 KernelDefBuilder& TypeConstraint(const char* attr_name, 47 gtl::ArraySlice<DataType> allowed); 48 49 // Like TypeConstraint but supports just a single type. 50 KernelDefBuilder& TypeConstraint(const char* attr_name, DataType allowed); 51 52 // Like TypeConstraint, but (a) gets the type from a template parameter 53 // and (b) only supports a constraint to a single type. 54 template <class T> 55 KernelDefBuilder& TypeConstraint(const char* attr_name); 56 // TODO(josh11b): Support other types of attr constraints as needed. 57 58 // Specify that this kernel requires/provides an input/output arg 59 // in host memory (instead of the default, device memory). 60 // Returns *this. 61 KernelDefBuilder& HostMemory(const char* arg_name); 62 63 // Specify that this kernel requires a particular value for the 64 // "_kernel" attr. May only be specified once. Returns *this. 65 KernelDefBuilder& Label(const char* label); 66 67 // Returns a pointer to a KernelDef with fields set based on the 68 // above calls to this instance. 69 // Caller takes ownership of the result. 70 const KernelDef* Build(); 71 72 private: 73 KernelDef* kernel_def_; 74 75 TF_DISALLOW_COPY_AND_ASSIGN(KernelDefBuilder); 76 }; 77 78 // IMPLEMENTATION 79 80 template <class T> 81 KernelDefBuilder& KernelDefBuilder::TypeConstraint(const char* attr_name) { 82 return this->TypeConstraint(attr_name, DataTypeToEnum<T>::v()); 83 } 84 85 } // namespace tensorflow 86 87 #endif // TENSORFLOW_FRAMEWORK_KERNEL_DEF_BUILDER_H_ 88
