android-15.0.0_r1/s

/* Copyright 2019 The TensorFlow 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.
==============================================================================*/

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/inlined_vector.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/tf2tensorrt/common/datavec.h"
#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_instance.pb.h"  // NOLINT
#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT

namespace tensorflow {
namespace tensorrt {

struct TestParam {
  nvinfer1::Dims dims;
  bool dynamic_shape;
  int n_inputs;
};

class TRTEngineResourceOpsTest
    : public OpsTestBase,
      public ::testing::WithParamInterface<TestParam> {
 public:
  TRTEngineResourceOpsTest() : param_(GetParam()) {}

 protected:
  void Reset() {
    for (auto& temp : tensors_) {
      delete temp;
    }
    for (auto& temp : managed_outputs_) {
      delete temp;
    }
    tensors_.clear();
    managed_outputs_.clear();
    inputs_.clear();
  }

  ITensorProxyPtr NetworkWith1Input(nvinfer1::INetworkDefinition* network,
                                    ITensorProxyPtr input) {
    // Add a unary layer.
    nvinfer1::IUnaryLayer* layer =
        network->addUnary(*input->trt_tensor(), nvinfer1::UnaryOperation::kEXP);
    EXPECT_NE(nullptr, layer);
    return layer->getOutput(0);
  }

  // Constructs a network with two inputs, where the second input is a shape
  // tensor. We take a slice of the first input with the size of the slice
  // specified by the second input, assuming the first input is a 2D tensor.
  // We then add the slice to itself to produce the output of the network.
  ITensorProxyPtr NetworkWith2Inputs(nvinfer1::INetworkDefinition* network,
                                     ITensorProxyPtr input) {
    nvinfer1::Dims dims2{1, {2}};
    ITensorProxyPtr input2 =
        network->addInput(absl::StrCat(IONamePrefixes::kInputPHName, 1).c_str(),
                          nvinfer1::DataType::kINT32, dims2);
    EXPECT_NE(nullptr, input2->trt_tensor());

    nvinfer1::Dims start{2, {0, 0}};
    nvinfer1::Dims stride{2, {1, 1}};
    auto slice_layer =
        network->addSlice(*input->trt_tensor(), start, stride, stride);
    EXPECT_NE(nullptr, slice_layer);

    slice_layer->setInput(2, *input2->trt_tensor());
    ITensorProxyPtr sliced_input = slice_layer->getOutput(0);
    EXPECT_NE(nullptr, sliced_input->trt_tensor());

    auto layer = network->addElementWise(*sliced_input->trt_tensor(),
                                         *sliced_input->trt_tensor(),
                                         nvinfer1::ElementWiseOperation::kSUM);
    EXPECT_NE(nullptr, layer);
    return layer->getOutput(0);
  }

  TrtUniquePtrType<nvinfer1::ICudaEngine> CreateTRTEngine() {
    TrtUniquePtrType<nvinfer1::IBuilder> builder(
        nvinfer1::createInferBuilder(logger_));
    TrtUniquePtrType<nvinfer1::INetworkDefinition> network;
#if IS_TRT_VERSION_GE(8, 0, 0, 0)
    network =
        TrtUniquePtrType<nvinfer1::INetworkDefinition>(builder->createNetworkV2(
            1U << static_cast<int>(
                nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
#else
    network =
        TrtUniquePtrType<nvinfer1::INetworkDefinition>(builder->createNetworkV2(
            1U << static_cast<int>(
                nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
#endif

    // Add the input.
    nvinfer1::Dims dims = this->param_.dims;
    if (this->param_.dynamic_shape) {
      std::fill(dims.d, dims.d + dims.nbDims, -1);
    }
    const std::string in_name = StrCat(IONamePrefixes::kInputPHName, 0);
    ITensorProxyPtr input =
        network->addInput(in_name.c_str(), nvinfer1::DataType::kFLOAT, dims);
    EXPECT_NE(nullptr, input->trt_tensor());
    // Mark the output.
    ITensorProxyPtr output =
        this->param_.n_inputs == 1
            ? this->NetworkWith1Input(network.get(), input)
            : this->NetworkWith2Inputs(network.get(), input);
    output->setName("output");
    network->markOutput(*output->trt_tensor());

    // Build the engine
    TrtUniquePtrType<nvinfer1::IBuilderConfig> builder_config(
        builder->createBuilderConfig());
    builder_config->setMaxWorkspaceSize(1 << 10);
    builder->setMaxBatchSize(1);

    if (this->param_.dynamic_shape) {
      TrtShapeOptimizationProfile profile;
      profile.SetShapeTensorMask(network.get());
      const int n_input = param_.n_inputs;
      // Set the input mask to true (no resource input)
      std::vector<bool> input_mask(n_input, true);
      profile.SetInputMask(input_mask);
      // The for loop defines three optimization profiles for the network.
      for (int i = 1; i <= 3; i++) {
        std::vector<TensorShape> shape_vec(n_input);
        // Define a shape with all dimensions set to 3*i.
        std::vector<int> dimvec(this->param_.dims.nbDims, 3 * i);
        TensorShape shape;
        TF_CHECK_OK(
            TensorShapeUtils::MakeShape(dimvec.data(), dimvec.size(), &shape));

        const ITensorProxyPtr input = network->getInput(0);
        const char* name = input->getName();
        VLOG(2) << "Defining profile for input " << name;
        shape_vec[0] = shape;
        if (this->param_.n_inputs == 2) {
          // The shape of the shape tensor.
          TF_CHECK_OK(TensorShapeUtils::MakeShape(
              std::vector<int32>{param_.dims.nbDims}, &shape));
          shape_vec[1] = shape;
          // Values of the shape tensor
          Tensor shape_tensor(DT_INT32, shape);
          // Define shape values {1, i}, where 1 is the value of the first dim,
          // and i is the value of the second dimension.
          std::vector<int32> vals{1, i};
          std::copy_n(vals.data(), vals.size(),
                      shape_tensor.flat<int32_t>().data());
          DataVec shape_values{{"one", {}}, {"two", shape_tensor}};
          TF_CHECK_OK(profile.CollectShapeValues(shape_values));
        } else {
          TF_CHECK_OK(profile.CollectShapeValues({{"one", {}}}));
        }
        profile.AddShape(shape_vec);
      }
      std::vector<PartialTensorShape> input_partial_shapes;
      TF_CHECK_OK(GetNetworkInputShapes(network.get(), &input_partial_shapes));
      profile.InitProfiles(input_partial_shapes,
                           ProfileStrategy::kImplicitBatchModeCompatible);
      // Configure and build engine
      TF_CHECK_OK(profile.ConfigureBuilder(builder.get(), builder_config.get(),
                                           network.get()));
    }
    VLOG(2) << "ConfigureBuilder Finished";
    TrtUniquePtrType<nvinfer1::ICudaEngine> engine(
        builder->buildEngineWithConfig(*network, *builder_config));
    VLOG(2) << "Engine constructed";
    EXPECT_NE(nullptr, engine);
    return engine;
  }
  Logger& logger_ = *Logger::GetLogger();
  TestParam param_;
};

#if IS_TRT_VERSION_GE(7, 1, 3, 0)
constexpr std::array<TestParam, 3> TestParameters = {
    TestParam{nvinfer1::Dims{1, {1}}, false, 1},
    TestParam{nvinfer1::Dims{1, {1}}, true, 1},
    TestParam{nvinfer1::Dims{2, {3, 3}}, true, 2}};
#else
constexpr std::array<TestParam, 2> TestParameters = {
    TestParam{nvinfer1::Dims{1, {1}}, false, 1},
    TestParam{nvinfer1::Dims{1, {1}}, true, 1}};
#endif

INSTANTIATE_TEST_CASE_P(EngineResourceOpsTestInstantiation,
                        TRTEngineResourceOpsTest,
                        ::testing::ValuesIn(TestParameters));

TEST_P(TRTEngineResourceOpsTest, Basic) {
  // Create the GPU device.
  std::unique_ptr<Device> device(
      DeviceFactory::NewDevice("GPU", {}, "/job:worker/replica:0/task:0"));
  ResourceMgr* rm = device->resource_manager();
  SetDevice(DEVICE_GPU, std::move(device));

  // Create a resource handle.
  const string container(kTfTrtContainerName);
  const string resource_name = "myresource";
  Reset();
  TF_ASSERT_OK(NodeDefBuilder("op", "CreateTRTResourceHandle")
                   .Attr("resource_name", resource_name)
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  TF_ASSERT_OK(RunOpKernel());
  ResourceHandle handle =
      context_->mutable_output(0)->scalar<ResourceHandle>()();

  // Check that a resource hasn't been created yet.
  TRTEngineCacheResource* resource = nullptr;
  EXPECT_TRUE(
      errors::IsNotFound(rm->Lookup(container, resource_name, &resource)));

  // Create a resource and use an empty file to initialize the resource.
  Reset();
  Env* env = Env::Default();
  const string filename = io::JoinPath(testing::TmpDir(), "trt_engine_file");
  {
    std::unique_ptr<WritableFile> file;
    TF_ASSERT_OK(env->NewWritableFile(filename, &file));
  }
  TF_ASSERT_OK(NodeDefBuilder("op", "InitializeTRTResource")
                   .Input(FakeInput(DT_RESOURCE))
                   .Input(FakeInput(DT_STRING))
                   .Attr("max_cached_engines_count", 1)
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
  AddInputFromArray<tstring>(TensorShape({}), {filename});
  TF_ASSERT_OK(RunOpKernel());

  // Check that the resource is registered with the resource manager and the
  // cache of the resource is empty.
  EXPECT_TRUE(rm->Lookup(container, resource_name, &resource).ok());
  EXPECT_EQ(0, resource->cache_.size());

  // Create an engine and add it to the cache of the resource.
  TrtUniquePtrType<nvinfer1::ICudaEngine> engine = CreateTRTEngine();
  ExecutionContext context = ExecutionContext::Create(engine.get());

  std::vector<TensorShape> engine_input_shape(1);
  TF_ASSERT_OK(DimsAdapter(param_.dims).TensorShape(&(engine_input_shape[0])));
  if (param_.n_inputs > 1) {
    engine_input_shape.push_back(TensorShape({1, 1}));
  }
  resource->cache_.emplace(
      engine_input_shape,
      std::make_unique<EngineContext>(std::move(engine), std::move(context)));
  // Check that the resource has multiple references before it is unregistered
  // from the resource manager.
  EXPECT_FALSE(resource->RefCountIsOne());

  // Serialize the engine to a file and unregistered the resource from the
  // resource manager.
  Reset();
  TF_ASSERT_OK(NodeDefBuilder("op", "SerializeTRTResource")
                   .Attr("delete_resource", true)
                   .Input(FakeInput(DT_STRING))
                   .Input(FakeInput(DT_STRING))
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<tstring>(TensorShape({}), {resource_name});
  AddInputFromArray<tstring>(TensorShape({}), {filename});
  TF_ASSERT_OK(RunOpKernel());
  // Check that the resource now has only one reference. Detach the reference
  // to the resource to destroy the resource.
  EXPECT_TRUE(resource->RefCountIsOne());
  resource->Unref();

  // Check that unregistering the resource from the resource manager returns
  // an error as the resource has already been unregistered.
  Reset();
  TF_ASSERT_OK(NodeDefBuilder("op", "DestroyResourceOp")
                   .Attr("ignore_lookup_error", false)
                   .Input(FakeInput(DT_RESOURCE))
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
  EXPECT_TRUE(errors::IsNotFound(RunOpKernel()));

  // Verify the file for the serialized engine.
  std::unique_ptr<RandomAccessFile> file;
  TF_ASSERT_OK(env->NewRandomAccessFile(filename, &file));
  auto reader = std::make_unique<io::RecordReader>(file.get());
  uint64 offset = 0;
  tstring record;
  TF_ASSERT_OK(reader->ReadRecord(&offset, &record));
  TRTEngineInstance engine_instance;
  engine_instance.ParseFromString(record);
  EXPECT_EQ(param_.n_inputs, engine_instance.input_shapes_size());
  EXPECT_EQ(param_.dims.nbDims, engine_instance.input_shapes(0).dim_size());
  for (int i = 0; i < param_.dims.nbDims; i++) {
    EXPECT_EQ(param_.dims.d[i], engine_instance.input_shapes(0).dim(i).size());
  }
  EXPECT_TRUE(errors::IsOutOfRange(reader->ReadRecord(&offset, &record)));

  // Recreate the resource and use the file with the serialized engine to
  // initialize the resource.
  Reset();
  TF_ASSERT_OK(NodeDefBuilder("op", "InitializeTRTResource")
                   .Input(FakeInput(DT_RESOURCE))
                   .Input(FakeInput(DT_STRING))
                   .Attr("max_cached_engines_count", 1)
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
  AddInputFromArray<tstring>(TensorShape({}), {filename});
  TF_ASSERT_OK(RunOpKernel());

  // Check that the resource is registered with the resource manager again and
  // the cache of the resource is not empty.
  EXPECT_TRUE(rm->Lookup(container, resource_name, &resource).ok());
  EXPECT_EQ(1, resource->cache_.size());
  if (this->param_.dynamic_shape) {
    EXPECT_EQ(3, resource->profiles_.GetNumProfiles());
    EXPECT_EQ(3, resource->cache_.begin()->second->GetNumContexts());

    if (this->param_.n_inputs == 1) {
      // Check if profiles are restored correctly.
      std::vector<TensorShape> shapes(1);
      // We create a shape vector that matches only profile 1.
      TF_CHECK_OK(
          TensorShapeUtils::MakeShape(std::vector<int32>{6}, &shapes[0]));
      EXPECT_EQ(1, resource->profiles_.GetProfileNumber(shapes));
    } else {
      // Check if shape values are restored corretly.
      std::vector<TensorShape> shapes(2);
      // We create a shape vector that matches only profile 2.
      TF_CHECK_OK(
          TensorShapeUtils::MakeShape(std::vector<int32>{9, 9}, &shapes[0]));
      TF_CHECK_OK(
          TensorShapeUtils::MakeShape(std::vector<int32>{2}, &shapes[1]));
      Tensor shape_tensor(DT_INT32, shapes[1]);
      std::vector<int32> vals{1, 3};
      std::copy_n(vals.data(), vals.size(),
                  shape_tensor.flat<int32_t>().data());
      // DataVec names are not in used CollectShapeValues, only the order
      // matters.
      DataVec shape_values{{"one", {}}, {"two", shape_tensor}};
      TF_CHECK_OK(resource->profiles_.CollectShapeValues(shape_values));
      EXPECT_EQ(2, resource->profiles_.GetProfileNumber(shapes));
    }
  }
  // Check that the resource has multiple references before it is unregistered
  // from the resource manager.
  EXPECT_FALSE(resource->RefCountIsOne());

  // Unregister the resource from the resource manager two times, expect that
  // the second time produces an error.
  Reset();
  TF_ASSERT_OK(NodeDefBuilder("op", "DestroyResourceOp")
                   .Attr("ignore_lookup_error", false)
                   .Input(FakeInput(DT_RESOURCE))
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
  TF_ASSERT_OK(RunOpKernel());
  EXPECT_TRUE(errors::IsNotFound(RunOpKernel()));

  // Check that the resource now has only one reference. Detach the reference
  // to the resource to destroy resource.
  EXPECT_TRUE(resource->RefCountIsOne());
  resource->Unref();
}

}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT