engine.h 7.1 KB
Newer Older
Y
Yan Chunwei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

Licensed 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. */

#pragma once

#include <NvInfer.h>
#include <memory>
19
#include <string>
Y
Yan Chunwei 已提交
20
#include <unordered_map>
21
#include <vector>
Y
Yan Chunwei 已提交
22 23
#include "paddle/fluid/inference/engine.h"
#include "paddle/fluid/inference/tensorrt/helper.h"
24
#include "paddle/fluid/inference/utils/singleton.h"
Y
Yan Chunwei 已提交
25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40

namespace paddle {
namespace inference {
namespace tensorrt {

/*
 * TensorRT Engine.
 *
 * There are two alternative ways to use it, one is  to build from a paddle
 * protobuf model, another way is to manully construct the network.
 */
class TensorRTEngine : public EngineBase {
 public:
  // Weight is model parameter.
  class Weight {
   public:
41
    Weight(nvinfer1::DataType dtype, void* value, size_t num_elem) {
Y
Yan Chunwei 已提交
42 43 44 45 46 47
      w_.type = dtype;
      w_.values = value;
      w_.count = num_elem;
    }
    const nvinfer1::Weights& get() { return w_; }

48 49
    std::vector<int64_t> dims;

Y
Yan Chunwei 已提交
50 51 52 53
   private:
    nvinfer1::Weights w_;
  };

Y
Yan Chunwei 已提交
54 55
  TensorRTEngine(int max_batch, int max_workspace,
                 cudaStream_t* stream = nullptr,
Y
Yan Chunwei 已提交
56 57 58
                 nvinfer1::ILogger& logger = NaiveLogger::Global())
      : max_batch_(max_batch),
        max_workspace_(max_workspace),
Y
Yan Chunwei 已提交
59
        stream_(stream ? stream : &default_stream_),
60 61 62
        logger_(logger) {
    cudaStreamCreate(&default_stream_);
  }
Y
Yan Chunwei 已提交
63 64 65 66

  virtual ~TensorRTEngine();

  // TODO(Superjomn) implement it later when graph segmentation is supported.
67
  void Build(const DescType& paddle_model) override;
Y
Yan Chunwei 已提交
68

69
  void Execute(int batch_size) override;
Y
Yan Chunwei 已提交
70 71 72 73

  // Initialize the inference network, so that TensorRT layers can add to this
  // network.
  void InitNetwork() {
74
    infer_builder_.reset(createInferBuilder(&logger_));
Y
Yan Chunwei 已提交
75 76 77 78 79 80 81 82 83 84 85 86 87 88
    infer_network_.reset(infer_builder_->createNetwork());
  }
  // After finishing adding ops, freeze this network and creates the executation
  // environment.
  void FreezeNetwork();

  // Add an input and set its name, data type and dimention.
  nvinfer1::ITensor* DeclareInput(const std::string& name,
                                  nvinfer1::DataType dtype,
                                  const nvinfer1::Dims& dim);
  // Set the offset-th output from a layer as the network's output, and set its
  // name.
  void DeclareOutput(const nvinfer1::ILayer* layer, int offset,
                     const std::string& name);
L
Luo Tao 已提交
89 90
  // Set the itensor_map_[name] as the network's output, and set its name.
  void DeclareOutput(const std::string& name);
Y
Yan Chunwei 已提交
91 92 93 94 95

  // GPU memory address for an ITensor with specific name. One can operate on
  // these memory directly for acceleration, for example, output the converted
  // data directly to the buffer to save data copy overhead.
  // NOTE this should be used after calling `FreezeNetwork`.
Y
Yan Chunwei 已提交
96 97 98
  Buffer& buffer(const std::string& name) override;

  cudaStream_t* stream() { return stream_; }
Y
Yan Chunwei 已提交
99 100

  // Fill an input from CPU memory with name and size.
101
  void SetInputFromCPU(const std::string& name, const void* data, size_t size);
Y
Yan Chunwei 已提交
102 103
  // TODO(Superjomn) is this method necessary given that buffer(xxx) can be
  // accessed directly. Fill an input from GPU memory with name and size.
104
  void SetInputFromGPU(const std::string& name, const void* data, size_t size);
Y
Yan Chunwei 已提交
105
  // Get an output called name, the output of tensorrt is in GPU, so this method
106
  // Return the output's GPU memory address without copy.
Y
Yan Chunwei 已提交
107
  void* GetOutputInGPU(const std::string& name);
108
  // Copy data into dst inside the GPU device.
109
  void GetOutputInGPU(const std::string& name, void* dst);
Y
Yan Chunwei 已提交
110 111
  // LOW EFFICENCY! Get output to CPU, this will trigger a memory copy from GPU
  // to CPU.
112
  void GetOutputInCPU(const std::string& name, void* dst);
L
Luo Tao 已提交
113 114 115 116
  // Fill an ITensor into map itensor_map_.
  void SetITensor(const std::string& name, nvinfer1::ITensor* tensor);
  // Get an ITensor called name.
  nvinfer1::ITensor* GetITensor(const std::string& name);
Y
Yan Chunwei 已提交
117 118 119

