Caffe2 - C++ API
A deep learning, cross platform ML framework
operator_schema.h
1 #ifndef CAFFE2_CORE_OPERATOR_SCHEMA_H_
2 #define CAFFE2_CORE_OPERATOR_SCHEMA_H_
3 
4 #include <climits>
5 #include <functional>
6 #include <initializer_list>
7 #include <ostream>
8 #include <set>
9 #include <vector>
10 #include <unordered_map>
11 
12 #include "c10/util/Registry.h"
13 #include "caffe2/core/common.h"
14 #include "caffe2/core/logging.h"
15 #include "caffe2/proto/caffe2_pb.h"
16 #include "caffe2/utils/filler.h"
17 
18 namespace caffe2 {
19 
20 // A const value returned by OpSchema::CalculateOutput() if the number of
21 // output cannot be determined.
22 constexpr int kCannotComputeNumOutputs = -1;
23 
39 class CAFFE2_API OpSchema {
40  public:
41  OpSchema() : type_("unknown"), file_("unknown"), line_(0) {}
42  OpSchema(const string& type, const string& file, const int line)
43  : type_(type), file_(file), line_(line) {}
44 
48  inline const string& file() const {
49  return file_;
50  }
51 
55  inline int line() const {
56  return line_;
57  }
58 
62  inline const char* doc() const {
63  return doc_.empty() ? nullptr : doc_.c_str();
64  }
65 
70  bool Verify(const OperatorDef& def) const;
71 
72  // Functions to set the property of the operator schemas.
73  // Sets the number of inputs, either a fixed number or a min and a max.
74 
78  OpSchema& NumInputs(int n);
82  OpSchema& NumInputs(int min, int max);
86  OpSchema& NumInputs(set<int> allowed_input_nums);
90  OpSchema& NumInputs(std::function<bool(int)> func);
91 
92  // Sets the number of outputs, either a fixed number, a min and a max,
93  // or a function that takes in the input number and produces an output
94  // number. Use only one function in the set below.
98  OpSchema& NumOutputs(int n);
102  OpSchema& NumOutputs(int min, int max);
106  OpSchema& NumOutputs(set<int> allowed_output_nums);
110  OpSchema& NumOutputs(std::function<bool(int)> func);
111 
116  OpSchema& NumInputsOutputs(std::function<bool(int, int)> func);
117 
118  // Set the function that can calculate the number of output based on the
119  // number of input. Use only one function in the set below.
123  OpSchema& OutputCalculator(std::function<int(int)> calc);
127  OpSchema& SameNumberOfOutput();
128 
129  // Sets the rule to allow optional in-place operation.
130  OpSchema& AllowInplace(std::function<bool(int, int)> inplace);
131  OpSchema& AllowInplace(set<std::pair<int, int>> inplace);
132  OpSchema& AllowOneToOneInplace();
133  // Sets the rule to enforce in-place opeartion.
134  OpSchema& EnforceInplace(std::function<bool(int, int)> inplace);
135  OpSchema& EnforceInplace(set<std::pair<int, int>> inplace);
136  OpSchema& EnforceOneToOneInplace();
137 
138  // Functions to deal with type and shape inference. Basically, this registers
139  // a function that takes in an OperatorDef and a series of input type and
140  // shape specified by TensorProto objects (whose data fields are empty), and
141  // produces a series of output type and shape.
142  typedef std::function<
143  vector<TensorShape>(const OperatorDef&, const vector<TensorShape>&)>
144  TensorInferenceFunctionType;
145 
150  OpSchema& TensorInferenceFunction(TensorInferenceFunctionType function);
151 
157  static TensorInferenceFunctionType NeedsAllInputShapes(
158  TensorInferenceFunctionType f);
159 
163  OpSchema& InheritOnnxSchema(const std::string& onnx_schema_name);
164 
169  return InheritOnnxSchema(type_);
170  }
171 
176  OpSchema& IdenticalTypeAndShape();
177  OpSchema& IdenticalTypeAndShapeOfInput(int idx);
178  OpSchema& IdenticalTypeAndShapeOfInputDim(int idx, int dim);
179  OpSchema& IdenticalTypeAndShapeOfMultipleInputs(const vector<int>& indices);
180  OpSchema& ScalarType(::caffe2::TensorProto_DataType dt);
181 
186  inline vector<TensorShape> InferTensor(
187  const OperatorDef& def,
188  const vector<TensorShape>& input_type_shape) const {
189  return tensor_inference_function_(def, input_type_shape);
190  }
191 
192  /*
193  * @brief A struct to store various cost information about
194  * an operator such as FLOPs, total memory use and parameters.
195  */
196  struct Cost {
197  uint64_t flops{0}; // Floating point operations.
198  uint64_t bytes_read{0}; // Total memory read.
199  uint64_t bytes_written{0}; // Total memory written.
200  uint64_t params_bytes{0}; // Memory read for parameters.
201  };
207  typedef std::function<
208  struct Cost(const OperatorDef&, const vector<TensorShape>&)>
210 
214  OpSchema& CostInferenceFunction(CostInferenceFunctionType function);
215 
216 #if 0 // def _MSC_VER
217 
220  template <typename T,
221  typename = std::enable_if<
222  std::is_same<CostInferenceFunctionType&&, T>:value
223  >:type>
224  inline OpSchema& CostInferenceFunction(T func) {
225  // Note: This is here in order to resolve an MSVC compiler issue: it
226  // does not automatically convert a function pointer to a std::function,
227  // and needs an explicit conversion.
228  return CostInferenceFunction(CostInferenceFunctionType(func));
229  }
230 #endif // _MSC_VER
231 
232  bool HasCostInferenceFunction() const {
233  return !!cost_inference_function_;
234  }
235 
236  inline struct Cost InferCost(
237  const OperatorDef& def,
238  const vector<TensorShape>& input_tensor_shape) const {
239  CAFFE_ENFORCE(
240  cost_inference_function_, "Cost inference function not defined.");
241  return (*cost_inference_function_)(def, input_tensor_shape);
242  }
243 
244  // Functions to do documentation for the operator schema.
245  OpSchema& SetDoc(const string& doc);
246 
247  struct Argument {
248  Argument(const char* name, const char* description, bool required)
249  : name_{name}, description_{description}, required_{required} {}
250 
251  const char* name() const {
252  return name_;
253  }
254 
255  const char* description() const {
256  return description_;
257  }
258 
259  bool is_required() const {
260  return required_;
261  }
262 
263  private:
264  const char* name_;
265  const char* description_;
266  const bool required_;
267  };
268 
269  OpSchema&
270  Arg(const char* name, const char* description, bool required = false);
271 
272 #define DECLARE_STANDARD_ARG(name, str) \
273  static const char* Arg_##name; \
274  OpSchema& Arg##name(const char* description);
275 
276  DECLARE_STANDARD_ARG(IsTest, is_test)
277 
278 #undef DECLARE_STANDARD_ARG
279 
280  OpSchema& Input(const int n, const char* name, const char* description);
281  OpSchema& Output(const int n, const char* name, const char* description);
282  // Calls the passed function with `this` as an argument. Useful for
283  // adding docs for templated/macro ops.
284  OpSchema& FillUsing(std::function<void(OpSchema&)> populator);
285 
286  // Remove from documentation
287  OpSchema& Private();
288 
289  // This op can pass data across devices
290  OpSchema& InputsCanCrossDevices();
291 
296  int CalculateOutput(int num_input) const;
297 
298  const std::string& onnx_schema() const {
299  return onnx_schema_;
300  }
301 
302  int min_input() const {
303  return min_input_;
304  }
305 
306  int max_input() const {
307  return max_input_;
308  }
309 
310  int min_output() const {
311  return min_output_;
312  }
313 
314  int max_output() const {
315  return max_output_;
316  }
317 
318  bool num_inputs_allowed(int x) const {
319  return num_inputs_allowed_(x);
320  }
321 
322  bool num_outputs_allowed(int x) const {
323  return num_outputs_allowed_(x);
324  }
325 
326  bool num_inputs_outputs_allowed(int x, int y) const {
327  return num_inputs_outputs_allowed_(x, y);
328  }
329 
330  int inf() const {
331  return std::numeric_limits<int>::max();
332  }
333 
334  bool inplace_enforced(int x, int y) const {
335  return inplace_enforced_(x, y);
336  }
337 
338  CAFFE2_API friend std::ostream& operator<<(std::ostream& out, const OpSchema& schema);
339 
340  const std::vector<Argument>& args() const {
341  return args_;
342  }
343 
344  const std::vector<std::pair<const char*, const char*>>& input_desc() const {
345  return input_desc_;
346  }
347  const std::vector<std::pair<const char*, const char*>>& output_desc() const {
348  return output_desc_;
349  }
350  bool private_op() {
351  return private_;
352  }
353  bool inputs_can_cross_devices() const {
354  return inputs_can_cross_devices_;
355  }
356 
360  using DeviceInferenceFunctionType = std::function<
361  std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>(
362  const OperatorDef& def)>;
363 
364  OpSchema& DeviceInferenceFunction(DeviceInferenceFunctionType function);
365 
369  inline std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
370  InferDevice(const OperatorDef& def) const {
371  return device_inference_function_(def);
372  }
373 
374  // The helper is build sparse input with values, keys, weights and lengths;
375  // e.g.:
376  // values = [1, 2, 3, 2, 4, 6, 7, 3, 6]
377  // keys = [0, 1, 4, 0, 1, 2, 5, 1, 2]
378  // weights = [1, 2, 3, 4, 5, 6, 7, 8, 9]
379  // \_____/ \________/ \__/
380  // lengths = [3, 4, 2]
381  OpSchema& WeightedValueKeyLengthInputFillers(
382  size_t value_index,
383  size_t key_index,
384  size_t length_index,
385  size_t weight_index);
386 
387  // The helper is build sparse input with values, keys, weights and lengths;
388  // e.g.:
389  // values = [1, 2, 3, 2, 4, 6, 7, 3, 6]
390  // keys = [0, 1, 4, 0, 1, 2, 5, 1, 2]
391  // \_____/ \________/ \__/
392  // lengths = [3, 4, 2]
393  OpSchema& ValueKeyLengthInputFillers(
394  size_t value_index,
395  size_t key_index,
396  size_t length_index);
397 
398  // The helper is build sparse input with values and lengths; e.g.:
399  // values = [1, 2, 3, 2, 4, 6, 7, 3, 6]
400  // \_____/ \________/ \__/
401  // lengths = [3, 4, 2]
402  OpSchema& ValueLengthInputFillers(size_t value_index, size_t length_index);
403 
404  OpSchema& DisallowInputFillers();
405 
406  std::vector<TensorFiller> InputFillers(
407  const std::vector<std::vector<int64_t>>& shapes) const;
408 
409  private:
410  std::vector<TensorFiller> SupplyDenseFillers(
411  const std::vector<std::vector<int64_t>>& shapes);
412 
413  private:
414  string type_;
415  string file_;
416  string doc_;
417  string onnx_schema_;
418  std::vector<Argument> args_{};
419  std::vector<std::pair<const char*, const char*>> input_desc_{};
420  std::vector<std::pair<const char*, const char*>> output_desc_{};
421  int line_ = 0;
422  int min_input_ = 0;
423  int max_input_ = std::numeric_limits<int>::max();
424  int min_output_ = 0;
425  int max_output_ = std::numeric_limits<int>::max();
426  bool private_ = false;
427  bool inputs_can_cross_devices_ = false;
428  std::function<bool(int)> num_inputs_allowed_ = [](int) { return true; };
429  std::function<bool(int)> num_outputs_allowed_ = [](int) { return true; };
430  std::function<bool(int, int)> num_inputs_outputs_allowed_ = [](int, int) {
431  return true;
432  };
433  std::function<int(int)> calculate_output_;
434  // In default, any in-place operation is neither allowed nor enforced.
435  std::function<bool(int, int)> inplace_allowed_ = [](int, int) {
436  return false;
437  };
438  std::function<bool(int, int)> inplace_enforced_ = [](int, int) {
439  return false;
440  };
441  TensorInferenceFunctionType tensor_inference_function_ =
442  [](const OperatorDef& def, const vector<TensorShape>&) {
443  vector<TensorShape> out;
444  for (int i = 0; i < def.output_size(); i++) {
445  TensorShape ts;
446  ts.set_unknown_shape(true);
447  out.push_back(ts);
448  }
449  return out;
450  };
451  std::unique_ptr<CostInferenceFunctionType> cost_inference_function_ = nullptr;
452  DeviceInferenceFunctionType device_inference_function_ =
453  [](const OperatorDef& def) {
454  auto op_device =
455  def.has_device_option() ? def.device_option() : DeviceOption();
456  vector<DeviceOption> in_dev(def.input_size(), op_device);
457  vector<DeviceOption> out_dev(def.output_size(), op_device);
458  return std::make_pair(in_dev, out_dev);
459  };
460 
461  std::function<std::vector<TensorFiller>(
462  const std::vector<std::vector<int64_t>>&)>
463  filler_supplier_ =
464  [this](const std::vector<std::vector<int64_t>>& shapes) {
465  return SupplyDenseFillers(shapes);
466  };
467 };
468 
472 class CAFFE2_API OpSchemaRegistry {
473  public:
474  static OpSchema&
475  NewSchema(const string& key, const string& file, const int line) {
476  auto& m = map();
477  auto it = m.find(key);
478  if (it != m.end()) {
479  const auto& schema = it->second;
480  std::ios_base::Init init;
481  std::cerr << "Trying to register schema with name " << key
482  << " from file " << file << " line " << line
483  << ", but it is already registered from file " << schema.file()
484  << " line " << schema.line();
485  abort();
486  }
487  m.emplace(std::make_pair(key, OpSchema(key, file, line)));
488  return m[key];
489  }
490 
491  static const OpSchema* Schema(const string& key) {
492  auto& m = map();
493  auto it = m.find(key);
494  if (it != m.end()) {
495  return &it->second;
496  } else {
497  return nullptr;
498  }
499  }
500 
501  private:
502  // OpSchemaRegistry should not need to be instantiated.
503  OpSchemaRegistry() = delete;
504 
515  static CaffeMap<string, OpSchema>& map();
516 };
517 
518 // Helper function for creating simple tensorproto with dimension and type
519 template <typename T_I = int>
520 inline TensorShape CreateTensorShape(
521  vector<T_I> dims,
522  ::caffe2::TensorProto_DataType dt) {
523  TensorShape ts;
524  for (T_I d : dims) {
525  ts.add_dims(d);
526  }
527  ts.set_data_type(dt);
528  return ts;
529 }
530 
531 // Helper function
532 inline vector<int64_t> GetDimsVector(const TensorShape& shape) {
533  vector<int64_t> dims;
534  for (auto d : shape.dims()) {
535  dims.push_back(d);
536  }
537  return dims;
538 }
539 
540 // Helper function
541 inline uint64_t nElemFromDim(const TensorShape& X, int dim = 0) {
542  CAFFE_ENFORCE_GE(dim, 0, "Invalid maximum index specified");
543 
544  uint64_t nElem = 1;
545  for (int i = dim; i < X.dims_size(); ++i) {
546  nElem *= X.dims(i);
547  }
548  return nElem;
549 }
550 
551 // Helper function
552 inline uint64_t nElemBetweenDim(const TensorShape& X, int start, int stop) {
553  CAFFE_ENFORCE_GE(start, 0, "Invalid maximum index specified");
554  CAFFE_ENFORCE_LE(stop, X.dims_size(), "Invalid maximum index specified");
555 
556  uint64_t nElem = 1;
557  for (int i = start; i < stop; ++i) {
558  nElem *= X.dims(i);
559  }
560  return nElem;
561 }
562 
563 // Helper function for infer op inputs and outputs device information.
564 inline std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
565 InferOpInputOutputDevice(const OperatorDef& op) {
566  auto op_schema = OpSchemaRegistry::Schema(op.type());
567  if (op_schema) {
568  // op_schema found
569  return op_schema->InferDevice(op);
570 
571  } else {
572  // No schema for op.type registered
573  auto temp_schema = OpSchema();
574  return temp_schema.InferDevice(op);
575  }
576 }
577 
578 template <uint64_t OpsPerPoint>
579 OpSchema::Cost PointwiseCostInference(
580  const OperatorDef& /* unused */,
581  const vector<TensorShape>& inputs) {
582  struct OpSchema::Cost c;
583  const TensorShape X = inputs[0];
584  uint64_t nElemX = nElemFromDim(X);
585  uint64_t nElemRead = 0;
586  for (size_t i = 0; i < inputs.size(); ++i) {
587  nElemRead += nElemFromDim(inputs[i]);
588  }
589 
590  c.flops = nElemX * OpsPerPoint;
591  c.bytes_read = nElemRead * sizeof(X.data_type());
592  c.bytes_written = nElemX * sizeof(X.data_type());
593  return c;
594 }
595 
596 } // namespace caffe2
597 
598 #ifndef CAFFE2_NO_OPERATOR_SCHEMA
599 
600 #define OPERATOR_SCHEMA(name) \
601  C10_EXPORT void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
602  static OpSchema* C10_ANONYMOUS_VARIABLE(name) CAFFE2_UNUSED = \
603  &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
604 
605 #else // CAFFE2_NO_OPERATOR_SCHEMA
606 
607 #define OPERATOR_SCHEMA(name) \
608  C10_EXPORT void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
609  static OpSchema* C10_ANONYMOUS_VARIABLE(name) CAFFE2_UNUSED = \
610  1 ? nullptr : &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
611 
612 #endif // CAFFE2_NO_OPERATOR_SCHEMA
613 
614 #ifdef CAFFE2_NO_GRADIENT_OPS
615 
616 #define GRADIENT_OPERATOR_SCHEMA(name) \
617  C10_EXPORT void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
618  static OpSchema* C10_ANONYMOUS_VARIABLE(name) CAFFE2_UNUSED = \
619  1 ? nullptr : &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
620 
621 #else
622 
623 #define GRADIENT_OPERATOR_SCHEMA(name) OPERATOR_SCHEMA(name)
624 
625 #endif
626 #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.
A class to record the schema of an op.
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.
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.
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.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
OpSchema & InheritOnnxSchema()
Shortcut to InheritOnnxSchema(type_)
const string & file() const
Returns the file that the op schema is registered from.
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 ...