pytorch/caffe2/opt/backend_transformer_base.cc

#include "caffe2/opt/backend_transformer_base.h"
#include "caffe2/onnx/onnx_exporter.h"
#include "caffe2/utils/proto_utils.h"

namespace caffe2 {

// Populate 'net_pos' argument for any ops that don't already have it. 'net_pos'
// we populate here starts after the max 'net_pos' value we encountered.
void BackendTransformerBase::annotateOpIndex(NetDef* net) {
  // find the max net_pos that we have so far.
  int i = -1;
  for (const auto& op : net->op()) {
    ArgumentHelper helper(op);
    int old_index = helper.GetSingleArgument(op, kNetPos, -1);
    i = std::max(i, old_index);
  }

  // populate net_pos for any op that doesn't already have it.
  for (auto& op : *(net->mutable_op())) {
    if (!ArgumentHelper::HasArgument(op, kNetPos)) {
      AddArgument(kNetPos, ++i, &op);
    }
  }
}

std::string BackendTransformerBase::getModelId(const NetDef& net) {
  static std::atomic<size_t> seq_id{0};
  std::string model_id;
  for (const auto& arg : net.arg()) {
    if (arg.name() == kModelId) {
      if (arg.has_s()) {
        model_id = arg.s();
      } else if (arg.has_i()) {
        model_id = c10::to_string(arg.i());
      }
      break;
    }
  }

  if (model_id.empty()) {
    model_id = "unnamed_" + c10::to_string(seq_id++);
  }
  return model_id;
}

TensorProto wrapShapeInfoIntoTensorProto(
    const std::string& name,
    const ShapeInfo& shape_info) {
  TensorProto t;
  t.set_name(name);
  t.set_data_type(shape_info.shape.data_type());
  for (const auto i : shape_info.shape.dims()) {
    t.add_dims(i);
  }
  for (const auto& dimType : shape_info.getDimType()) {
    t.add_int32_data(static_cast<int32_t>(dimType));
  }
  return t;
}

QTensorProto wrapShapeInfoIntoQTensorProto(
    const std::string& name,
    const ShapeInfo& shape_info) {
  QTensorProto t;
  CAFFE_ENFORCE(
      shape_info.is_quantized == true,
      "Only quantized shapeinfo can be extracted into QTensor!");
  t.set_name(name);
  t.set_data_type(shape_info.shape.data_type());
  t.set_axis(shape_info.q_info.axis);
  t.set_is_multiparam(true);
  for (const auto i : shape_info.q_info.scale) {
    t.add_scales(i);
  }
  t.set_scale(1.0);
  for (const auto i : shape_info.q_info.offset) {
    t.add_biases(i);
  }
  t.set_bias(0.0);
  // precision and is_signed is not used in onnxifi workflow, but it is required
  // field
  t.set_precision(0);
  // NOLINTNEXTLINE(modernize-use-bool-literals)
  t.set_is_signed(0);
  for (const auto i : shape_info.shape.dims()) {
    t.add_dims(i);
  }
  for (const auto& dimType : shape_info.getDimType()) {
    t.add_data(static_cast<int32_t>(dimType));
  }
  return t;
}

ShapeInfoMap BackendTransformerBase::ssaRewriteAndMapNames(
    Workspace* ws,
    NetDef* pred_net,
    const ShapeInfoMap& input_shape_hints) {
  input_mapping_ = onnx::SsaRewrite(nullptr, pred_net);
  // Annote the ops with net position
  annotateOpIndex(pred_net);

  // Since we are going to create a mapped workspace, we need to make sure that
  // the parent workspace has the mapped blob names. If the blobs don't exist
  // (usually such blobs are input tensor names), we exclude them from mapping.
  std::vector<std::string> exclude_mapping;
  for (const auto& kv : input_mapping_) {
    if (!ws->HasBlob(kv.second)) {
      exclude_mapping.emplace_back(kv.first);
    }
  }
  for (const auto& i : exclude_mapping) {
    input_mapping_.erase(i);
  }

  ShapeInfoMap shape_hints_mapped;
  for (const auto& kv : input_shape_hints) {
    shape_hints_mapped.emplace(kv.first, kv.second);
  }
  return shape_hints_mapped;
}

ShapeInfoMap BackendTransformerBase::inferShapes(
    Workspace* ws,
    NetDef* pred_net,
    const ShapeInfoMap& shape_hints_mapped,
    const BoundShapeSpec& spec) {
  ShapeInfoMap shape_map;

  // Populate shapes from workplace
  const std::vector<std::string> ws_blobs = ws->Blobs();
  for (const auto& s : ws_blobs) {
    auto shape_info = getShapeInfoFromBlob(ws->GetBlob(s));
    if (shape_info.dimTypeIsSet()) {
      shape_map.emplace(s, shape_info);
    }
  }
  for (const auto& s : shape_hints_mapped) {
    shape_map.insert(s);
  }
  auto eng = BoundShapeInferencerRegistry()->Create("C10", spec);
  eng->InferBoundShapeAndType(*pred_net, shape_map, ws);
  const auto& out_map = eng->shape_info();
  shape_map.clear();
  for (const auto& kv : out_map) {
    shape_map.emplace(
        std::piecewise_construct,
        std::forward_as_tuple(kv.first),
        std::forward_as_tuple(
            kv.second.getDimType(),
            kv.second.shape,
            kv.second.is_quantized,
            kv.second.q_info));
  }
  return shape_map;
}

void BackendTransformerBase::addShapeToNet(
    NetDef& shape_net,
    const ShapeInfoMap& shape_hints) const {
  auto* shape_arg = shape_net.add_arg();
  auto* qshape_arg = shape_net.add_arg();
  shape_arg->set_name("shape_info");
  qshape_arg->set_name("qshape_info");
  for (const auto& kv : shape_hints) {
    if (!kv.second.is_quantized) {
      auto t = wrapShapeInfoIntoTensorProto(kv.first, kv.second);
      shape_arg->mutable_tensors()->Add()->CopyFrom(t);
    } else {
      auto t = wrapShapeInfoIntoQTensorProto(kv.first, kv.second);
      qshape_arg->mutable_qtensors()->Add()->CopyFrom(t);
    }
  }
}

void BackendTransformerBase::dumpNet(
    const NetDef& pred_net,
    const ShapeInfoMap& shape_hints,
    const std::string& fname) const {
  NetDef shape_net(pred_net);
  addShapeToNet(shape_net, shape_hints);
  WriteProtoToTextFile(shape_net, fname, false);
}
} // namespace caffe2