operator_schema.h

#ifndef CAFFE2_CORE_OPERATOR_SCHEMA_H_
#define CAFFE2_CORE_OPERATOR_SCHEMA_H_

#include <climits>
#include <functional>
#include <initializer_list>
#include <ostream>
#include <set>
#include <vector>
#include <unordered_map>

#include "caffe2/core/common.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/registry.h"
#include "caffe2/proto/caffe2.pb.h"

namespace caffe2 {

// A const value returned by OpSchema::CalculateOutput() if the number of
// output cannot be determined.
constexpr int kCannotComputeNumOutputs = -1;

class OpSchema {
 public:
  OpSchema() : file_("unknown"), line_(0) {}
  OpSchema(const string& file, const int line) : file_(file), line_(line) {}

  inline const string& file() const {
    return file_;
  }

  inline int line() const {
    return line_;
  }

  inline const char* doc() const {
    return doc_.empty() ? nullptr : doc_.c_str();
  }

  bool Verify(const OperatorDef& def) const;

  // Functions to set the property of the operator schemas.
  // Sets the number of inputs, either a fixed number or a min and a max.

  OpSchema& NumInputs(int n);
  OpSchema& NumInputs(int min, int max);
  OpSchema& NumInputs(set<int> allowed_input_nums);
  OpSchema& NumInputs(std::function<bool(int)> func);

  // Sets the number of outputs, either a fixed number, a min and a max,
  // or a function that takes in the input number and produces an output
  // number. Use only one function in the set below.
  OpSchema& NumOutputs(int n);
  OpSchema& NumOutputs(int min, int max);
  OpSchema& NumOutputs(set<int> allowed_output_nums);
  OpSchema& NumOutputs(std::function<bool(int)> func);

  OpSchema& NumInputsOutputs(std::function<bool(int, int)> func);

  // Set the function that can calculate the number of output based on the
  // number of input. Use only one function in the set below.
  OpSchema& OutputCalculator(std::function<int(int)> calc);
  OpSchema& SameNumberOfOutput();

  // Sets the rule to allow optional in-place operation.
  OpSchema& AllowInplace(std::function<bool(int, int)> inplace);
  OpSchema& AllowInplace(set<std::pair<int, int>> inplace);
  OpSchema& AllowOneToOneInplace();
  // Sets the rule to enforce in-place opeartion.
  OpSchema& EnforceInplace(std::function<bool(int, int)> inplace);
  OpSchema& EnforceInplace(set<std::pair<int, int>> inplace);
  OpSchema& EnforceOneToOneInplace();

  // Functions to deal with type and shape inference. Basically, this registers
  // a function that takes in an OperatorDef and a series of input type and
  // shape specified by TensorProto objects (whose data fields are empty), and
  // produces a series of output type and shape.
  typedef std::function<
      vector<TensorShape>(const OperatorDef&, const vector<TensorShape>&)>
      TensorInferenceFunctionType;

  OpSchema& TensorInferenceFunction(TensorInferenceFunctionType function);

  OpSchema& InheritOnnxSchema(const std::string& onnx_schema_name);

  OpSchema& IdenticalTypeAndShape();
  OpSchema& IdenticalTypeAndShapeOfInput(int idx);
  OpSchema& IdenticalTypeAndShapeOfInputDim(int idx, int dim);
  OpSchema& ScalarType(::caffe2::TensorProto_DataType dt);

  inline vector<TensorShape> InferTensor(
      const OperatorDef& def,
      const vector<TensorShape>& input_type_shape) const {
    return tensor_inference_function_(def, input_type_shape);
  }

  /*
   * @brief A struct to store various cost information about
   * an operator such as FLOPs, total memory use and parameters.
   */
  struct Cost {
    uint64_t flops; // Floating point operations.
    uint64_t bytes_moved; // Total memory used.
    uint64_t params_bytes; // Memory footprint of parameters
  };
  typedef std::function<
      struct Cost(const OperatorDef&, const vector<TensorShape>&)>
      CostInferenceFunctionType;

  OpSchema& CostInferenceFunction(CostInferenceFunctionType function);

#if 0 // def _MSC_VER

  template <typename T,
            typename = std::enable_if<
                std::is_same<CostInferenceFunctionType&&, T>:value
                >:type>
  inline OpSchema& CostInferenceFunction(T func) {
    // Note: This is here in order to resolve an MSVC compiler issue: it
    // does not automatically convert a function pointer to a std::function,
    // and needs an explicit conversion.
    return CostInferenceFunction(CostInferenceFunctionType(func));
  }
#endif // _MSC_VER

  bool HasCostInferenceFunction() const {
    return !!cost_inference_function_;
  }

  inline struct Cost InferCost(
      const OperatorDef& def,
      const vector<TensorShape>& input_tensor_shape) const {
    CAFFE_ENFORCE(
        cost_inference_function_, "Cost inference function not defined.");
    return (*cost_inference_function_)(def, input_tensor_shape);
  }

  // Functions to do documentation for the operator schema.
  OpSchema& SetDoc(const string& doc);

  struct Argument {
    Argument(const char* name, const char* description, bool required)
        : name_{name}, description_{description}, required_{required} {}

    const char* name() const {
      return name_;
    }

    const char* description() const {
      return description_;
    }

    bool is_required() const {
      return required_;
    }

   private:
    const char* name_;
    const char* description_;
    const bool required_;
  };

  OpSchema&
  Arg(const char* name, const char* description, bool required = false);

#define DECLARE_STANDARD_ARG(name, str)     \
  CAFFE2_API static const char* Arg_##name; \
  CAFFE2_API OpSchema& Arg##name(const char* description);

  DECLARE_STANDARD_ARG(IsTest, is_test)

#undef DECLARE_STANDARD_ARG

  OpSchema& Input(const int n, const char* name, const char* description);
  OpSchema& Output(const int n, const char* name, const char* description);
  // Calls the passed function with `this` as an argument. Useful for
  // adding docs for temlated/macro ops.
  OpSchema& FillUsing(std::function<void(OpSchema&)> populator);

  // Remove from documentation
  OpSchema& Private();

  // This op can pass data across devices
  OpSchema& InputsCanCrossDevices();

  int CalculateOutput(int num_input) const;

  const std::string& onnx_schema() const {
    return onnx_schema_;
  }

  int min_input() const {
    return min_input_;
  }

  int max_input() const {
    return max_input_;
  }

  int min_output() const {
    return min_output_;
  }

  int max_output() const {
    return max_output_;
  }

  bool num_inputs_allowed(int x) const {
    return num_inputs_allowed_(x);
  }

  bool num_outputs_allowed(int x) const {
    return num_outputs_allowed_(x);
  }

  bool num_inputs_outputs_allowed(int x, int y) const {
    return num_inputs_outputs_allowed_(x, y);
  }

  int inf() const {
    return std::numeric_limits<int>::max();
  }

  friend std::ostream& operator<<(std::ostream& out, const OpSchema& schema);

  const std::vector<Argument>& args() const {
    return args_;
  }

  const std::vector<std::pair<const char*, const char*>>& input_desc() const {
    return input_desc_;
  }
  const std::vector<std::pair<const char*, const char*>>& output_desc() const {
    return output_desc_;
  }
  bool private_op() {
    return private_;
  }
  bool inputs_can_cross_devices() const {
    return inputs_can_cross_devices_;
  }

  using DeviceInferenceFunctionType = std::function<
      std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>(
          const OperatorDef& def)>;

  OpSchema& DeviceInferenceFunction(DeviceInferenceFunctionType function);

  inline std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
  InferDevice(const OperatorDef& def) const {
    return device_inference_function_(def);
  }

 private:
  string file_;
  string doc_;
  string onnx_schema_;
  std::vector<Argument> args_{};
  std::vector<std::pair<const char*, const char*>> input_desc_{};
  std::vector<std::pair<const char*, const char*>> output_desc_{};
  int line_ = 0;
  int min_input_ = 0;
  int max_input_ = std::numeric_limits<int>::max();
  int min_output_ = 0;
  int max_output_ = std::numeric_limits<int>::max();
  bool private_ = false;
  bool inputs_can_cross_devices_ = false;
  std::function<bool(int)> num_inputs_allowed_ = [](int) { return true; };
  std::function<bool(int)> num_outputs_allowed_ = [](int) { return true; };
  std::function<bool(int, int)> num_inputs_outputs_allowed_ = [](int, int) {
    return true;
  };
  std::function<int(int)> calculate_output_;
  // In default, any in-place operation is neither allowed nor enforced.
  std::function<bool(int, int)> inplace_allowed_ = [](int, int) {
    return false;
  };
  std::function<bool(int, int)> inplace_enforced_ = [](int, int) {
    return false;
  };
  TensorInferenceFunctionType tensor_inference_function_ =
      [](const OperatorDef& def, const vector<TensorShape>&) {
        vector<TensorShape> out;
        for (int i = 0; i < def.output_size(); i++) {
          TensorShape ts;
          ts.set_unknown_shape(true);
          out.push_back(ts);
        }
        return out;
      };
  std::unique_ptr<CostInferenceFunctionType> cost_inference_function_ = nullptr;
  DeviceInferenceFunctionType device_inference_function_ =
      [](const OperatorDef& def) {
        auto op_device =
            def.has_device_option() ? def.device_option() : DeviceOption();
        vector<DeviceOption> in_dev(def.input_size(), op_device);
        vector<DeviceOption> out_dev(def.output_size(), op_device);
        return std::make_pair(in_dev, out_dev);
      };
};

class OpSchemaRegistry {
 public:
  static OpSchema&
  NewSchema(const string& key, const string& file, const int line) {
    auto& m = map();
    auto it = m.find(key);
    if (it != m.end()) {
      const auto& schema = it->second;
      std::ios_base::Init init;
      std::cerr << "Trying to register schema with name " << key
                << " from file " << file << " line " << line
                << ", but it is already registered from file " << schema.file()
                << " line " << schema.line();
      abort();
    }
    m.emplace(std::make_pair(key, OpSchema(file, line)));
    return m[key];
  }

  static const OpSchema* Schema(const string& key) {
    auto& m = map();
    auto it = m.find(key);
    if (it != m.end()) {
      return &it->second;
    } else {
      return nullptr;
    }
  }

 private:
  // OpSchemaRegistry should not need to be instantiated.
  OpSchemaRegistry() = delete;

  static CaffeMap<string, OpSchema>& map();
};

// Helper function for creating simple tensorproto with dimension and type
template <typename T_I = int>
inline TensorShape CreateTensorShape(
    vector<T_I> dims,
    ::caffe2::TensorProto_DataType dt) {
  TensorShape ts;
  for (int d : dims) {
    ts.add_dims(d);
  }
  ts.set_data_type(dt);
  return ts;
}

// Helper function
inline vector<TIndex> GetDimsVector(const TensorShape& shape) {
  vector<TIndex> dims;
  for (auto d : shape.dims()) {
    dims.push_back(d);
  }
  return dims;
}

// Helper function for infer op inputs and outputs device information.
inline std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
InferOpInputOutputDevice(const OperatorDef& op) {
  auto op_schema = OpSchemaRegistry::Schema(op.type());
  CAFFE_ENFORCE(
      op_schema, "Device inference failed. No schema for: ", op.type());
  // TODO(wyiming) : add try catch here.
  return op_schema->InferDevice(op);
}

template <uint64_t OpsPerPoint>
OpSchema::Cost PointwiseCostInference(
    const OperatorDef& /* unused */,
    const vector<TensorShape>& inputs) {
  struct OpSchema::Cost c;
  const TensorShape X = inputs[0];
  uint64_t size = 1;

  for (auto i = 0; i < X.dims().size(); ++i) {
    size *= X.dims(i);
  }

  c.flops = size * OpsPerPoint;
  c.bytes_moved = size * sizeof(X.data_type());
  return c;
}

} // namespace caffe2

#ifndef CAFFE2_NO_OPERATOR_SCHEMA

#define OPERATOR_SCHEMA(name)                            \
  void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(name) =      \
      &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
#define OPERATOR_SCHEMA_STR(name)                                  \
  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(schema_registration) = \
      &OpSchemaRegistry::NewSchema(name, __FILE__, __LINE__)

#else // CAFFE2_NO_OPERATOR_SCHEMA

#define OPERATOR_SCHEMA(name)                            \
  void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(){}; \
  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(name) =      \
      1 ? nullptr : &OpSchemaRegistry::NewSchema(#name, __FILE__, __LINE__)
#define OPERATOR_SCHEMA_STR(name)                                  \
  static OpSchema* CAFFE_ANONYMOUS_VARIABLE(schema_registration) = \
      1 ? nullptr : &OpSchemaRegistry::NewSchema(name, __FILE__, __LINE__)

#endif // CAFFE2_NO_OPERATOR_SCHEMA

#endif // CAFFE2_CORE_OPERATOR_SCHEMA_H_