Builder class passed to the REGISTER_KERNEL_BUILDER() macro.
| 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 | template <typename T> |
| 47 | KernelDefBuilder& AttrConstraint(const char* attr_name, |
| 48 | gtl::ArraySlice<T> allowed); |
| 49 | |
| 50 | // Like AttrConstraint above but supports just a single value. |
| 51 | template <typename T> |
| 52 | KernelDefBuilder& AttrConstraint(const char* attr_name, T allowed); |
| 53 | |
| 54 | // Specify that this kernel supports a limited set of values for a |
| 55 | // particular type or list(type) attr (a further restriction than |
| 56 | // what the Op allows). |
| 57 | // Returns *this. |
| 58 | KernelDefBuilder& TypeConstraint(const char* attr_name, |
| 59 | gtl::ArraySlice<DataType> allowed); |
| 60 | |
| 61 | // Like TypeConstraint but supports just a single type. |
| 62 | KernelDefBuilder& TypeConstraint(const char* attr_name, DataType allowed); |
| 63 | |
| 64 | // Like TypeConstraint, but (a) gets the type from a template parameter |
| 65 | // and (b) only supports a constraint to a single type. |
| 66 | template <class T> |
| 67 | KernelDefBuilder& TypeConstraint(const char* attr_name); |
| 68 | // TODO(josh11b): Support other types of attr constraints as needed. |
| 69 | |
| 70 | // Specify that this kernel requires/provides an input/output arg |
| 71 | // in host memory (instead of the default, device memory). |
| 72 | // Returns *this. |
| 73 | KernelDefBuilder& HostMemory(const char* arg_name); |
| 74 | |
| 75 | // Specify that this kernel requires a particular value for the |
| 76 | // "_kernel" attr. May only be specified once. Returns *this. |
| 77 | KernelDefBuilder& Label(const char* label); |
| 78 | |
| 79 | // Specify a priority number for this kernel. |
| 80 | KernelDefBuilder& Priority(int32 priority); |
| 81 | |
| 82 | // Returns a pointer to a KernelDef with fields set based on the |
| 83 | // above calls to this instance. |
| 84 | // Caller takes ownership of the result. |
| 85 | const KernelDef* Build(); |
| 86 | |
| 87 | private: |