Method GetEngine

tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc:665–765

}

StatusOr<EngineContext*> TRTEngineOp::GetEngine(
    const std::vector<TensorShape>& input_shapes, OpKernelContext* ctx,
    TRTEngineCacheResource* cache_res) {
  static EngineContext empty_context;

  mutex_lock lock(engine_mutex_);
  // Using first input to get batch size is reliable - VerifyInputShapes() has
  // verified that.
  const int batch_size = input_shapes[0].dim_size(0);
  auto& cache = cache_res->cache_;
  auto allocator = cache_res->allocator_.get();
  if (allocator == nullptr) {
    return &empty_context;
  }

  // Handle the static engine case. For static engines, the cache will have a
  // single element containing the only engine.
  if (static_engine_) {
    if (cache.size()) {
      if (AreShapesCompatible(input_shapes, cache.begin()->first)) {
        return cache.begin()->second.get();
      }
      return &empty_context;
    }

    TrtUniquePtrType<IRuntime> infer(nvinfer1::createInferRuntime(logger));
    infer->setGpuAllocator(allocator);
    TrtUniquePtrType<nvinfer1::ICudaEngine> static_engine(
        infer->deserializeCudaEngine(serialized_segment_.c_str(),
                                     serialized_segment_.size(), nullptr));
    auto raw_static_engine = static_engine.get();
    const auto max_batch_size = raw_static_engine->getMaxBatchSize();
    // Static engine will have max_batch_size for batch size so that all inputs
    // will map to this single engine.
    std::vector<TensorShape> engine_input_shapes(input_shapes);
    for (int i = 0; i < engine_input_shapes.size(); i++) {
      // TODO(tmorris): will all inputs have batch size as first dimension??
      engine_input_shapes[i].set_dim(0, max_batch_size);
    }
    // TODO(laigd): here we assume engine_input_shapes matches the actual input
    // shapes of the engine, we should verify that.
    cache.emplace(engine_input_shapes,
                  absl::make_unique<EngineContext>(
                      std::move(static_engine),
                      TrtUniquePtrType<nvinfer1::IExecutionContext>(
                          raw_static_engine->createExecutionContext())));
    // Runtime is safe to delete after engine creation
    VLOG(1) << "Size of serialized TRT engine: "
            << serialized_segment_.capacity();
    string tmp;
    // Swap with temporary empty string to deallocate the CPU memory.
    serialized_segment_.swap(tmp);
    if (max_batch_size < batch_size) {
      return &empty_context;
    }
    return cache.at(engine_input_shapes).get();
  }  // static_engine_