engine.h 8.8 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>
N
nhzlx 已提交
22
#include "paddle/fluid/framework/tensor.h"
Y
Yan Chunwei 已提交
23 24
#include "paddle/fluid/inference/engine.h"
#include "paddle/fluid/inference/tensorrt/helper.h"
25
#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
26
#include "paddle/fluid/inference/utils/singleton.h"
Y
Yan Chunwei 已提交
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42

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:
43
    Weight(nvinfer1::DataType dtype, void* value, size_t num_elem) {
Y
Yan Chunwei 已提交
44 45 46 47 48 49
      w_.type = dtype;
      w_.values = value;
      w_.count = num_elem;
    }
    const nvinfer1::Weights& get() { return w_; }

50 51
    std::vector<int64_t> dims;

Y
Yan Chunwei 已提交
52 53 54 55
   private:
    nvinfer1::Weights w_;
  };

Y
Yan Chunwei 已提交
56
  TensorRTEngine(int max_batch, int max_workspace,
N
nhzlx 已提交
57
                 cudaStream_t* stream = nullptr, int device = 0,
Y
Yan Chunwei 已提交
58 59 60
                 nvinfer1::ILogger& logger = NaiveLogger::Global())
      : max_batch_(max_batch),
        max_workspace_(max_workspace),
Y
Yan Chunwei 已提交
61
        stream_(stream ? stream : &default_stream_),
N
nhzlx 已提交
62 63 64 65
        logger_(logger),
        device_(device) {
    freshDeviceId();
    cudaStreamCreate(stream_);
66
  }
Y
Yan Chunwei 已提交
67 68 69 70

  virtual ~TensorRTEngine();

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

73
  void Execute(int batch_size) override;
