Caffe2 - C++ API
A deep learning, cross platform ML framework
operator_schema.h
1 
17 #ifndef CAFFE2_CORE_OPERATOR_SCHEMA_H_
18 #define CAFFE2_CORE_OPERATOR_SCHEMA_H_
19 
20 #include <climits>
21 #include <functional>
22 #include <initializer_list>
23 #include <ostream>
24 #include <set>
25 #include <vector>
26 
27 #include "caffe2/core/common.h"
28 #include "caffe2/core/logging.h"
29 #include "caffe2/core/registry.h"
30 #include "caffe2/proto/caffe2.pb.h"
31 
32 namespace caffe2 {
33 
34 // A const value returned by OpSchema::CalculateOutput() if the number of
35 // output cannot be determined.
36 constexpr int kCannotComputeNumOutputs = -1;
37 
53 class OpSchema {
54  public:
55  OpSchema() : file_("unknown"), line_(0) {}
56  OpSchema(const string& file, const int line) : file_(file), line_(line) {}
57 
61  inline const string& file() const {
62  return file_;
63  }
64 
68  inline int line() const {
69  return line_;
70  }
71 
75  inline const char* doc() const {
76  return doc_.empty() ? nullptr : doc_.c_str();
77  }
78 
83  bool Verify(const OperatorDef& def) const;
84 
85  // Functions to set the property of the operator schemas.
86  // Sets the number of inputs, either a fixed number or a min and a max.
87 
91  OpSchema& NumInputs(int n);
95  OpSchema& NumInputs(int min, int max);
99  OpSchema& NumInputs(set<int> allowed_input_nums);
103  OpSchema& NumInputs(std::function<bool(int)> func);
104 
105  // Sets the number of outputs, either a fixed number, a min and a max,
106  // or a function that takes in the input number and produces an output
107  // number. Use only one function in the set below.
111  OpSchema& NumOutputs(int n);
115  OpSchema& NumOutputs(int min, int max);
119  OpSchema& NumOutputs(set<int> allowed_output_nums);
123  OpSchema& NumOutputs(std::function<bool(int)> func);
124 
129  OpSchema& NumInputsOutputs(std::function<bool(int, int)> func);
130 
131  // Set the function that can calculate the number of output based on the
132  // number of input. Use only one function in the set below.
136  OpSchema& OutputCalculator(std::function<int(int)> calc);
141 
142  // Sets the rule to allow optional in-place operation.
143  OpSchema& AllowInplace(std::function<bool(int, int)> inplace);
144  OpSchema& AllowInplace(set<std::pair<int, int>> inplace);
145  OpSchema& AllowOneToOneInplace();
146  // Sets the rule to enforce in-place opeartion.
147  OpSchema& EnforceInplace(std::function<bool(int, int)> inplace);
148  OpSchema& EnforceInplace(set<std::pair<int, int>> inplace);
149  OpSchema& EnforceOneToOneInplace();
150 
151  // Functions to deal with type and shape inference. Basically, this registers
152  // a function that takes in an OperatorDef and a series of input type and
153  // shape specified by TensorProto objects (whose data fields are empty), and
154  // produces a series of output type and shape.
155  typedef std::function<
156  vector<TensorShape>(const OperatorDef&, const vector<TensorShape>&)>
157  TensorInferenceFunctionType;
162  OpSchema& TensorInferenceFunction(TensorInferenceFunctionType function);
168  OpSchema& IdenticalTypeAndShapeOfInput(int idx);
169  OpSchema& IdenticalTypeAndShapeOfInputDim(int idx, int dim);
170  OpSchema& ScalarType(::caffe2::TensorProto_DataType dt);
171 
176  inline vector<TensorShape> InferTensor(
177  const OperatorDef& def,
178  const vector<TensorShape> input_type_shape) const {
179  return tensor_inference_function_(def, input_type_shape);
180  }
181 
182  /*
183  * @brief A struct to store various cost information about
184  * an operator such as FLOPs, total memory use and parameters.
185  */
186  struct Cost {
187  uint64_t flops; // Floating point operations.
188  uint64_t bytes_moved; // Total memory used.
189  uint64_t params_bytes; // Memory footprint of parameters
190  };
196  typedef std::function<
197  struct Cost(const OperatorDef&, const vector<TensorShape>&)>
199 
204 
205 #if 0 // def _MSC_VER
206 
209  template <typename T,
210  typename = std::enable_if<
211  std::is_same<CostInferenceFunctionType&&, T>:value
212  >:type>
213  inline OpSchema& CostInferenceFunction(T func) {
214  // Note: This is here in order to resolve an MSVC compiler issue: it
215  // does not automatically convert a function pointer to a std::function,
216  // and needs an explicit conversion.
218  }
219 #endif // _MSC_VER
220 
221  bool HasCostInferenceFunction() const {
222  return !!cost_inference_function_;
223  }
224 
225  inline struct Cost InferCost(
226  const OperatorDef& def,
227  const vector<TensorShape>& input_tensor_shape) const {
228  CAFFE_ENFORCE(
229  cost_inference_function_, "Cost inference function not defined.");
230  return (*cost_inference_function_)(def, input_tensor_shape);
231  }
232 
233  // Functions to do documentation for the operator schema.
234  OpSchema& SetDoc(const string& doc);
235 
236  struct Argument {
237  Argument(const char* name, const char* description, bool required)
238  : name_{name}, description_{description}, required_{required} {}
239 
240  const char* name() const {
241  return name_;
242  }
243 
244  const char* description() const {
245  return description_;
246  }
247 
248  bool is_required() const {
249  return required_;
250  }
251 
252  private:
253  const char* name_;
254  const char* description_;
255  const bool required_;
256  };
257 
258  OpSchema&
259  Arg(const char* name, const char* description, bool required = false);
260 
261 #define DECLARE_STANDARD_ARG(name, str) \
262  CAFFE2_API static const char* Arg_##name; \
263  CAFFE2_API OpSchema& Arg##name(const char* description);
264 
265  DECLARE_STANDARD_ARG(IsTest, is_test)
266 
267 #undef DECLARE_STANDARD_ARG
268 
269  OpSchema& Input(const int n, const char* name, const char* description);
270  OpSchema& Output(const int n, const char* name, const char* description);
271  // Calls the passed function with `this` as an argument. Useful for
272  // adding docs for temlated/macro ops.
273  OpSchema& FillUsing(std::function<void(OpSchema&)> populator);
274 
275  // Remove from documentation
276  OpSchema& Private();
277 
278  // This op can pass data across devices
279  OpSchema& InputsCanCrossDevices();
280 
285  int CalculateOutput(int num_input) const;
286 
287  int min_input() const {
288  return min_input_;
289  }
290 
291  int max_input() const {
292  return max_input_;
293  }
294 
295  int min_output() const {
296  return min_output_;
297  }
298 
299  int max_output() const {
300  return max_output_;
301  }
302 
303  bool num_inputs_allowed(int x) const {
304  return num_inputs_allowed_(x);
305  }
306 
307  bool num_outputs_allowed(int x) const {
308  return num_outputs_allowed_(x);
309  }
310 
311  bool num_inputs_outputs_allowed(int x, int y) const {
312  return num_inputs_outputs_allowed_(x, y);
313  }
314 
315  int inf() const {
316  return std::numeric_limits<int>::max();
317  }
318 
319  friend std::ostream& operator<<(std::ostream& out, const OpSchema& schema);
320 
321  const std::vector<Argument>& args() const {
322  return args_;
323  }
324 
325  const std::vector<std::pair<const char*, const char*>>& input_desc() const {
326  return input_desc_;
327  }
328  const std::vector<std::pair<const char*, const char*>>& output_desc() const {
329  return output_desc_;
330  }
331  bool private_op() {
332  return private_;
333  }
334  bool inputs_can_cross_devices() const {
335  return inputs_can_cross_devices_;
336  }
337 
341  using DeviceInferenceFunctionType = std::function<
342  std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>(
343  const OperatorDef& def)>;
344 
345  OpSchema& DeviceInferenceFunction(DeviceInferenceFunctionType function);
346 
350  inline std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
351  InferDevice(const OperatorDef& def) const {
352  return device_inference_function_(def);
353  }
354 
355  private:
356  string file_;
357  string doc_;
358  std::vector<Argument> args_{};
359  std::vector<std::pair<const char*, const char*>> input_desc_{};
360  std::vector<std::pair<const char*, const char*>> output_desc_{};
361  int line_ = 0;
362  int min_input_ = 0;
363  int max_input_ = std::numeric_limits<int>::max();
364  int min_output_ = 0;
365  int max_output_ = std::numeric_limits<int>::max();
366  bool private_ = false;
367  bool inputs_can_cross_devices_ = false;
368  std::function<bool(int)> num_inputs_allowed_ = [](int) { return true; };
369  std::function<bool(int)> num_outputs_allowed_ = [](int) { return true; };
370  std::function<bool(int, int)> num_inputs_outputs_allowed_ = [](int, int) {
371  return true;
372  };
373  std::function<int(int)> calculate_output_;
374  // In default, any in-place operation is neither allowed nor enforced.
375  std::function<bool(int, int)> inplace_allowed_ = [](int, int) {
376  return false;
377  };
378  std::function<bool(int, int)> inplace_enforced_ = [](int, int) {
379  return false;
380  };
381  TensorInferenceFunctionType tensor_inference_function_ =
382  [](const OperatorDef& def, const vector<TensorShape>&) {
383  vector<TensorShape> out;
384  for (int i = 0; i < def.output_size(); i++) {
385  TensorShape ts;
386  ts.set_unknown_shape(true);
387  out.push_back(ts);
388  }
389  return out;
390  };
391  std::unique_ptr<CostInferenceFunctionType> cost_inference_function_ = nullptr;
392  DeviceInferenceFunctionType device_inference_function_ =
393  [](const OperatorDef& def) {
394  auto op_device =
395  def.has_device_option() ? def.device_option() : DeviceOption();
396  vector<DeviceOption> in_dev(def.input_size(), op_device);
397  vector<DeviceOption> out_dev(def.output_size(), op_device);
398  return std::make_pair(in_dev, out_dev);
399  };
400 };
401 
406  public:
407  static OpSchema&
408  NewSchema(const string& key, const string& file, const int line) {
409  auto& m = map();
410  if (m.count(key)) {
411  const auto& schema = m[key];
412  std::ios_base::Init init;
413  std::cerr << "Trying to register schema with name " << key
414  << " from file " << file << " line " << line
415  << ", but it is already registered from file " << schema.file()
416  << " line " << schema.line();
417  abort();
418  }
419  m.emplace(std::make_pair(key, OpSchema(file, line)));
420  return m[key];
421  }
422 
423  static const OpSchema* Schema(const string& key) {
424  auto& m = map();
425  if (m.count(key)) {
426  return &m[key];
427  } else {
428  return nullptr;
429  }
430  }
431 
432  private:
433  // OpSchemaRegistry should not need to be instantiated.
434  OpSchemaRegistry() = delete;
435 
446  static CaffeMap<string, OpSchema>& map();
447 };
448 
449 // Helper function for creating simple tensorproto with dimension and type
450 template <typename T_I = int>
451 inline TensorShape CreateTensorShape(
452  vector<T_I> dims,
453  ::caffe2::TensorProto_DataType dt) {
454  TensorShape ts;
455  for (int d : dims) {
456  ts.add_dims(d);
457  }
458  ts.set_data_type(dt);
459  return ts;
460 }
461 
462 // Helper function
463 inline vector<TIndex> GetDimsVector(const TensorShape& shape) {
464  vector<TIndex> dims;
465  for (auto d : shape.dims()) {
466  dims.push_back(d);
467  }
468  return dims;
469 }
470 
471 // Helper function for infer op inputs and outputs device information.
472 inline std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
473 InferOpInputOutputDevice(const OperatorDef& op) {
474  auto op_schema = OpSchemaRegistry::Schema(op.type());
475  CAFFE_ENFORCE(
476  op_schema, "Device inference failed. No schema for: ", op.type());
477  // TODO(wyiming) : add try catch here.
478  return op_schema->InferDevice(op);
479 }
480 
481 template <uint64_t OpsPerPoint>
482 OpSchema::Cost PointwiseCostInference(
483  const OperatorDef& /* unused */,
484  const vector<TensorShape>& inputs) {
485  struct OpSchema::Cost c;
486  const TensorShape X = inputs[0];
487  uint64_t size = 1;
488 
489  for (auto i = 0; i < X.dims().size(); ++i) {
490  size *= X.dims(i);
491  }
492 
493  c.flops = size * OpsPerPoint;
494  c.bytes_moved = size;
495  return c;
496 }
497 
498 } // namespace caffe2
499 
500 #ifndef CAFFE2_NO_OPERATOR_SCHEMA
501 
502 #define OPERATOR_SCHEMA(name) \
503  void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
504  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(name) = \
505  &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
506 #define OPERATOR_SCHEMA_STR(name) \
507  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(schema_registration) = \
508  &OpSchemaRegistry::NewSchema(name, __FILE__, __LINE__)
509 
510 #else // CAFFE2_NO_OPERATOR_SCHEMA
511 
512 #define OPERATOR_SCHEMA(name) \
513  void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
514  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(name) = \
515  1 ? nullptr : &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
516 #define OPERATOR_SCHEMA_STR(name) \
517  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(schema_registration) = \
518  1 ? nullptr : &OpSchemaRegistry::NewSchema(name, __FILE__, __LINE__)
519 
520 #endif // CAFFE2_NO_OPERATOR_SCHEMA
521 
522 #endif // CAFFE2_CORE_OPERATOR_SCHEMA_H_
std::function< std::pair< std::vector< DeviceOption >, std::vector< DeviceOption >>(const OperatorDef &def)> DeviceInferenceFunctionType
Returns the required device location of inputs and outputs.
OpSchema & NumInputs(int n)
A single input.
A class to record the schema of an op.
bool Verify(const OperatorDef &def) const
Verifies if an operator definition protobuf matches the pattern specified in the schema.
A registry to hold all the operator schemas.
int line() const
Returns the line in file that the op schema is registered from.
const char * doc() const
Returns the docstring of the op schema.
vector< TensorShape > InferTensor(const OperatorDef &def, const vector< TensorShape > input_type_shape) const
A function to allow one to infer the type and shape from the op schema.
OpSchema & OutputCalculator(std::function< int(int)> calc)
Set the output calculator to a user-defined function.
OpSchema & IdenticalTypeAndShape()
Sets the tensor inference function to produce the same output as the input.
OpSchema & SameNumberOfOutput()
Set the number of outputs to be the same as the number of inputs.
std::pair< std::vector< DeviceOption >, std::vector< DeviceOption > > InferDevice(const OperatorDef &def) const
Infer required device location of an op&#39;s inputs and outputs.
Copyright (c) 2016-present, Facebook, Inc.
OpSchema & CostInferenceFunction(CostInferenceFunctionType function)
Register the Cost inference function.
OpSchema & NumInputsOutputs(std::function< bool(int, int)> func)
Relationship between inputs and outputs is checked with a specified function.
const string & file() const
Returns the file that the op schema is registered from.
OpSchema & TensorInferenceFunction(TensorInferenceFunctionType function)
Sets the tensor inference function, which is a std::function object defined in operator_schema.h.
int CalculateOutput(int num_input) const
A function to allow one to get the number of outputs based on the number of inputs, if this schema supports it.
OpSchema & NumOutputs(int n)
A single output.
std::function< struct Cost(const OperatorDef &, const vector< TensorShape > &)> CostInferenceFunctionType
Registers a function that takes in an OperatorDef and a series of input shapes and returns the total ...