  nvinfer1::ICudaEngine* engine() { return infer_engine_.get(); }
  nvinfer1::INetworkDefinition* network() { return infer_network_.get(); }
120 121
  void SetRuntimeBatch(size_t batch_size);
  int GetRuntimeBatch();
Y
Yan Chunwei 已提交
122 123 124 125

 private:
  // the max batch size
  int max_batch_;
126 127
  // the runtime batch size
  static int runtime_batch_;
Y
Yan Chunwei 已提交
128 129
  // the max memory size the engine uses
  int max_workspace_;
130 131 132

  // batch size of the current data, will be updated each Executation.
  int batch_size_{-1};
Y
Yan Chunwei 已提交
133
  cudaStream_t* stream_;
Y
Yan Chunwei 已提交
134 135
  // If stream_ is not set from outside, hold its own stream.
  cudaStream_t default_stream_;
Y
Yan Chunwei 已提交
136 137
  nvinfer1::ILogger& logger_;

Y
Yan Chunwei 已提交
138
  std::vector<Buffer> buffers_;
Y
Yan Chunwei 已提交
139 140
  // max data size for the buffers.
  std::unordered_map<std::string /*name*/, size_t /*max size*/> buffer_sizes_;
L
Luo Tao 已提交
141 142
  std::unordered_map<std::string /*name*/, nvinfer1::ITensor* /*ITensor*/>
      itensor_map_;
Y
Yan Chunwei 已提交
143 144 145 146

  // TensorRT related internal members
  template <typename T>
  struct Destroyer {
147 148 149 150 151
    void operator()(T* x) {
      if (x) {
        x->destroy();
      }
    }
Y
Yan Chunwei 已提交
152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174
  };
  template <typename T>
  using infer_ptr = std::unique_ptr<T, Destroyer<T>>;
  infer_ptr<nvinfer1::IBuilder> infer_builder_;
  infer_ptr<nvinfer1::INetworkDefinition> infer_network_;
  infer_ptr<nvinfer1::ICudaEngine> infer_engine_;
  infer_ptr<nvinfer1::IExecutionContext> infer_context_;
};  // class TensorRTEngine

// Add an layer__ into engine__ with args ARGS.
// For example:
//   TRT_ENGINE_ADD_LAYER(xxx, FullyConnected, input, dim, weights, bias)
//
// Reference
// https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#charRNN_define_network
//
// will add a fully connected layer into the engine.
// TensorRT has too many layers, so that is not wise to add member functions for
// them, and an macro like this is more extensible when underlying TensorRT
// library add new layer supports.
#define TRT_ENGINE_ADD_LAYER(engine__, layer__, ARGS...) \
  engine__->network()->add##layer__(ARGS);

175 176 177 178 179
/*
 * Helper to control the TensorRT engine's creation and deletion.
 */
class TRT_EngineManager {
 public:
Y
Yan Chunwei 已提交
180 181 182 183 184 185 186 187 188 189 190 191 192 193 194
  bool HasEngine(const std::string& name) const {
    return engines_.count(name) != 0;
  }

  // Get an engine called `name`.
  TensorRTEngine* Get(const std::string& name) const {
    return engines_.at(name).get();
  }

  // Create or get an engine called `name`
  TensorRTEngine* Create(int max_batch, int max_workspace, cudaStream_t* stream,
                         const std::string& name) {
    auto* p = new TensorRTEngine(max_batch, max_workspace, stream);
    engines_[name].reset(p);
    return p;
195 196 197
  }

  void DeleteALl() {
Y
Yan Chunwei 已提交
198 199
    for (auto& item : engines_) {
      item.second.reset(nullptr);
200 201 202 203
    }
  }

 private:
Y
Yan Chunwei 已提交
204
  std::unordered_map<std::string, std::unique_ptr<TensorRTEngine>> engines_;
205 206
};

Y
Yan Chunwei 已提交
207 208 209
}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle