transform.h

#pragma once

#include "caffe2/core/common.h"
#include "caffe2/core/graph.h"
#include "caffe2/core/workspace.h"
#include "caffe2/proto/caffe2_pb.h"
#include "caffe2/utils/proto_utils.h"

namespace caffe2 {

class CAFFE2_API Transform {
 public:
  Transform() {}

  NetDef ApplyTo(const NetDef& orig_net_def);

  virtual ~Transform() {}

  enum PatternMatchType {
    CONNECTED_SUBGRAPH,
    SORTED_WRT_EXECUTION_ORDER,
    GENERAL
  };

  std::vector<std::vector<int>> PatternMatch(const transform::Graph& graph);

  void ReplacePattern(
      const std::vector<std::vector<int>>& matches,
      transform::Graph* graph);

 protected:
  virtual bool PatternRule(
      const transform::Graph& g,
      const std::vector<int>& subgraph,
      int /*idx*/) {
    CAFFE_NOT_IMPLEMENTED;
  }

  virtual bool ValidatorRule(
      const transform::Graph& g,
      const std::vector<int>& subgraph) {
    CAFFE_NOT_IMPLEMENTED;
  }

  virtual bool ReplaceRule(
      const std::vector<int>& subgraph,
      transform::Graph* g_ptr) {
    CAFFE_NOT_IMPLEMENTED;
  }

  void SetPatternMatchType(PatternMatchType type) {
    pattern_match_type_ = type;
  }

 private:
  void PatternMatchHelper(
      const transform::Graph& graph,
      const std::vector<bool>& matched,
      std::vector<int>* subgraph_ptr,
      std::vector<int>* best_subgraph_ptr);
  void TryNeighbors(
      const transform::Graph& graph,
      const std::map<int, std::vector<string>>& neighbors,
      const std::vector<bool>& matched,
      std::vector<int>* subgraph_ptr,
      std::vector<int>* best_subgraph_ptr);

  PatternMatchType pattern_match_type_ = CONNECTED_SUBGRAPH;
};

// Creates a Transform based on a key, which should be defined in registry.
CAFFE2_API unique_ptr<Transform> CreateTransform(string key);

C10_DECLARE_REGISTRY(TransformRegistry, Transform);
#define REGISTER_TRANSFORM(name, ...) \
  C10_REGISTER_CLASS(TransformRegistry, name, __VA_ARGS__)

// Create a Transform object from registry,
// and immediately apply it to a Netdef.
CAFFE2_API NetDef ApplyTransform(const string& key, const NetDef& netdef);

// Create a Transform object from registry, apply it to a NetDef.
// Will only return the transformed net if it is faster than the old net.
// This will run the init net first, will run the two nets warmup_runs times.
// Then, we will take the average time of main_runs runs, and only keep the
// transformed net if it is faster by a factor of improvement_threshold.
CAFFE2_API NetDef ApplyTransformIfFaster(
    const string& key,
    const NetDef& netdef,
    const NetDef& init_netdef,
    const int warmup_runs,
    const int main_runs,
    const double improvement_threshold);

} // namespace