Y
Yan Chunwei 已提交
74 75 76 77

  // Initialize the inference network, so that TensorRT layers can add to this
  // network.
  void InitNetwork() {
78
    infer_builder_.reset(createInferBuilder(&logger_));
Y
Yan Chunwei 已提交
79 80 81 82 83 84 85 86 87 88 89 90 91 92
    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 已提交
93 94
  // Set the itensor_map_[name] as the network's output, and set its name.
  void DeclareOutput(const std::string& name);
N
nhzlx 已提交
95 96
  // Check if the ITensor has been declared
  bool HasDeclared(const std::string& name);
Y
Yan Chunwei 已提交
97 98 99 100 101

  // 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 已提交
102 103 104
  Buffer& buffer(const std::string& name) override;

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

  // Fill an input from CPU memory with name and size.
107
  void SetInputFromCPU(const std::string& name, const void* data, size_t size);
Y
Yan Chunwei 已提交
108 109
  // TODO(Superjomn) is this method necessary given that buffer(xxx) can be
  // accessed directly. Fill an input from GPU memory with name and size.
110
  void SetInputFromGPU(const std::string& name, const void* data, size_t size);
Y
Yan Chunwei 已提交
111
  // Get an output called name, the output of tensorrt is in GPU, so this method
112
  // Return the output's GPU memory address without copy.
Y
Yan Chunwei 已提交
113
  void* GetOutputInGPU(const std::string& name);
114
  // Copy data into dst inside the GPU device.
N
nhzlx 已提交
115
  void GetOutputInGPU(const std::string& name, void* dst, size_t max_size);
Y
Yan Chunwei 已提交
116 117
  // LOW EFFICENCY! Get output to CPU, this will trigger a memory copy from GPU
  // to CPU.
N
nhzlx 已提交
118
  void GetOutputInCPU(const std::string& name, void* dst, size_t max_size);
L
Luo Tao 已提交
119 120 121 122
  // 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 已提交
123 124 125

  nvinfer1::ICudaEngine* engine() { return infer_engine_.get(); }
  nvinfer1::INetworkDefinition* network() { return infer_network_.get(); }
126 127
  void SetRuntimeBatch(size_t batch_size);
  int GetRuntimeBatch();
N
nhzlx 已提交
128
  int GetDevice() { return device_; }
129 130
  nvinfer1::IPluginLayer* addPlugin(nvinfer1::ITensor* const* inputs,
                                    int nbInputs, PluginTensorRT*);
N
nhzlx 已提交
131 132 133 134 135 136 137 138

  // A pointer to CPU memory is needed of the TRT weight.
  // Before TRT runs, fluid loads weight into GPU storage.
  // so we need to copy the weights from GPU to CPU in our op converter.
  // We use a map to store these weights for the weight memory is not released
  // in advance, which affecting the construction of TRT Op.
  std::unordered_map<std::string /*name*/, std::unique_ptr<framework::Tensor>>
      weight_map;
Y
Yan Chunwei 已提交
139

140
  // TODO(NHZLX)
N
nhzlx 已提交
141 142 143 144 145 146 147 148 149
  // In the normal case, the paddle-trt exists bug when runing the googlenet.
  // When there are more than two convolutions of 1 * 1 with the same input, the
  // paddle-tensorrt will do the merging optimization, which fuse those conv
  // into
  // one conv, and then trigger bug. So,  We should use strategy to avoid this
  // optimization for the time being. This bug will be fixed in the future.
  std::unordered_map<std::string /*name*/, int /*ITensor_quote_num*/>
      itensor_quote_num;

Y
Yan Chunwei 已提交
150 151 152
 private:
  // the max batch size
  int max_batch_;
153 154
  // the runtime batch size
  static int runtime_batch_;
Y
Yan Chunwei 已提交
155 156
  // the max memory size the engine uses
  int max_workspace_;
157 158 159

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

Y
Yan Chunwei 已提交
165
  std::vector<Buffer> buffers_;
Y
Yan Chunwei 已提交
166 167
  // max data size for the buffers.
  std::unordered_map<std::string /*name*/, size_t /*max size*/> buffer_sizes_;
L
Luo Tao 已提交
168 169
  std::unordered_map<std::string /*name*/, nvinfer1::ITensor* /*ITensor*/>
      itensor_map_;
170

N
nhzlx 已提交
171 172
  // The specific GPU id that the TensorRTEngine bounded to.
  int device_;
173
  std::vector<std::unique_ptr<PluginTensorRT>> owned_plugin_;
Y
Yan Chunwei 已提交
174 175 176 177

  // TensorRT related internal members
  template <typename T>
  struct Destroyer {
178 179 180 181 182
    void operator()(T* x) {
      if (x) {
        x->destroy();
      }
    }
Y
Yan Chunwei 已提交
183 184 185 186 187 188 189
  };
  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_;
N
nhzlx 已提交
190 191 192 193
  // Each ICudaEngine object is bound to a specific GPU when it is instantiated,
  // ensure that the thread is associated with the correct device by calling
  // freshDeviceId().
  void freshDeviceId();
Y
Yan Chunwei 已提交
194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209
};  // 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);

210 211 212 213 214
/*
 * Helper to control the TensorRT engine's creation and deletion.
 */
class TRT_EngineManager {
 public:
Y
Yan Chunwei 已提交
215 216 217 218 219 220 221 222 223 224 225
  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,
N
nhzlx 已提交
226 227
                         const std::string& name, int gpu_device = 0) {
    auto* p = new TensorRTEngine(max_batch, max_workspace, stream, gpu_device);
Y
Yan Chunwei 已提交
228 229
    engines_[name].reset(p);
    return p;
230 231 232
  }

  void DeleteALl() {
Y
Yan Chunwei 已提交
233 234
    for (auto& item : engines_) {
      item.second.reset(nullptr);
235 236 237 238
    }
  }

 private:
Y
Yan Chunwei 已提交
239
  std::unordered_map<std::string, std::unique_ptr<TensorRTEngine>> engines_;
240 241
};

Y
Yan Chunwei 已提交
242 243 244
}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle