pattern_net_transform.h

#pragma once

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

namespace caffe2 {

class CAFFE2_API PatternNetTransform : public Transform {
 public:
  PatternNetTransform(const NetDef& pattern_net, const NetDef& replace_net)
      : p_(transform::Graph(pattern_net)), r_(transform::Graph(replace_net)) {
    // external input and output must match!
    CAFFE_ENFORCE(
        p_.external_input() == r_.external_input(),
        "External inputs do not match!");
    CAFFE_ENFORCE(
        p_.external_output() == r_.external_output(),
        "External outputs do not match!");
    ordered_ops_ = GetPatternTraversalOrder(p_);
    inverse_ops_.resize(ordered_ops_.size());
    for (size_t i = 0; i < ordered_ops_.size(); i++) {
      inverse_ops_[ordered_ops_[i]] = i;
    }
  }

  void EnableArgumentMatching() {
    argument_match_ = true;
  }

  void DisableArgumentMatching() {
    argument_match_ = false;
  }

 protected:
  bool PatternRule(
      const transform::Graph& g,
      const std::vector<int>& subgraph,
      int idx) override;
  bool ValidatorRule(
      const transform::Graph& g,
      const std::vector<int>& subgraph) override;
  bool ReplaceRule(const std::vector<int>& subgraph, transform::Graph* g_ptr)
      override;

 private:
  std::vector<int> GetPatternTraversalOrder(const transform::Graph& g);

  // Graph of Pattern NetDef
  transform::Graph p_;

  // The Traversal Order of the Pattern Net's Operators
  // This is a permutation of the numbers from {0, ..., p.size()-1}
  std::vector<int> ordered_ops_;

  // The Inverse of the Traversal Order of the Pattern Net's Operators
  // That is, inverse_ops[ordered_ops[i]] == i is always true.
  std::vector<int> inverse_ops_;

  // Graph of Replace NetDef
  transform::Graph r_;

  // This flag determines if the transform will match operator arguments.
  bool argument_match_ = false;

  const string TransformBlobWrapper(const string& blob_name) {
    return "transform/" + blob_name + "_" + c10::to_string(ssa_id_);
  }

  int ssa_id_ = 0;
};

} // namespace caffe2