subgraph/tensorrt/tensorrt.cc

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file tensorrt.cc
 * \brief TensorRT operation registration
 * \author Marek Kolodziej, Clement Fuji Tsang, Serge Panev
*/

#if MXNET_USE_TENSORRT

#include "./tensorrt-inl.h"

#include <NvInfer.h>

namespace mxnet {
namespace op {

inline uint32_t TRTNumInputs(const nnvm::NodeAttrs& attrs) {
  const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
  const auto inputs_to_idx = param.inputs_to_idx;
  return inputs_to_idx.size();
}

inline std::vector<std::string> TRTListInputNames(const nnvm::NodeAttrs& attrs) {
  std::vector<std::string> outputs;
  const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
  const auto inputs_to_idx = param.inputs_to_idx;
  for (auto& p : inputs_to_idx) {
    outputs[p.second] = p.first;
  }
  return outputs;
}

inline bool TRTInferShape(const nnvm::NodeAttrs& attrs,
                          std::vector<TShape> *in_shapes,
                          std::vector<TShape> *out_shapes) {
  using namespace exec;
  const nnvm::Symbol subgraph_sym = *(attrs.subgraphs[0]);
  const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
  auto params_map = param.params_map;
  auto inputs_to_idx = param.inputs_to_idx;
  nnvm::Graph g;
  g.outputs = subgraph_sym.outputs;
  const auto& idx_g = g.indexed_graph();
  CHECK_EQ(idx_g.input_nodes().size(), in_shapes->size() + params_map.size());
  CHECK_EQ(idx_g.outputs().size(), out_shapes->size());

  // Put the input and output shapes to the shape vector.
  mxnet::ShapeVector shapes(idx_g.num_node_entries());
  const auto &input_nids = idx_g.input_nodes();
  CHECK_EQ(input_nids.size(), in_shapes->size() + params_map.size());
  for (size_t i = 0; i < input_nids.size(); i++) {
    auto node = idx_g[input_nids[i]].source;
    auto eid = idx_g.entry_id(input_nids[i], 0);
    auto it_params = params_map.find(node->attrs.name);
    auto it_inputs = inputs_to_idx.find(node->attrs.name);
    if (it_params != params_map.end()) {
      shapes[eid] = it_params->second.shape();
    } else if (it_inputs != inputs_to_idx.end()) {
      shapes[eid] = in_shapes->at(it_inputs->second);
    } else {
      LOG(FATAL) << node->attrs.name << " shape information is missing for attributes inference";
    }
  }
  CHECK_EQ(g.outputs.size(), out_shapes->size());
  for (size_t i = 0; i < out_shapes->size(); i++) {
    auto eid = idx_g.entry_id(g.outputs[i]);
    shapes[eid] = out_shapes->at(i);
  }

  // Infer shape of the graph.
  g.attrs["shape"] = std::make_shared<dmlc::any>(std::move(shapes));
  g = exec::InferShape(std::move(g));
  // Copy the inferred shape back to the input shapes and the output shapes.
  shapes = g.GetAttr<mxnet::ShapeVector>("shape");
  // assign to in_shapes
  for (size_t i = 0; i < input_nids.size(); ++i) {
    const auto node = idx_g[input_nids[i]].source;
    const auto eid = idx_g.entry_id(input_nids[i], 0);
    auto it = inputs_to_idx.find(node->attrs.name);
    if (it != inputs_to_idx.end()) {
      SHAPE_ASSIGN_CHECK(*in_shapes, it->second, shapes[eid]);
    }
  }
  // assign to out_shapes
  for (size_t i = 0; i < g.outputs.size(); ++i) {
    const auto eid = idx_g.entry_id(g.outputs[i]);
    SHAPE_ASSIGN_CHECK(*out_shapes, i, shapes[eid]);
  }
  // Check if we have inferred the shapes correctly.
  return g.GetAttr<size_t>("shape_num_unknown_nodes") == 0;
}

inline bool TRTInferType(const nnvm::NodeAttrs& attrs,
                    std::vector<int> *in_types,
                    std::vector<int> *out_types) {
  const nnvm::Symbol subgraph_sym = *(attrs.subgraphs[0]);
  const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
  auto params_map = param.params_map;
  auto inputs_to_idx = param.inputs_to_idx;

  nnvm::Graph g;
  g.outputs = subgraph_sym.outputs;
  const auto& idx_g = g.indexed_graph();
  CHECK_EQ(idx_g.input_nodes().size(), in_types->size() + params_map.size());
  CHECK_EQ(idx_g.outputs().size(), out_types->size());

  // Put the input and output data types to the dtype vector.
  nnvm::DTypeVector types(idx_g.num_node_entries(), -1);
  const auto &input_nids = idx_g.input_nodes();
  CHECK_EQ(input_nids.size(), in_types->size() + params_map.size());
  for (size_t i = 0; i < input_nids.size(); i++) {
    auto node = idx_g[input_nids[i]].source;
    auto eid = idx_g.entry_id(input_nids[i], 0);
    auto it_params = params_map.find(node->attrs.name);
    auto it_inputs = inputs_to_idx.find(node->attrs.name);
    if (it_params != params_map.end()) {
      types[eid] = -1;
    } else if (it_inputs != inputs_to_idx.end()) {
      types[eid] = in_types->at(it_inputs->second);
    } else {
      LOG(FATAL) << node->attrs.name
                 << " dtype information is missing for attributes inference";
    }
  }
  CHECK_EQ(g.outputs.size(), out_types->size());
  for (size_t i = 0; i < out_types->size(); i++) {
    auto eid = idx_g.entry_id(g.outputs[i]);
    types[eid] = out_types->at(i);
  }

  // Infer data type of the graph.
  g.attrs["dtype"] = std::make_shared<dmlc::any>(std::move(types));
  g = exec::InferType(std::move(g));

  types = g.GetAttr<nnvm::DTypeVector>("dtype");
  // assign to in_types
  for (size_t i = 0; i < input_nids.size(); ++i) {
    const auto node = idx_g[input_nids[i]].source;
    const auto eid = idx_g.entry_id(input_nids[i], 0);
    auto it = inputs_to_idx.find(node->attrs.name);
    if (it != inputs_to_idx.end()) {
      TYPE_ASSIGN_CHECK(*in_types, it->second, types[eid]);
    }
  }
  // assign to out_types
  for (size_t i = 0; i < g.outputs.size(); ++i) {
    const auto eid = idx_g.entry_id(g.outputs[i]);
    TYPE_ASSIGN_CHECK(*out_types, i, types[eid]);
  }

  // Check if we have inferred the dtypes correctly.
  return g.GetAttr<size_t>("dtype_num_unknown_nodes") == 0;
}

inline bool TRTInferStorageType(const nnvm::NodeAttrs& attrs,
                           const int dev_mask,
                           DispatchMode* dispatch_mode,
                           std::vector<int>* in_stypes,
                           std::vector<int>* out_stypes) {
  const nnvm::Symbol subgraph_sym = *(attrs.subgraphs[0]);
  const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
  auto params_map = param.params_map;
  auto inputs_to_idx = param.inputs_to_idx;
  nnvm::Graph g;
  g.outputs = subgraph_sym.outputs;
  const auto& idx_g = g.indexed_graph();
  CHECK_EQ(idx_g.input_nodes().size(), in_stypes->size() + params_map.size());
  CHECK_EQ(idx_g.outputs().size(), out_stypes->size());
  exec::DevMaskVector dev_masks(idx_g.num_node_entries(), dev_mask);

  // Put the input and output storages to the storage vector.
  StorageTypeVector stypes(idx_g.num_node_entries(), kUndefinedStorage);
  const auto &input_nids = idx_g.input_nodes();
  CHECK_EQ(input_nids.size(), in_stypes->size() + params_map.size());
  for (size_t i = 0; i < input_nids.size(); i++) {
    auto node = idx_g[input_nids[i]].source;
    auto eid = idx_g.entry_id(input_nids[i], 0);
    auto it_params = params_map.find(node->attrs.name);
    auto it_inputs = inputs_to_idx.find(node->attrs.name);
    if (it_params != params_map.end()) {
      stypes[eid] = it_params->second.storage_type();
    } else if (it_inputs != inputs_to_idx.end()) {
      stypes[eid] = in_stypes->at(it_inputs->second);
    } else {
      LOG(FATAL) << node->attrs.name
                 << " storage type information is missing for attributes inference";
    }
  }
  CHECK_EQ(g.outputs.size(), out_stypes->size());
  for (size_t i = 0; i < out_stypes->size(); i++) {
    auto eid = idx_g.entry_id(g.outputs[i]);
    stypes[eid] = out_stypes->at(i);
  }

  // Infer storage type of the graph.
  bool dev_match = g.attrs.count("dev_mask") &&
                   g.GetAttr<exec::DevMaskVector>("dev_mask") == dev_masks;
  if (!dev_match) {
    g.attrs["dev_mask"] = std::make_shared<dmlc::any>(std::move(dev_masks));
  }
  g.attrs["storage_type"] = std::make_shared<dmlc::any>(std::move(stypes));
  g = exec::InferStorageType(std::move(g));

  stypes = g.GetAttr<StorageTypeVector>("storage_type");
  // assign to in_types
  for (size_t i = 0; i < input_nids.size(); ++i) {
    const auto node = idx_g[input_nids[i]].source;
    const auto eid = idx_g.entry_id(input_nids[i], 0);
    auto it = inputs_to_idx.find(node->attrs.name);
    if (it != inputs_to_idx.end()) {
      STORAGE_TYPE_ASSIGN_CHECK(*in_stypes, it->second, stypes[eid]);
    }
  }

  DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
  // assign to out_types
  for (size_t i = 0; i < g.outputs.size(); ++i) {
    const auto eid = idx_g.entry_id(g.outputs[i]);
    STORAGE_TYPE_ASSIGN_CHECK(*out_stypes, i, stypes[eid]);
  }
  // Check if we have inferred the storages correctly.
  return g.GetAttr<size_t>("storage_type_num_unknown_nodes") == 0;
}

void TRTParamParser(nnvm::NodeAttrs* attrs) {
  TRTParam& _param = nnvm::get<TRTParam>(attrs->parsed);
  std::string prefix = "subgraph_param_";
  std::string str_dtype, str_shape, str_pointer, _tmp;
  for (auto it = attrs->dict.begin(); it != attrs->dict.end();) {
    std::string attrs_name = it->first;
    if (std::equal(prefix.begin(), prefix.end(), attrs_name.begin())) {
      std::string param_name = attrs_name.substr(prefix.size(),
                                                 attrs_name.size() - prefix.size());
      // TODO(cfujitsang): find a less dirty way to give weights
      NDArray *cache = reinterpret_cast<NDArray*>(stol(it->second));
      _param.params_map.emplace(param_name, cache->Copy(Context()));
      _param.params_map[param_name].WaitToRead();
      it = attrs->dict.erase(it);
    } else {
      ++it;
    }
  }
  attrs->parsed = std::move(_param);
}

OpStatePtr TRTCreateState(const nnvm::NodeAttrs& attrs, Context ctx,
                          const std::vector<TShape>& in_shape,
                          const std::vector<int>& in_type) {
  const auto& node_param = nnvm::get<TRTParam>(attrs.parsed);
  const bool tensorrt_int8 = node_param.int8_mode;
  nnvm::Graph graph;
  graph.outputs = attrs.subgraphs[0]->outputs;
  uint32_t max_batch_size = dmlc::GetEnv("MXNET_TENSORRT_MAX_BATCH_SIZE", in_shape[0][0]);
  if (max_batch_size < in_shape[0][0]) {
    LOG(INFO) << "Warning: max batch size changed to be is: " << in_shape[0][0]
              << " instead of: " << max_batch_size;
    max_batch_size = in_shape[0][0];
  }
  std::unordered_map<std::string, NDArray> params_map = node_param.params_map;
  const auto& inputs_to_idx = node_param.inputs_to_idx;
  const auto& outputs_to_idx = node_param.outputs_to_idx;
  const auto& idx_g = graph.indexed_graph();
  const auto& input_nids = idx_g.input_nodes();

  // needed by the int8 calibrator
  std::unordered_map<std::string, std::pair<void*, size_t>> input_buffers;
  mxnet::ShapeVector shape_inputs(input_nids.size());
  nnvm::DTypeVector dtype_inputs(input_nids.size());
  for (size_t i = 0; i < input_nids.size(); ++i) {
    auto node = idx_g[input_nids[i]].source;
    auto it_params = params_map.find(node->attrs.name);
    auto it_inputs = inputs_to_idx.find(node->attrs.name);
    if (it_params != params_map.end()) {
      shape_inputs[i] = it_params->second.shape();
      dtype_inputs[i] = it_params->second.dtype();
    } else if (it_inputs != inputs_to_idx.end()) {
      shape_inputs[i] = in_shape[it_inputs->second];
      dtype_inputs[i] = in_type[it_inputs->second];
      if (tensorrt_int8) {
        int dtype_size;
        if (dtype_inputs[i] == mshadow::kFloat32) {
          dtype_size = 4;
        } else if (dtype_inputs[i] == mshadow::kFloat16) {
          dtype_size = 2;
        } else {
          LOG(FATAL) << "TensorRT op supports only float32 and float16 inputs.";
        }
        size_t buffer_size = shape_inputs[i].Size() * dtype_size;
        void *ptr;
        MSHADOW_CUDA_CALL(cudaMalloc(&ptr, buffer_size));
        input_buffers.emplace(node->attrs.name,
                              std::make_pair(ptr, buffer_size));
      }
    } else {
      LOG(FATAL) << node->attrs.name << " attribute is missing for attributes inference";
    }
  }
  mxnet::ShapeVector out_shape(graph.outputs.size());
  nnvm::DTypeVector out_type(graph.outputs.size(), -1);
  mxnet::ShapeVector _in_shape(in_shape.begin(), in_shape.end());
  nnvm::DTypeVector _in_type(in_type.begin(), in_type.end());
  TRTInferShape(attrs, &_in_shape, &out_shape);
  TRTInferType(attrs, &_in_type, &out_type);
  nnvm::DTypeVector dtypes(idx_g.num_node_entries());
  mxnet::ShapeVector shapes(idx_g.num_node_entries());
  for (size_t i = 0; i < graph.outputs.size(); ++i) {
    auto eid = idx_g.entry_id(graph.outputs[i]);
    dtypes[eid] = out_type[i];
    shapes[eid] = out_shape[i];
  }
  graph.attrs["dtype_inputs"] = std::make_shared<nnvm::any>(std::move(dtype_inputs));
  graph.attrs["shape_inputs"] = std::make_shared<nnvm::any>(std::move(shape_inputs));
  graph.attrs["dtype"]        = std::make_shared<nnvm::any>(std::move(dtypes));
  graph.attrs["shape"]        = std::make_shared<nnvm::any>(std::move(shapes));

  std::unique_ptr<::onnx_to_tensorrt::TRTInt8Calibrator> calibrator;
  if (tensorrt_int8) {
    calibrator.reset(
      new ::onnx_to_tensorrt::TRTInt8Calibrator(params_map, std::move(input_buffers),
                                                max_batch_size, node_param.calibration_iters));
  }
  auto onnx_graph = op::nnvm_to_onnx::ConvertNnvmGraphToOnnx(graph, &params_map);
  uint32_t verbose = dmlc::GetEnv("MXNET_TENSORRT_VERBOSE", 0);
  auto log_lvl = nvinfer1::ILogger::Severity::kWARNING;
  if (verbose != 0) {
    log_lvl = nvinfer1::ILogger::Severity::kVERBOSE;
  }

  auto trt_tuple = ::onnx_to_tensorrt::onnxToTrtCtx(onnx_graph, node_param.fp16_mode,
                                                    max_batch_size, 1 << 30,
                                                    calibrator.get(),
                                                    log_lvl);

  return OpStatePtr::Create<TRTEngineParam>(std::move(std::get<0>(trt_tuple)),
                                            std::move(std::get<1>(trt_tuple)),
                                            std::move(std::get<2>(trt_tuple)),
                                            inputs_to_idx, outputs_to_idx,
                                            max_batch_size,
                                            std::move(calibrator),
                                            std::move(std::get<3>(trt_tuple)));
}

NNVM_REGISTER_OP(_TensorRT)
    .describe(R"code(TRT operation (one engine)
)code" ADD_FILELINE)
    .set_num_inputs(TRTNumInputs)
    .set_num_outputs(DefaultSubgraphOpNumOutputs)
    .set_attr_parser(TRTParamParser)
    .set_attr<mxnet::FInferShape>("FInferShape", TRTInferShape)
    .set_attr<nnvm::FInferType>("FInferType", TRTInferType)
    .set_attr<nnvm::FListInputNames>("FListInputNames", TRTListInputNames)
    .set_attr<nnvm::FListOutputNames>("FListOutputNames", DefaultSubgraphOpListOutputs)
    .set_attr<FCreateOpState>("FCreateOpState", TRTCreateState)
    .set_attr<FIsCUDAGraphsCompatible>("FIsCUDAGraphsCompatible",
        [](const NodeAttrs& attrs, const bool) {
          const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
          return !param.int8_mode;
        })
    .set_attr<FInferStorageType>("FInferStorageType", TRTInferStorageType);

MXNET_REGISTER_SUBGRAPH_BACKEND(TensorRT);

MXNET_REGISTER_SUBGRAPH_PROPERTY(TensorRT, TensorrtProperty);
}  // namespace op
}  // namespace mxnet

#endif  // MXNET_USE_TENSORRT