chore: import upstream snapshot with attribution

tensorflow/compiler/tf2tensorrt/utils/py_utils.cc

@@ -0,0 +1,57 @@
+/* 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 "tensorflow/compiler/tf2tensorrt/utils/py_utils.h"
+
+#include <string>
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tsl/platform/dso_loader.h"
+#include "third_party/tensorrt/NvInfer.h"
+#endif
+
+namespace tensorflow {
+namespace tensorrt {
+
+bool IsGoogleTensorRTEnabled() {
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if TF_USE_TENSORRT_STATIC
+  LOG(INFO) << "TensorRT libraries are statically linked, skip dlopen check";
+  return true;
+#else   // TF_USE_TENSORRT_STATIC
+  auto handle_or = tsl::internal::DsoLoader::TryDlopenTensorRTLibraries();
+  if (!handle_or.ok()) {
+    LOG_WARNING_WITH_PREFIX << "Could not find TensorRT";
+  }
+  return handle_or.ok();
+#endif  // TF_USE_TENSORRT_STATIC
+#else   // GOOGLE_CUDA && GOOGLE_TENSORRT
+  return false;
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+}
+
+std::vector<std::string> GetRegisteredOpConverters() {
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+  auto* registry = tensorflow::tensorrt::convert::GetOpConverterRegistry();
+  return registry->ListRegisteredOps();
+#else
+  return {"undef"};
+#endif
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow

tensorflow/compiler/tf2tensorrt/utils/py_utils.h

@@ -0,0 +1,32 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_PY_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_PY_UTILS_H_
+
+#include <string>
+#include <vector>
+
+namespace tensorflow {
+namespace tensorrt {
+
+bool IsGoogleTensorRTEnabled();
+
+std::vector<std::string> GetRegisteredOpConverters();
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_PY_UTILS_H_

tensorflow/compiler/tf2tensorrt/utils/py_utils_wrapper.cc

@@ -0,0 +1,44 @@
+/* Copyright 2020 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 <tuple>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/py_utils.h"
+
+std::tuple<int, int, int> get_linked_tensorrt_version() {
+  return tensorflow::tensorrt::GetLinkedTensorRTVersion();
+}
+
+std::tuple<int, int, int> get_loaded_tensorrt_version() {
+  return tensorflow::tensorrt::GetLoadedTensorRTVersion();
+}
+
+PYBIND11_MODULE(_pywrap_py_utils, m, pybind11::mod_gil_not_used()) {
+  m.doc() = "_pywrap_py_utils: Various TensorRT utilities";
+  m.def("get_linked_tensorrt_version", get_linked_tensorrt_version,
+        "Return the compile time TensorRT library version as the tuple "
+        "(Major, Minor, Patch).");
+  m.def("get_loaded_tensorrt_version", get_loaded_tensorrt_version,
+        "Return the runtime time TensorRT library version as the tuple "
+        "(Major, Minor, Patch).");
+  m.def("is_tensorrt_enabled", tensorflow::tensorrt::IsGoogleTensorRTEnabled,
+        "Returns True if TensorRT is enabled.");
+  m.def("get_registered_op_converters",
+        tensorflow::tensorrt::GetRegisteredOpConverters,
+        "Return a list of registered op converters by operation name");
+}

tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc

@@ -0,0 +1,115 @@
+/* Copyright 2018 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 "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
+
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/gpus/cuda/include/cuda_runtime_api.h"
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+
+// std::align is not supported, so this method mimic its behavior.
+//
+// NOTE(aaroey): according to the TensorRT API,
+// nvinfer1::IGpuAllocator::allocate() uses uint64_t type for size and alignment
+// parameters, so here we use the same type to make it compatible.
+void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space) {
+  QCHECK_GT(alignment, 0ul) << "alignment must be greater than 0.";
+  QCHECK_EQ(0, alignment & (alignment - 1)) << "Alignment must be power of 2.";
+  QCHECK_GT(size, 0ul) << "size must be greater than 0.";
+  QCHECK(ptr) << "ptr must not be nullptr.";
+  QCHECK_GT(space, 0ul) << "space must be greater than 0.";
+  const uintptr_t ptr_val = reinterpret_cast<uintptr_t>(ptr);
+  QCHECK_GE(ptr_val + space, ptr_val) << "Provided space overflows.";
+
+  if (size > space) return nullptr;
+  const uintptr_t aligned_ptr_val = ((ptr_val + alignment - 1) & -alignment);
+  if (aligned_ptr_val > ptr_val + space - size) return nullptr;
+  ptr = reinterpret_cast<void*>(aligned_ptr_val);
+  const uintptr_t diff = aligned_ptr_val - ptr_val;
+  space -= diff;
+  return ptr;
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+
+void* TRTDeviceAllocator::allocate(uint64_t size, uint64_t alignment,
+                                   uint32_t flags) noexcept {
+  if (size == 0) return nullptr;
+  // WAR for allocator alignment requirement. Certain cuda API calls require GPU
+  // memory with alignment to cudaDeviceProp::textureAlignment.
+  // See issue #20856
+  alignment = 512;
+  assert((alignment & (alignment - 1)) == 0);  // zero or a power of 2.
+  uint64_t total_size = size + alignment;
+  // TODO(aaroey): AllocateRaw takes size_t size as input, so it'll produce
+  // unexpected result when TRT tries to allocate more bytes than size_t can
+  // carry. Fix this.
+  //
+  // Fail immediately if allocation fails, rather than waiting 10 seconds and
+  // failing then anyway.
+  // TensorRT 7 can also switch to a different algorithm for a layer if an
+  // algorithm uses too much memory. If we don't fail immediately building the
+  // engine can be *very* slow with TensorRT7 when GPU memory is limited.
+  AllocationAttributes attributes;
+  attributes.retry_on_failure = false;
+  void* mem = allocator_->AllocateRaw(alignment, total_size, attributes);
+  if (!mem) return nullptr;
+
+  void* alloc_mem = mem;
+  QCHECK(Align(alignment, size, mem, total_size));
+  mutex_lock lock(mu_);
+  if (mem != alloc_mem) {
+    QCHECK(mem_map_.insert({mem, alloc_mem}).second);
+  }
+  VLOG(2) << "Allocated " << total_size << " bytes memory @" << alloc_mem
+          << "; aligned to " << size << " bytes @" << mem << " with alignment "
+          << alignment;
+  return mem;
+}
+
+TRTDeviceAllocator::TRTDeviceAllocator(Allocator* allocator)
+    : allocator_(allocator) {
+  VLOG(1) << "Using " << allocator->Name() << " allocator from TensorFlow";
+}
+
+void TRTDeviceAllocator::free(void* memory) noexcept {
+  mutex_lock lock(mu_);
+  VLOG(2) << "Deallocating @ " << memory;
+  // allocated memory adjusted for alignment, restore the original pointer
+  if (memory) {
+    auto alloc_mem = mem_map_.find(memory);
+    if (alloc_mem != mem_map_.end()) {
+      memory = alloc_mem->second;
+      mem_map_.erase(alloc_mem->first);
+    }
+    allocator_->DeallocateRaw(memory);
+  }
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h

@@ -0,0 +1,73 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_
+
+#include <unordered_map>
+
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/platform/mutex.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+// std::align is not supported, so this function mimic its behavior.
+void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space);
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+
+class TRTBaseAllocator : public nvinfer1::IGpuAllocator {
+  // Base allocator class so we can have a virtual destructor;
+ public:
+  // python wrapper seems to be not happy with an pure virtual destructor;
+  virtual ~TRTBaseAllocator() = default;
+};
+
+class TRTDeviceAllocator : public TRTBaseAllocator {
+  // Allocator implementation wrapping TF device allocators.
+ public:
+  TRTDeviceAllocator(Allocator* allocator);
+
+  // TODO(aaroey): base class doesn't have a virtual destructor, work with
+  // Nvidia to fix it.
+  virtual ~TRTDeviceAllocator() {
+    VLOG(1) << "Destroying allocator attached to " << allocator_->Name();
+  }
+  void* allocate(uint64_t size, uint64_t alignment,
+                 uint32_t flags) noexcept override;
+  void free(void* memory) noexcept override;
+
+ private:
+  mutex mu_;
+  Allocator* allocator_;
+
+  // supporting alignment from allocation request requires a map to free;
+  std::unordered_map<void*, void*> mem_map_ TF_GUARDED_BY(mu_);
+};
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_

tensorflow/compiler/tf2tensorrt/utils/trt_allocator_test.cc

@@ -0,0 +1,87 @@
+/* Copyright 2018 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 "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
+
+#include <cstdint>
+
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+bool RunTest(const uint64_t alignment, const uint64_t size,
+             const intptr_t orig_ptr_val, const uint64_t orig_space) {
+  void* const orig_ptr = reinterpret_cast<void*>(orig_ptr_val);
+  void* ptr = orig_ptr;
+  uint64_t space = orig_space;
+  void* result = Align(alignment, size, ptr, space);
+  if (result == nullptr) {
+    EXPECT_EQ(orig_ptr, ptr);
+    EXPECT_EQ(orig_space, space);
+    return false;
+  } else {
+    EXPECT_EQ(result, ptr);
+    const intptr_t ptr_val = reinterpret_cast<intptr_t>(ptr);
+    EXPECT_EQ(0, ptr_val % alignment);
+    EXPECT_GE(ptr_val, orig_ptr_val);
+    EXPECT_GE(space, size);
+    EXPECT_LE(space, orig_space);
+    EXPECT_EQ(ptr_val + space, orig_ptr_val + orig_space);
+    return true;
+  }
+}
+
+TEST(TRTAllocatorTest, Align) {
+  for (const uint64_t space :
+       {1ul, 2ul, 3ul, 4ul, 7ul, 8ul, 9ul, 10ul, 16ul, 32ul, 511ul, 512ul,
+        513ul, 700ul, 12345ul, 1ul << 32}) {
+    for (uint64_t alignment = 1; alignment <= space * 4; alignment *= 2) {
+      for (const uintptr_t ptr_val :
+           {static_cast<uint64_t>(1),
+            alignment == 1 ? static_cast<uint64_t>(1) : alignment - 1,
+            alignment, alignment + 1, alignment + (alignment / 2)}) {
+        if (ptr_val % alignment == 0) {
+          for (const uint64_t size :
+               {static_cast<uint64_t>(1),
+                space == 1 ? static_cast<uint64_t>(1) : space - 1, space,
+                space + 1}) {
+            EXPECT_EQ(space >= size, RunTest(alignment, size, ptr_val, space));
+          }
+        } else {
+          EXPECT_FALSE(RunTest(alignment, space, ptr_val, space));
+          const uint64_t diff = alignment - ptr_val % alignment;
+          if (space > diff) {
+            EXPECT_TRUE(
+                RunTest(alignment, space - diff, ptr_val + diff, space - diff));
+            for (const uint64_t size :
+                 {static_cast<uint64_t>(1),
+                  space - diff > 1 ? space - diff - 1
+                                   : static_cast<uint64_t>(1),
+                  space - diff, space - diff + 1, space - 1}) {
+              EXPECT_EQ(space - diff >= size,
+                        RunTest(alignment, size, ptr_val, space));
+            }
+          } else {
+            EXPECT_FALSE(RunTest(alignment, 1, ptr_val, space));
+          }
+        }
+      }
+    }
+  }
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow

tensorflow/compiler/tf2tensorrt/utils/trt_engine_instance.proto

@@ -0,0 +1,34 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tensorrt;
+
+import "tensorflow/core/framework/tensor_shape.proto";
+
+// Containing information for a serialized TensorRT engine.
+message TRTEngineInstance {
+  // The input shapes of the TRT engine.
+  repeated TensorShapeProto input_shapes = 1;
+
+  // The serialized TRT engine.
+  //
+  // TODO(laigd): consider using a more efficient in-memory representation
+  // instead of string which is the default here.
+  bytes serialized_engine = 2;
+
+  // TODO(laigd): consider adding calibration stats, precision_modes, etc.
+}

tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc

@@ -0,0 +1,286 @@
+/* Copyright 2020 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 "tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h"
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_execution_context.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/profiler/lib/traceme.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+using absl::StrCat;
+
+ExecutionContext ExecutionContext::Create(nvinfer1::ICudaEngine* cuda_engine) {
+  bool has_int32_output = false;
+  for (int i = 0; i < cuda_engine->getNbBindings(); i++) {
+    if (!cuda_engine->bindingIsInput(i) &&
+        cuda_engine->getBindingDataType(i) == nvinfer1::DataType::kINT32) {
+      has_int32_output = true;
+      break;
+    }
+  }
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && has_int32_output) {
+    // TODO(nvbugs/3390469): Remove this workaround when the bug is fixed.
+    nvinfer1::IExecutionContext* execution_context =
+        cuda_engine->createExecutionContext();
+    return ExecutionContext(execution_context, true);
+  }
+
+  nvinfer1::IExecutionContext* execution_context =
+      cuda_engine->createExecutionContextWithoutDeviceMemory();
+  return ExecutionContext(execution_context, false);
+}
+
+Status GetTrtBindingShape(const nvinfer1::ICudaEngine* cuda_engine,
+                          const nvinfer1::IExecutionContext* execution_context,
+                          int binding_index, bool use_implicit_batch,
+                          int batch_size, TensorShape& shape) {
+  tensorflow::profiler::TraceMe activity(
+      "getBindingDimensions", tensorflow::profiler::TraceMeLevel::kInfo);
+  nvinfer1::Dims dims =
+      use_implicit_batch
+          ? cuda_engine->getBindingDimensions(binding_index)
+          : execution_context->getBindingDimensions(binding_index);
+  if (!use_implicit_batch) {
+    if (dims.nbDims == -1) {
+      return errors::Internal(
+          "Binding index out of range. This can happen if profile is not set, "
+          "or the network is invalid for the current profile.");
+    }
+  }
+  TF_RETURN_IF_ERROR(DimsAdapter(dims).TensorShape(
+      &shape,
+      use_implicit_batch ? std::optional<int>(batch_size) : std::nullopt));
+  return OkStatus();
+}
+
+Status SetupBindings(nvinfer1::ICudaEngine* cuda_engine, const Tensor& tensor,
+                     std::vector<void*>& buffers, int binding_index) {
+  tensorflow::profiler::TraceMe activity(
+      "SetBindingPointers", tensorflow::profiler::TraceMeLevel::kInfo);
+  const auto dtype = cuda_engine->getBindingDataType(binding_index);
+  VLOG(2) << "<<<<<<<<< SetupBindings with dtype = " << (int)dtype;
+  switch (dtype) {
+    case nvinfer1::DataType::kFLOAT:
+      buffers[binding_index] = const_cast<float*>(tensor.flat<float>().data());
+      break;
+    case nvinfer1::DataType::kHALF:
+      buffers[binding_index] =
+          const_cast<Eigen::half*>(tensor.flat<Eigen::half>().data());
+      break;
+    case nvinfer1::DataType::kINT8:
+      return errors::Internal("INT8 inputs are not supported yet!");
+    case nvinfer1::DataType::kINT32:
+      buffers[binding_index] = const_cast<int32*>(tensor.flat<int32>().data());
+      break;
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    case nvinfer1::DataType::kBOOL:
+      buffers[binding_index] = const_cast<bool*>(tensor.flat<bool>().data());
+      break;
+#endif
+#if IS_TRT_VERSION_GE(8, 5, 0, 0)
+    case nvinfer1::DataType::kUINT8:
+      buffers[binding_index] = const_cast<uint8*>(tensor.flat<uint8>().data());
+      break;
+#endif
+#if IS_TRT_VERSION_GE(8, 6, 0, 0)
+    case nvinfer1::DataType::kFP8:
+      return errors::Internal("FP8 inputs are not supported yet!");
+#endif
+    default:
+      return errors::Internal("Unknown TRT data type: ",
+                              static_cast<int>(dtype));
+  }
+  return OkStatus();
+}
+
+// Sets up bindings.
+Status SetTrtEngineInputs(nvinfer1::ICudaEngine* cuda_engine,
+                          nvinfer1::IExecutionContext* execution_context,
+                          const int trt_profile_idx,
+                          std::vector<void*>& buffers, bool use_implicit_batch,
+                          int num_batch,
+                          const TrtShapeOptimizationProfile& profiles,
+                          OpKernelContext* ctx, const DataVec* input_vec) {
+  tensorflow::profiler::TraceMe activity(
+      "SetTrtEngineInputs", tensorflow::profiler::TraceMeLevel::kInfo);
+  int n_inputs = ctx ? ctx->num_inputs() : (input_vec ? input_vec->size() : 0);
+  // Setup engine inputs.
+  for (int i = 0; i < n_inputs; i++) {
+    const Tensor& input_tensor = ctx ? ctx->input(i) : input_vec->at(i).tensor;
+    const TensorShape& input_shape = input_tensor.shape();
+
+    // Skip resource inputs.
+    if (input_tensor.dtype() == DataType::DT_RESOURCE) {
+      continue;
+    }
+
+    const string input_name =
+        ctx ? StrCat(IONamePrefixes::kInputPHName, i) : input_vec->at(i).name;
+    int binding_index;
+    Status status = GetTrtBindingIndex(input_name.c_str(), trt_profile_idx,
+                                       cuda_engine, &binding_index);
+    if (IS_TRT_VERSION_GE(8, 0, 0, 0)) {
+      TF_RETURN_IF_ERROR(status);
+    } else if (!status.ok()) {
+      // Before TRT 8, an input tensor can be pruned if it is not used by the
+      // network (e.g. only its shape is used, but the shape is already defined
+      // by the optimization profile by setting min=max). nvbugs/3153064
+      VLOG(2) << "Skipping pruned input " << input_name;
+      continue;
+    }
+
+    if (use_implicit_batch && ctx) {
+      // Ensure all inputs have the same batch size
+      if (num_batch != input_shape.dim_size(0)) {
+        const string msg =
+            StrCat("Input data has inconsistent batch size: ", num_batch,
+                   " vs ", input_shape.dim_size(0));
+        return errors::NotFound(msg);
+      }
+    }
+    // Set known input dimensions. This is necessary because TRT network
+    // could be made with dynamic dimensions.
+    if (!use_implicit_batch) {
+      TF_RETURN_IF_ERROR(profiles.SetInputShapeBinding(
+          i, binding_index, cuda_engine, execution_context));
+
+      if (cuda_engine->isExecutionBinding(binding_index)) {
+        tensorflow::profiler::TraceMe activity(
+            "SetTrtEngineInputs::setBindingDimensions",
+            tensorflow::profiler::TraceMeLevel::kInfo);
+        auto adap = DimsAdapter::Create(input_shape);
+        TRT_ENSURE_OK(adap);
+        nvinfer1::Dims trt_dims = adap->AsTrtDims();
+        if (execution_context->getBindingDimensions(binding_index) !=
+            trt_dims) {
+          VLOG(2) << "Setting binding dimensions for idx " << binding_index;
+          bool ret =
+              execution_context->setBindingDimensions(binding_index, trt_dims);
+          if (!ret) {
+            VLOG(2) << "Error setting engine input " << binding_index << " "
+                    << DebugString(trt_dims);
+            return errors::Internal(
+                "Binding dimension does not fit selected profile.");
+          }
+        }
+      }
+    }
+    // Setup input bindings.
+    TF_RETURN_IF_ERROR(
+        SetupBindings(cuda_engine, input_tensor, buffers, binding_index));
+  }
+
+  // Ensure all network dynamic dimensions (if any) are set in execution
+  // context.
+  if (!execution_context->allInputDimensionsSpecified()) {
+    return errors::Internal(
+        "Failed to set dimensions for all dynamic input tensors");
+  }
+  if (!execution_context->allInputShapesSpecified()) {
+    return errors::Internal(
+        "Failed to set dimensions for all shape input tensors.");
+  }
+  return OkStatus();
+}
+
+Status SetTrtEngineOutputs(nvinfer1::ICudaEngine* cuda_engine,
+                           nvinfer1::IExecutionContext* execution_context,
+                           int trt_profile_idx, std::vector<void*>& buffers,
+                           bool use_implicit_batch, int batch_size,
+                           OpKernelContext* ctx, DataVec* outputs) {
+  tensorflow::profiler::TraceMe activity(
+      "SetTrtEngineOutputs", tensorflow::profiler::TraceMeLevel::kInfo);
+  // Either one of ctx or outpus should be specified
+  int n_outputs = ctx ? ctx->num_outputs() : (outputs ? outputs->size() : 0);
+  for (int i = 0; i < n_outputs; i++) {
+    const string output_name =
+        ctx ? StrCat(IONamePrefixes::kOutputPHName, i) : outputs->at(i).name;
+    int binding_index;
+    TF_RETURN_IF_ERROR(GetTrtBindingIndex(output_name.c_str(), trt_profile_idx,
+                                          cuda_engine, &binding_index));
+
+    // Get TRT output shapes for allocating output memory.
+    TensorShape output_shape;
+    TF_RETURN_IF_ERROR(GetTrtBindingShape(cuda_engine, execution_context,
+                                          binding_index, use_implicit_batch,
+                                          batch_size, output_shape));
+
+    // Allocate output tensor of TRTEngineOp.
+    Tensor* output_tensor = nullptr;
+    if (ctx) {
+      tensorflow::profiler::TraceMe activity(
+          "AllocateOutput", tensorflow::profiler::TraceMeLevel::kInfo);
+      TF_RETURN_IF_ERROR(ctx->allocate_output(i, output_shape, &output_tensor));
+    } else {
+      // This path is used for unit tests. The tensor is already allocated.
+      // Its shape is not necessarily set correctly, we fix that.
+      VLOG(2) << "Applying shape " << output_shape.DebugString()
+              << " on output.";
+      output_tensor = &(outputs->at(i).tensor);
+      bool status = output_tensor->CopyFrom(*output_tensor, output_shape);
+      if (!status) {
+        return errors::Internal(
+            "Buffer size (", output_tensor->NumElements(),
+            ") do not match while reshaping output tensors to shape ",
+            output_shape.DebugString());
+      }
+    }
+
+    // Set up output bindings.
+    TF_RETURN_IF_ERROR(
+        SetupBindings(cuda_engine, *output_tensor, buffers, binding_index));
+  }
+  return OkStatus();
+}
+
+Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
+                  std::vector<void*>& buffers, cudaStream_t stream,
+                  bool use_implicit_batch, int batch_size) {
+  tensorflow::profiler::TraceMe activity(
+      "TrtEnqueue", tensorflow::profiler::TraceMeLevel::kInfo);
+  bool ret = false;
+  if (use_implicit_batch) {
+    ret = execution_context->enqueue(batch_size, &buffers[0], stream, nullptr);
+    VLOG(1) << "Called IExecutionContext::enqueue";
+  } else {
+    ret = execution_context->enqueueV2(&buffers[0], stream, nullptr);
+    VLOG(1) << "Called IExecutionContext::enqueueV2";
+  }
+  if (!ret) {
+    return errors::Internal("Failed to enqueue batch for TRT engine");
+  }
+  // Synchronization will be done by TF.
+  return OkStatus();
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h

@@ -0,0 +1,82 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ENGINE_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ENGINE_UTILS_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/common/datavec.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/status.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+using ::tsl::StatusOr;
+
+// Creates a TensorRT execution context.
+ExecutionContext CreateExecutionContext(nvinfer1::ICudaEngine* cuda_engine);
+
+// Sets input buffers for TRT from a list of input tensors. The input tensors
+// are either defined by ctx or by input_vec.
+Status SetTrtEngineInputs(nvinfer1::ICudaEngine* cuda_engine,
+                          nvinfer1::IExecutionContext* execution_context,
+                          const int trt_profile_idx,
+                          std::vector<void*>& buffers, bool use_implicit_batch,
+                          int num_batch,
+                          const TrtShapeOptimizationProfile& profiles,
+                          OpKernelContext* ctx = nullptr,
+                          const DataVec* input_vec = nullptr);
+
+// Returns the shape of a binding from TensorRT.
+//
+// The binding is identified by its binding_index. The batch_size argument is
+// ignored if use_implicit_batch==false. The shape is returned in the last
+// argument.
+Status GetTrtBindingShape(const nvinfer1::ICudaEngine* cuda_engine,
+                          const nvinfer1::IExecutionContext* execution_context,
+                          int binding_index, bool use_implicit_batch,
+                          int batch_size, TensorShape& shape);
+
+// Defines output buffers for TRT. The buffers are allocated by ctx, if ctx is
+// not null. Otherwise it is expected that the outputs DataVec is not null, and
+// the Tensors in outputs are already allocated.
+Status SetTrtEngineOutputs(nvinfer1::ICudaEngine* cuda_engine,
+                           nvinfer1::IExecutionContext* execution_context,
+                           int trt_profile_idx, std::vector<void*>& buffers,
+                           bool use_implicit_batch, int batch_size = 0,
+                           OpKernelContext* ctx = nullptr,
+                           DataVec* outputs = nullptr);
+
+// Enqueues TensorRT inference job. The batch_size argument is only relevant in
+// implicit batch mode.
+Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
+                  std::vector<void*>& buffers, cudaStream_t stream,
+                  bool use_implicit_batch, int batch_size = 1);
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ENGINE_UTILS_H_

tensorflow/compiler/tf2tensorrt/utils/trt_execution_context.h

@@ -0,0 +1,43 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_EXECUTION_CONTEXT_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_EXECUTION_CONTEXT_H_
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// A wrapper for the TensorRT execution context which will destroy the TensorRT
+// execution context when the object goes out of scope.
+class ExecutionContext : public TrtUniquePtrType<nvinfer1::IExecutionContext> {
+ public:
+  ExecutionContext(nvinfer1::IExecutionContext* context, bool has_memory)
+      : TrtUniquePtrType<nvinfer1::IExecutionContext>(context),
+        has_device_memory_(has_memory) {}
+  static ExecutionContext Create(nvinfer1::ICudaEngine* cuda_engine);
+
+  bool HasDeviceMemory() { return has_device_memory_; }
+
+ private:
+  bool has_device_memory_;
+};
+
+};  // namespace tensorrt
+};  // namespace tensorflow
+#endif
+#endif

tensorflow/compiler/tf2tensorrt/utils/trt_experimental_features.cc

@@ -0,0 +1,35 @@
+/* Copyright 2022 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 "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/core/util/env_var.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+bool isExperimentalFeatureActivated(string feature_name) {
+  string envvar_str;
+  TF_CHECK_OK(
+      ReadStringFromEnvVar("TF_TRT_EXPERIMENTAL_FEATURES", "", &envvar_str));
+  return envvar_str.find(feature_name) != string::npos;
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_experimental_features.h

@@ -0,0 +1,31 @@
+/* Copyright 2022 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_EXPERIMENTAL_FEATURES_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_EXPERIMENTAL_FEATURES_H_
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+
+bool isExperimentalFeatureActivated(string feature_name);
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_EXPERIMENTAL_FEATURES_H_

tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc

@@ -0,0 +1,150 @@
+/* Copyright 2018 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 "tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h"
+
+#include <atomic>
+#include <unordered_map>
+
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/gpus/cuda/include/cuda_runtime_api.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// set the batch size before constructing the thread to execute engine
+int TRTInt8Calibrator::getBatchSize() const noexcept { return batch_size_; }
+
+TRTInt8Calibrator::TRTInt8Calibrator(
+    const std::unordered_map<string, std::pair<void*, size_t>>& dev_buffers,
+    int batch_size, string engine_name)
+    : batch_size_(batch_size),
+      done_(false),
+      dev_buffers_(dev_buffers),
+      // Make sure setBatch() waits until getBatch() is called (the first time).
+      calib_running_(true),
+      batch_is_set_(false),
+      engine_name_(engine_name) {}
+
+TRTInt8Calibrator::TRTInt8Calibrator(const string& calib_data)
+    : batch_size_(0),
+      done_(true),
+      calib_running_(false),
+      batch_is_set_(false),
+      calibration_table_(calib_data) {}
+
+bool TRTInt8Calibrator::setBatch(const std::unordered_map<string, void*>& data,
+                                 const cudaStream_t stream) {
+  mutex_lock lock(cond_mtx_);
+
+  // Wait while the queue is full or calibration is running.
+  while ((calib_running_ || batch_is_set_) && !done_) cond_.wait(lock);
+  if (done_) return false;
+  CHECK(!calib_running_ && !batch_is_set_);
+  VLOG(1) << "Set Batch Waiting finished";
+
+  // Sets the batch.
+  for (const auto& it : data) {
+    auto devptr = dev_buffers_.find(it.first);
+    if (devptr == dev_buffers_.end()) {
+      LOG(FATAL) << "FATAL " << engine_name_ << " input name '" << it.first
+                 << "' does not match with the buffer names";
+    }
+    const auto& d = devptr->second;
+
+    // TODO(sami,aaroey): Need to figure out a way to ensure synchronization
+    // between stream, perhaps using a tensor?
+    auto status = cudaMemcpyAsync(d.first, it.second, d.second,
+                                  cudaMemcpyDeviceToDevice, stream);
+    if (status != cudaSuccess) {
+      LOG(FATAL) << "cudaMemcpy " << engine_name_ << " for '" << it.first
+                 << "' failed with " << status;
+    }
+  }
+
+  // TODO(Sami, aaorey): Find an alternative way!
+  // we have to wait for the stream before returning!
+  cudaStreamSynchronize(stream);
+  batch_is_set_ = true;
+  cond_.notify_all();
+  return true;
+}
+
+bool TRTInt8Calibrator::getBatch(void** bindings, const char** names,
+                                 int num_bindings) noexcept {
+  mutex_lock lock(cond_mtx_);
+  // Notify finish of last round of calibration.
+  calib_running_ = false;
+  cond_.notify_all();
+
+  // Wait until new batch arrives
+  while ((!batch_is_set_ && !done_)) cond_.wait(lock);
+  if (done_) return false;
+
+  // Gets the batch
+  for (int i = 0; i < num_bindings; i++) {
+    auto it = dev_buffers_.find(names[i]);
+    if (it == dev_buffers_.end()) {
+      LOG(FATAL) << "Calibration engine asked for unknown tensor name '"
+                 << names[i] << "' at position " << i;
+    }
+    bindings[i] = it->second.first;
+  }
+  batch_is_set_ = false;
+  calib_running_ = true;
+  return true;
+}
+
+void TRTInt8Calibrator::waitAndSetDone() {
+  mutex_lock lock(cond_mtx_);
+  // Wait while the queue is full or calibration is running, so we don't miss
+  // the last batch.
+  while ((calib_running_ || batch_is_set_) && !done_) cond_.wait(lock);
+  if (!done_) {
+    done_ = true;
+    cond_.notify_all();
+    dev_buffers_.clear();
+  }
+}
+
+const void* TRTInt8Calibrator::readCalibrationCache(
+    std::size_t& length) noexcept {
+  if (calibration_table_.empty()) return nullptr;
+  length = calibration_table_.size();
+  return calibration_table_.data();
+}
+
+void TRTInt8Calibrator::setDone() {
+  mutex_lock lock(cond_mtx_);
+  done_ = true;
+  cond_.notify_all();
+}
+
+void TRTInt8Calibrator::writeCalibrationCache(const void* ptr,
+                                              std::size_t length) noexcept {
+  calibration_table_ = string(static_cast<const char*>(ptr), length);
+  VLOG(1) << "Got calibration data for " << engine_name_ << " @" << ptr
+          << " length=" << length;
+}
+TRTInt8Calibrator::~TRTInt8Calibrator() {
+  VLOG(1) << "Destroying calibrator for " << engine_name_;
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h

@@ -0,0 +1,102 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_
+
+#include <atomic>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+#include "tensorflow/core/platform/mutex.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "third_party/gpus/cuda/include/cuda_runtime_api.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+// This class provides a 1 element queue to match TFs push model to
+// TRTs pull model for calibration. When TRT implements a means for
+// a push calibration This class should be updated accordingly
+
+// IInt8EntropyCalibrator2 is preferred for TRT 5.1+.
+struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator2 {
+ public:
+  // Construct a calibrator for future calibration.
+  TRTInt8Calibrator(
+      const std::unordered_map<string, std::pair<void*, size_t>>& dev_buffers,
+      int batch_size, string engine_name);
+
+  // Construct a finalized calibrator where we don't need to run calibration any
+  // more, as the calibration data is provided.
+  TRTInt8Calibrator(const string& calibration_data);
+
+  ~TRTInt8Calibrator();
+
+  int getBatchSize() const noexcept override;
+
+  bool getBatch(void* bindings[], const char* names[],
+                int num_bindings) noexcept override;
+
+  // Feed calibration data to the calibrator, and return true if the data is
+  // accepted. Return false if the calibrator has been terminated.
+  bool setBatch(const std::unordered_map<string, void*>& data,
+                const cudaStream_t stream);
+
+  // Wait until the last batch is consumed by the calibrator and set done.
+  void waitAndSetDone();
+
+  // Notify that calibration is done and future batches provided by setBatch()
+  // will be ignored.
+  void setDone();
+
+  // If not null, calibration is skipped.
+  const void* readCalibrationCache(std::size_t& length) noexcept override;
+
+  void writeCalibrationCache(const void* ptr,
+                             std::size_t length) noexcept override;
+
+  const string& getCalibrationTableAsString() { return calibration_table_; }
+
+ private:
+  const int batch_size_;
+
+  // mutex for condition_variable
+  mutex cond_mtx_;
+
+  // condition variable to implement producer-consumer queue for calibration
+  condition_variable cond_;
+
+  // Is calibration finished?
+  bool done_;
+
+  // Map to keep tensorrt input buffers and sizes keyed with buffer names
+  std::unordered_map<string, std::pair<void*, size_t>> dev_buffers_;
+
+  bool calib_running_;
+  bool batch_is_set_;
+
+  string engine_name_;
+  string calibration_table_;
+};
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_

tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc

@@ -0,0 +1,132 @@
+/* Copyright 2018 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 "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_experimental_features.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+bool filter_string(string msg) {
+  // This function checks for known substrings that shall be ignored.
+
+  static const std::vector<string> substr_patterns{
+      // Automatic messages generated by TensorRT when combined with
+      // Automatic Mixed Precision - TensorRT 8.2
+      "Missing scale and zero-point for",
+      "Subnormal FP16 values detected",
+      "If this is not the desired behavior, please modify the weights",
+      "had the following issues when converted to FP16",
+      "Values less than smallest positive FP16 Subnormal value detected.",
+      // Deprecation Warnings
+      "The implicit batch dimension mode has been deprecated.",
+      "The getMaxBatchSize() function should not be used with an engine built",
+      // Input-Warnings
+      "[RemoveDeadLayers] Input Tensor input is unused or used only at",
+      "Unused Input:",
+      // Data Type Warnings
+      "Tensor DataType is determined at build time for tensors not marked as",
+      // Myelin Performance Warning in dynamic shape mode
+      "Myelin graph with multiple dynamic values may have poor performance",
+      "(# 0 (SHAPE",
+      "CUDA lazy loading is not enabled. Enabling it can significantly reduce",
+  };
+
+  for (int i = 0; i < substr_patterns.size(); i++) {
+    std::size_t is_found = msg.find(substr_patterns[i]);
+    if (is_found != string::npos) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// Use TF logging for TensorRT information
+void Logger::log(Severity severity, const char* msg) noexcept {
+  static const bool filter_messages = []() {
+    return !isExperimentalFeatureActivated("disable_logger_filtering");
+  }();
+
+  if (filter_messages && filter_string(msg)) return;
+
+  if (!isValidSeverity(severity, msg) || suppressedMsg_ & (1 << (int)severity))
+    return;
+
+  // Suppress info-level messages
+  switch (severity) {
+    case Severity::kVERBOSE:
+    case Severity::kINFO: {  // Mark TRT info messages as debug!
+      VLOG(2) << name_ << " " << msg;
+      break;
+    }
+    case Severity::kWARNING: {
+      LOG_WARNING_WITH_PREFIX << name_ << " " << msg;
+      break;
+    }
+    case Severity::kERROR: {
+      LOG(ERROR) << name_ << " " << msg;
+      break;
+    }
+    case Severity::kINTERNAL_ERROR: {
+      LOG(FATAL) << name_ << " " << msg;
+      break;
+    }
+  }
+}
+
+void Logger::suppressLoggerMsgs(Severity severity) {
+  if (isValidSeverity(severity)) {
+    suppressedMsg_ |= 1 << (int)severity;
+  }
+}
+
+void Logger::unsuppressLoggerMsgs(Severity severity) {
+  if (isValidSeverity(severity)) {
+    suppressedMsg_ &= (-1) ^ (1 << (int)severity);
+  }
+}
+
+bool Logger::isValidSeverity(Severity severity, const char* msg) noexcept {
+  switch (severity) {
+    case Severity::kVERBOSE:
+    case Severity::kINFO:
+    case Severity::kWARNING:
+    case Severity::kERROR:
+    case Severity::kINTERNAL_ERROR:
+      return true;
+  }
+  return false;
+}
+
+// static
+Logger* Logger::GetLogger() {
+  static Logger* logger = new Logger("DefaultLogger");
+  return logger;
+}
+
+REGISTER_TENSORRT_LOGGER("DefaultLogger", Logger::GetLogger());
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_logger.h

@@ -0,0 +1,50 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_
+
+#include "tensorflow/core/platform/types.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// Logger for GIE info/warning/errors
+class Logger : public nvinfer1::ILogger {
+ public:
+  Logger(string name = "DefaultLogger") : name_(name) {}
+  void log(nvinfer1::ILogger::Severity severity,
+           const char* msg) noexcept override;
+  void suppressLoggerMsgs(nvinfer1::ILogger::Severity severity);
+  void unsuppressLoggerMsgs(nvinfer1::ILogger::Severity severity);
+  void unsuppressAllLoggerMsgs() { suppressedMsg_ = 0; }
+  static Logger* GetLogger();
+
+ private:
+  bool isValidSeverity(nvinfer1::ILogger::Severity severity,
+                       const char* msg = nullptr) noexcept;
+  const string name_;
+  unsigned int suppressedMsg_ = 0;
+};
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_

tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc

@@ -0,0 +1,143 @@
+/* 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 "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
+
+#include <sstream>
+
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/mutex.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+string CalibrationContext::TerminateCalibration() {
+  mutex_lock l(mu_);
+  if (terminated_) return calibration_table_;
+
+  TRTInt8Calibrator* raw_calibrator = calibrator_.get();
+  raw_calibrator->waitAndSetDone();
+  terminated_ = true;
+
+  // At this point the calibration thread `thr_` is woken up and can
+  // transfer the ownership of `calibrator_` and `engine_` at any time, so
+  // it's not safe to use `calibrator_` below, but we can still access it
+  // using raw pointer.
+  // TODO(laigd): make TRTEngineOp::AllocateCalibrationResources() a member
+  // function of this class instead.
+
+  thr_->join();
+  calibration_table_ = raw_calibrator->getCalibrationTableAsString();
+  return calibration_table_;
+}
+
+const absl::string_view kTfTrtContainerName = "TF-TRT";
+
+Logger& TRTEngineCacheResource::GetLogger() {
+  static Logger* logger = new Logger();
+  return *logger;
+}
+
+TRTEngineCacheResource::TRTEngineCacheResource(OpKernelContext* ctx,
+                                               size_t capacity)
+    : cache_(capacity) {
+  auto device = ctx->device();
+  auto alloc = device->GetAllocator(AllocatorAttributes());
+  if (!alloc) {
+    LOG(ERROR) << "Can't find device allocator for gpu device "
+               << device->name();
+    allocator_ = nullptr;
+  } else {
+    allocator_.reset(new TRTDeviceAllocator(alloc));
+  }
+}
+
+TRTEngineCacheResource::~TRTEngineCacheResource() {
+  VLOG(1) << "Destroying TRTEngineCacheResource...";
+}
+
+string TRTEngineCacheResource::DebugString() const {
+  std::stringstream oss;
+  using std::dec;
+  using std::endl;
+  using std::hex;
+  oss << "TRTEngineCacheResource: ";
+  oss << "TRTBaseAllocator = " << hex << allocator_.get() << dec << ", ";
+  oss << "LRUCache = " << hex << &cache_ << dec << endl;
+  oss << "Containing " << cache_.size() << " entries: " << endl;
+  for (const auto& item : cache_) {
+    mutex_lock lock(item.second->mu);
+    oss << TensorShapeUtils::ShapeListString(item.first) << ": " << hex
+        << "ICudaEngine: " << item.second->GetCudaEngine() << ", "
+        << "IExecutionContext: ";
+    absl::c_for_each(
+        item.second->execution_contexts,
+        [&](const ExecutionContext& ctx) { oss << ctx.get() << ","; });
+    oss << dec << endl;
+  }
+  return oss.str();
+}
+
+EngineContext* TRTEngineCacheResource::GetEngineContext(
+    const std::vector<TensorShape>& input_shapes) {
+  EngineContext* engine_context = nullptr;
+  int64 min_matched_batch_size = std::numeric_limits<int64_t>::max();
+  for (const auto& pair : cache_) {
+    const std::vector<TensorShape>& cached_input_shapes = pair.first;
+    // This should not happen, but just for safety.
+    if (input_shapes.size() != cached_input_shapes.size()) {
+      LOG(ERROR) << "Input shape list size mismatch"
+                 << ", cached size: " << cached_input_shapes.size()
+                 << " vs. input size: " << input_shapes.size();
+    }
+    if (AreShapesCompatible(input_shapes, cached_input_shapes)) {
+      const int cached_batch_size = cached_input_shapes[0].dim_size(0);
+      if (min_matched_batch_size > cached_batch_size) {
+        min_matched_batch_size = cached_batch_size;
+        engine_context = pair.second.get();
+      }
+    }
+  }
+  return engine_context;
+}
+
+EngineContext* TRTEngineCacheResource::GetEngineContext(const int profile_id) {
+  if (profiles_.NeedProfiles() && profile_id >= profiles_.GetNumProfiles()) {
+    LOG(ERROR) << "Out of range: profile_id " << profile_id
+               << " is larger than number of profiles "
+               << profiles_.GetNumProfiles();
+    return nullptr;
+  }
+  if (cache_.size() > 1) {
+    LOG(ERROR) << "Cache is expected to have at most "
+               << "1 engine in explicit batch mode where profiles are used.";
+    return nullptr;
+  }
+  if (cache_.size() == 0) {
+    return nullptr;
+  }
+  return cache_.begin()->second.get();
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h

@@ -0,0 +1,261 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LRU_CACHE_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LRU_CACHE_H_
+
+#include <list>
+#include <thread>
+#include <unordered_map>
+
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+
+template <class Key, class Value, class HashFunction>
+class LRUCache {
+ public:
+  typedef Value value_type;
+  typedef Key key_type;
+  typedef HashFunction hasher;
+  typedef typename std::unordered_map<key_type, value_type, hasher> map_type;
+  typedef typename map_type::iterator iterator;
+  typedef typename map_type::const_iterator const_iterator;
+
+  LRUCache() : capacity_(0) {}
+  explicit LRUCache(size_t capacity) : capacity_(capacity) {}
+
+  size_t capacity() const { return capacity_; }
+
+  void reserve(size_t capacity) {
+    capacity_ = capacity;
+    DiscardOld();
+  }
+
+  size_t size() const { return objects_.size(); }
+
+  size_t count(const key_type& key) const { return objects_.count(key); }
+
+  value_type& at(const key_type& key) { return Touch(key); }
+
+  const_iterator begin() const { return objects_.begin(); }
+  const_iterator end() const { return objects_.end(); }
+
+  iterator begin() { return objects_.begin(); }
+  iterator end() { return objects_.end(); }
+
+  template <typename... Args>
+  std::pair<iterator, bool> emplace(Args&&... args) {
+    DiscardOld(1);
+    std::pair<iterator, bool> result =
+        objects_.emplace(std::forward<Args>(args)...);
+    key_type key = result.first->first;
+    if (result.second) {
+      keys_.push_front(key);
+    } else {
+      TouchNoCheck(key);  // The key must exist in this case.
+    }
+    return result;
+  }
+
+ private:
+  std::unordered_map<key_type, value_type, hasher> objects_;
+  std::list<key_type> keys_;
+  size_t capacity_;
+  value_type not_found_value_;
+
+  value_type& Touch(const key_type& key) {
+    // Check that the key exists, and let it return std::out_of_range error if
+    // not.
+    value_type& value = objects_.at(key);
+    TouchNoCheck(key);
+    return value;
+  }
+
+  void TouchNoCheck(const key_type& key) {
+    auto rank = std::find(keys_.begin(), keys_.end(), key);
+    if (rank != keys_.begin()) {
+      keys_.erase(rank);
+      keys_.push_front(key);
+    }
+  }
+
+  // Creates n free positions in cache
+  void DiscardOld(size_t n = 0) {
+    DCHECK(capacity_ >= n) << "Insufficient capacity in cache (capacity = "
+                           << capacity_ << ", requested " << n << ")";
+    while (objects_.size() > (capacity_ - n)) {
+      key_type discard_key = keys_.back();
+      keys_.pop_back();
+      objects_.erase(discard_key);
+    }
+  }
+};
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+struct EngineContext {
+  EngineContext() {}  // Creates an empty context.
+  EngineContext(TrtUniquePtrType<nvinfer1::ICudaEngine>&& cuda_engine,
+                ExecutionContext&& execution_context)
+      : cuda_engine_(std::move(cuda_engine)) {
+    execution_contexts.push_back(std::move(execution_context));
+    device_memory_size_ =
+        cuda_engine_ ? cuda_engine_->getDeviceMemorySize() : 0;
+  }
+  EngineContext(TrtUniquePtrType<nvinfer1::ICudaEngine>&& cuda_engine,
+                std::vector<ExecutionContext>&& execution_contexts)
+      : cuda_engine_(std::move(cuda_engine)),
+        execution_contexts(std::move(execution_contexts)) {
+    device_memory_size_ =
+        cuda_engine_ ? cuda_engine_->getDeviceMemorySize() : 0;
+  }
+
+  mutex mu;
+
+  nvinfer1::ICudaEngine* GetCudaEngine() { return cuda_engine_.get(); }
+
+  Status GetExecutionContext(int idx, nvinfer1::IExecutionContext** exec_ctx,
+                             bool* has_device_memory)
+      TF_EXCLUSIVE_LOCKS_REQUIRED(mu) {
+    if (idx >= execution_contexts.size()) {
+      return errors::Internal("Requested engine context with index ", idx,
+                              ", but only ", execution_contexts.size(),
+                              "contexts are present.");
+    }
+    *exec_ctx = execution_contexts[idx].get();
+    *has_device_memory = execution_contexts[idx].HasDeviceMemory();
+    return OkStatus();
+  }
+
+  int GetNumContexts() {
+    mutex_lock lock(mu);
+    return execution_contexts.size();
+  }
+
+  size_t GetDeviceMemorySize() { return device_memory_size_; }
+
+ private:
+  // Note: declaration has to come before execution_contexts, to ensure proper
+  // order of destruction.
+  TrtUniquePtrType<nvinfer1::ICudaEngine> cuda_engine_;
+
+ public:
+  // In explicit batch mode, we maintain a vector of contexts for each engine,
+  // where each context is created for a specific profile. This is because it is
+  // either not possible or non-trivial to change the profile of a context for
+  // the following reasons:
+  // - To switch profiles (from TRT 7), one must first ensure that all inference
+  //   calls in that context are finished. This would require an additional
+  //   synchronization before we call setOptimizationProfile. To avoid this
+  //   extra sync call, we maintain separate execution context for each profile.
+  // IExecutionContext object is not thread safe: only one thread should use it
+  // for inference at a time therefore we need a mutex. More details at
+  // https://docs.nvidia.com/deeplearning/sdk/tensorrt-best-practices/index.html#thread-safety
+  // Additional discussion about execution context management and thread safety
+  // at https://github.com/tensorflow/tensorflow/issues/36959
+  std::vector<ExecutionContext> execution_contexts TF_GUARDED_BY(mu);
+
+ private:
+  // Until TRT 8.4 ICudaEngine::getDeviceMemorySize() has a non-negligible
+  // latency. Since its value remains constant, we can cache it.
+  size_t device_memory_size_;
+};
+// Contains the context required to build the calibration data.
+class CalibrationContext {
+ public:
+  string TerminateCalibration();
+
+  // Lookup table for temporary staging areas of input tensors for calibration.
+  std::unordered_map<string, std::pair<void*, size_t>> device_buffers_;
+
+  // Temporary staging areas for calibration inputs.
+  std::vector<Tensor> device_tensors_;
+
+  std::unique_ptr<TRTInt8Calibrator> calibrator_;
+  TrtUniquePtrType<nvinfer1::IBuilder> builder_;
+  TrtUniquePtrType<nvinfer1::ICudaEngine> engine_;
+  // TODO(sami): Use threadpool threads!
+  std::unique_ptr<std::thread> thr_;
+
+ private:
+  mutex mu_;
+  bool terminated_ TF_GUARDED_BY(mu_) = false;
+  std::string calibration_table_ TF_GUARDED_BY(mu_);
+};
+
+ABSL_CONST_INIT extern const absl::string_view kTfTrtContainerName;
+
+class TRTEngineCacheResource : public ResourceBase {
+ public:
+  // According to the TensorRT API, the logger is considered a singleton by the
+  // TensorRT library, and multiple instances of IRuntime and/or IBuilder must
+  // all use the same logger. So here we make it a singleton.
+  //
+  // TODO(laigd): use this logger in all places where conversion happens.
+  static Logger& GetLogger();
+
+  TRTEngineCacheResource(OpKernelContext* ctx, size_t capacity);
+
+  ~TRTEngineCacheResource() override;
+
+  string DebugString() const override;
+
+  // Returns the EngineContext that is compatible with input_shapes.
+  // Returns nullptr if no compatible EngineContexts is found in cache.
+  EngineContext* GetEngineContext(const std::vector<TensorShape>& input_shapes);
+
+  // Returns the EngineContext that is compatible with profile_id.
+  // This function should be only called in explicit batch mode where
+  // cache size is expected to be at most one.
+  // Returns nullptr if no compatible EngineContexts is found in cache.
+  EngineContext* GetEngineContext(const int profile_id);
+
+  // Keep device allocator for TRT.
+  std::unique_ptr<TRTBaseAllocator> allocator_;
+
+  // Declare cache after allocator so that it is destroyed before allocator is.
+  LRUCache<std::vector<TensorShape>, std::unique_ptr<EngineContext>,
+           VectorTensorShapeHasher>
+      cache_;
+
+  // TODO(hinsu): Use different calibration context for the available shapes and
+  // attach it to each item of the cache.
+  std::unique_ptr<CalibrationContext> calib_ctx_;
+
+  // This object maintains all the optimization profiles during profile
+  // generation and engine build. During runtime the list of profiles is used to
+  // look up a matching profile for the input data.
+  TrtShapeOptimizationProfile profiles_;
+};
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LRU_CACHE_H_

tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache_test.cc

@@ -0,0 +1,59 @@
+/* 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 "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
+
+#include <functional>
+
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+TEST(LRUCacheTest, Basic) {
+  LRUCache<int, int, std::hash<int>> cache;
+  cache.reserve(2);
+  // Insert 10
+  cache.emplace(10, 100);
+  EXPECT_EQ(cache.size(), 1);
+  EXPECT_EQ(cache.count(10), 1);
+  EXPECT_EQ(cache.at(10), 100);
+  EXPECT_EQ(cache.count(100), 0);
+  // Insert 20
+  cache.emplace(20, 200);
+  EXPECT_EQ(cache.size(), 2);
+  EXPECT_EQ(cache.count(10), 1);
+  EXPECT_EQ(cache.count(20), 1);
+  EXPECT_EQ(cache.at(10), 100);
+  EXPECT_EQ(cache.at(20), 200);
+  EXPECT_EQ(cache.count(100), 0);
+  EXPECT_EQ(cache.count(200), 0);
+  // Insert 30, Evicting 10
+  cache.emplace(30, 300);
+  EXPECT_EQ(cache.count(10), 0);
+  EXPECT_EQ(cache.count(20), 1);
+  EXPECT_EQ(cache.count(30), 1);
+  // Touch 20
+  cache.at(20);
+  // Insert 40, Evicting 30
+  cache.emplace(40, 400);
+  EXPECT_EQ(cache.count(10), 0);
+  EXPECT_EQ(cache.count(20), 1);
+  EXPECT_EQ(cache.count(30), 0);
+  EXPECT_EQ(cache.count(40), 1);
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow

tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.cc

@@ -0,0 +1,661 @@
+/* Copyright 2020 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 "tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h"
+
+#include <algorithm>
+#include <functional>
+
+#include "absl/algorithm/container.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "tensorflow/core/profiler/lib/traceme.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "third_party/gpus/cuda/include/cuda_runtime_api.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// Returns a vector of nvinfer1::Dims for a vector of TensorShapes.
+template <typename TensorShapeType>
+std::vector<nvinfer1::Dims> GetDimVec(std::vector<TensorShapeType> shape_vec) {
+  std::vector<nvinfer1::Dims> dimvec(shape_vec.size());
+  absl::c_transform(shape_vec, dimvec.begin(), [](TensorShapeType shape) {
+    auto adap = DimsAdapter::Create(shape);
+    TF_CHECK_OK(adap.status());
+    return adap->AsTrtDims();
+  });
+  return dimvec;
+}
+
+// In dynamic shape mode the optimization profile dims are only allowed to
+// differ from the network input dims where the network input dims have -1
+// values. We enforce this condition by changing prof_dims if necessary.
+void EnforceCompatibility(nvinfer1::Dims* prof_dims,
+                          const PartialTensorShape& input_shape) {
+  for (int i = 0; i < input_shape.dims(); i++) {
+    if (input_shape.dim_size(i) != -1) {
+      prof_dims->d[i] = input_shape.dim_size(i);
+    }
+  }
+}
+
+void SetImplicitBatchModeCompatibleProfile(
+    const std::vector<nvinfer1::Dims>& dimvec, std::vector<nvinfer1::Dims>* min,
+    std::vector<nvinfer1::Dims>* opt, std::vector<nvinfer1::Dims>* max) {
+  *min = dimvec;
+  for (auto& dim : *min) {
+    // Shape value tensors can have -1 value as a wildcard. We do not change
+    // in that case.
+    if (dim.d[0] != -1) dim.d[0] = 1;  // Set min batch size to 1.
+  }
+  *opt = dimvec;
+  *max = dimvec;
+}
+
+void TrtShapeOptimizationProfile::ImplicitBatchModeCompatibleStrategy(
+    const std::vector<std::vector<nvinfer1::Dims>>& collected_shapes) {
+  for (auto& shape_vec : collected_shapes) {
+    std::vector<nvinfer1::Dims> min, opt, max;
+    SetImplicitBatchModeCompatibleProfile(shape_vec, &min, &opt, &max);
+    VLOG(2) << "Initializing optimization profile config with min="
+            << DebugString(min) << ", opt=max=" << DebugString(max);
+    OptimizationProfileConfig profConfig{min, opt, max};
+    profiles_.push_back(std::move(profConfig));
+  }
+}
+
+// Applies a binary operation for each dimension of the input shapes.
+// x[i].d[k] = op(x[i].d[k], y[i].d[k]), where i enumerates the input tensors,
+// and k enumerates the dimensions of the tensors. The BinaryOperation may be
+// std::min, std::max etc.
+template <typename BinaryOperation>
+Status ShapeProfileBinaryOp(std::vector<nvinfer1::Dims>* x,
+                            const std::vector<nvinfer1::Dims>& y,
+                            BinaryOperation op) {
+  if (x->size() != y.size())
+    return errors::InvalidArgument(
+        "Number of input tensors differ during profile creation");
+  for (int i = 0; i < x->size(); i++) {
+    if (x->at(i).nbDims != y[i].nbDims)
+      return errors::InvalidArgument(
+          "Number of input dimensions differ during profile creation at dim ",
+          i, ", values ", x->at(i).nbDims, y[i].nbDims);
+    for (int j = 0; j < x->at(i).nbDims; j++) {
+      x->at(i).d[j] = op(x->at(i).d[j], y[i].d[j]);
+    }
+  }
+  return OkStatus();
+}
+
+Status TrtShapeOptimizationProfile::RangeStrategy(
+    const std::vector<std::vector<nvinfer1::Dims>>& collected_shapes) {
+  if (collected_shapes.empty()) return OkStatus();
+
+  std::vector<nvinfer1::Dims> min = collected_shapes[0];
+  std::vector<nvinfer1::Dims> max = min;
+
+  for (int i = 1; i < collected_shapes.size(); i++) {
+    TF_RETURN_IF_ERROR(
+        ShapeProfileBinaryOp(&min, collected_shapes[i],
+                             [](int a, int b) { return std::min(a, b); }));
+    TF_RETURN_IF_ERROR(
+        ShapeProfileBinaryOp(&max, collected_shapes[i],
+                             [](int a, int b) { return std::max(a, b); }));
+  }
+  VLOG(2) << "Initializing optimization profile config with min="
+          << DebugString(min) << ", opt=max=" << DebugString(max);
+  OptimizationProfileConfig profConfig{min, max, max};
+  profiles_.push_back(std::move(profConfig));
+  return OkStatus();
+}
+
+void TrtShapeOptimizationProfile::OptimalStrategy(
+    const std::vector<std::vector<nvinfer1::Dims>>& collected_shapes) {
+  for (auto& shape_vec : collected_shapes) {
+    std::vector<nvinfer1::Dims> min = shape_vec;
+    std::vector<nvinfer1::Dims> opt = min;
+    std::vector<nvinfer1::Dims> max = min;
+    VLOG(2) << "Initializing optimization profile config with min=opt=max="
+            << DebugString(min);
+    OptimizationProfileConfig profConfig{min, opt, max};
+    profiles_.push_back(std::move(profConfig));
+  }
+}
+
+// Collects the values of tensors that are ShapeTensorCompatible to. The values
+// are stored in the actual_shape_values_ member variable.
+Status TrtShapeOptimizationProfile::CollectShapeValues(OpKernelContext* ctx) {
+  tensorflow::profiler::TraceMe activity(
+      "TrtShapeOptimizationProfile::CollectShapeValues",
+      tensorflow::profiler::TraceMeLevel::kInfo);
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(CHECK_NOTNULL(
+      ctx->op_device_context()->stream()->platform_specific_handle().stream));
+  actual_shape_values_.resize(ctx->num_inputs());
+  if (is_shape_tensor_.empty()) {
+    is_shape_tensor_.resize(ctx->num_inputs());
+    for (int i = 0; i < ctx->num_inputs(); i++) {
+      is_shape_tensor_[i] = IsTrtShapeTensorCompatible(ctx->input(i));
+    }
+  }
+  int n_shape_val = 0;
+  // First copy all the shape value candidates into actual_shape_values_ vector.
+  for (int i = 0; i < ctx->num_inputs(); i++) {
+    if (is_shape_tensor_[i]) {
+      if (ctx->input_dtype(i) != DT_INT32) {
+        // In case the is_shape_tensor mask was initialized with the input
+        // shapes only (without knowledge of dtype) then we apply correction.
+        is_shape_tensor_[i] = false;
+        continue;
+      }
+      if (input_shape_values_.size() > 0 &&
+          input_shape_values_[0][i].nbDims != ctx->input(i).NumElements()) {
+        // Shape tensor dims should not change. It must be a value tensor.
+        is_shape_tensor_[i] = false;
+        continue;
+      }
+      // We have to copy the shape values to the host, because TRT's
+      // ExecutionContext::setInputShapeBinding expects a host pointer.
+      n_shape_val++;
+      const Tensor& input = ctx->input(i);
+      actual_shape_values_[i].nbDims = input.NumElements();
+      auto ret = cudaMemcpyAsync(
+          actual_shape_values_[i].d, input.flat<int32>().data(),
+          input.NumElements() * sizeof(int32), cudaMemcpyDeviceToHost, stream);
+      if (ret != 0) {
+        return errors::Internal("Could not copy shape tensor values");
+      }
+      VLOG(2) << "Input " << i << " is (probably) a shape tensor, n_values="
+              << input.NumElements();
+    } else {
+      actual_shape_values_[i] = {0, {}};
+    }
+  }
+  if (n_shape_val > 0) {
+    // If we have any shape values candidates, then wait until data is copied
+    // to host.
+    cudaStreamSynchronize(stream);
+  }
+  return OkStatus();
+}
+
+// Collects the values of tensors that are ShapeTensorCompatible to. To be used
+// for unit tests.
+Status TrtShapeOptimizationProfile::CollectShapeValues(const DataVec& input) {
+  actual_shape_values_.resize(input.size());
+  for (int i = 0; i < input.size(); i++) {
+    if (is_shape_tensor_[i]) {
+      if (!IsTrtShapeTensorCompatible(input[i].tensor)) {
+        return errors::Internal("Inconsistent shape tensor ", input[i].name,
+                                ", ", i);
+      }
+      int n_elements = input[i].tensor.NumElements();
+      actual_shape_values_[i].nbDims = n_elements;
+      // During unit tests, the data is in unified memory
+      std::copy(input[i].tensor.flat<int32>().data(),
+                input[i].tensor.flat<int32>().data() + n_elements,
+                actual_shape_values_[i].d);
+      VLOG(2) << "Collected tensor shape values "
+              << DebugString(actual_shape_values_[i]);
+    } else {
+      actual_shape_values_[i] = {0, {}};
+    }
+  }
+  return OkStatus();
+}
+
+// Adjusts shape value profile to prevent TRT from removing shape value input
+// bindings whose value is redundant (only a single value matches the profile).
+// This should be removed once the NVIDIA bug 3153064 is fixed.
+void FixShapeValueProfile(OptimizationProfileConfig* prof,
+                          const std::vector<bool>& is_shape_tensor) {
+  int shape_value_offset = is_shape_tensor.size();
+  for (int i = 0; i < is_shape_tensor.size(); i++) {
+    if (is_shape_tensor[i] &&
+        std::equal(prof->min[shape_value_offset + i].d,
+                   prof->min[shape_value_offset + i].d +
+                       prof->min[shape_value_offset + i].nbDims,
+                   prof->max[shape_value_offset + i].d)) {
+      prof->max[shape_value_offset + i].d[0]++;
+      VLOG(2) << "Adjusted profile for shape value tensor " << i << " "
+              << DebugString(prof->max[shape_value_offset + i]);
+    } else {
+      VLOG(2) << i << " is not a shape tensor." << is_shape_tensor[i];
+    }
+  }
+}
+
+// Checks whether rhs is already contained in values.
+bool AlreadyCollected(const std::vector<std::vector<nvinfer1::Dims>>& values,
+                      const std::vector<nvinfer1::Dims>& rhs) {
+  for (auto& lhs : values) {
+    bool ret = lhs.size() == rhs.size();
+    for (int i = 0; ret && i < lhs.size(); i++) {
+      ret &= lhs[i].nbDims == rhs[i].nbDims;
+      for (int j = 0; ret && j < lhs[i].nbDims; j++) {
+        ret &= (lhs[i].d[j] == rhs[i].d[j]);
+      }
+    }
+    if (ret) return true;
+  }
+  return false;
+}
+
+void TrtShapeOptimizationProfile::InitProfiles(
+    const std::vector<PartialTensorShape>& input_partial_shapes,
+    ProfileStrategy strategy) {
+  strategy_ = strategy;
+  if (input_shapes_.size() == 0) {
+    VLOG(1) << "Not creating profiles without input_shapes. "
+               "You have to enable profile generation mode first (build).";
+    return;
+  }
+  // Preprocess the vector of input shapes and shape values:
+  // - Converts TensorShape -> nvinfer::Dims.
+  // - Concatenates the shape values after the input shapes:
+  //   dimvec = [dim0, dim1,..., shapeval0, shapval1, ...]
+  // - Ensures that the list is unique.
+  std::vector<std::vector<nvinfer1::Dims>> collected_shapes;
+  for (int i = 0; i < input_shapes_.size(); i++) {
+    auto shape_vec = input_shapes_[i];
+    VLOG(2) << "Initprofiles, processing shape " << i;
+    if (!shape_vec.empty()) {
+      // Correct for values that are mistakenly used as shape values
+      for (int k = 0; k < input_shape_values_[i].size(); k++) {
+        if (!is_shape_tensor_[k])
+          input_shape_values_[i][k] = nvinfer1::Dims{0, {}};
+      }
+      std::vector<nvinfer1::Dims> dimvec = GetDimVec(shape_vec);
+      dimvec.insert(dimvec.end(), input_shape_values_[i].begin(),
+                    input_shape_values_[i].end());
+      // TODO(tfeher): This condition should not apply for explicit profile. In
+      // that case consicutive elements in collected_shapes contain the user
+      // defined values of min, opt and max, and it is valid the have min = opt
+      // and opt = max.
+      if (!AlreadyCollected(collected_shapes, dimvec)) {
+        collected_shapes.push_back(dimvec);
+      }
+    }
+  }
+  switch (strategy_) {
+    case ProfileStrategy::kImplicitBatchModeCompatible:
+      VLOG(1) << "Creating profiles with ImplicitBatchModeCompatible strategy";
+      ImplicitBatchModeCompatibleStrategy(collected_shapes);
+      break;
+    // Treat all other strategies the same as kOptimal for now. Implementing
+    // those is outlined in the dynamic shape support implementation plan.
+    case ProfileStrategy::kRange:
+      VLOG(1) << "Creating profiles with Range strategy";
+      TF_CHECK_OK(RangeStrategy(collected_shapes));
+      break;
+    case ProfileStrategy::kRangeOptimal:
+      VLOG(1) << "Creating profiles with RangeOptimal strategy";
+      OptimalStrategy(collected_shapes);
+      TF_CHECK_OK(RangeStrategy(collected_shapes));
+      break;
+    case ProfileStrategy::kOptimal:
+      VLOG(1) << "Creating profiles with Optimal strategy";
+      OptimalStrategy(collected_shapes);
+      break;
+  }
+  // Define a mask that describe which input could be a shape tensor. Note
+  // that here we can have false positives. The shape tensor mask will be
+  // updated once the network is constructed.
+  SetShapeTensorMask(input_partial_shapes);
+  if (input_partial_shapes.size() > 0) {
+    for (OptimizationProfileConfig& prof : profiles_) {
+      // TODO: Remove this when the bug is fixed.
+#if !IS_TRT_VERSION_GE(8, 0, 0, 0)
+      FixShapeValueProfile(&prof, is_shape_tensor_);
+#endif
+      for (int i = 0; i < input_partial_shapes.size(); i++) {
+        auto network_input = input_partial_shapes[i];
+        EnforceCompatibility(&prof.min[i], network_input);
+        EnforceCompatibility(&prof.opt[i], network_input);
+        EnforceCompatibility(&prof.max[i], network_input);
+      }
+    }
+  }
+}
+
+void TrtShapeOptimizationProfile::InitCalibProfile(
+    const std::vector<TensorShape>& shapes) {
+  VLOG(1) << "Collected shape(s) " << DebugString(shapes) << " for "
+          << " calibration profile.";
+  auto shape_vec = shapes;
+  if (!shape_vec.empty()) {
+    std::vector<nvinfer1::Dims> dimvec = GetDimVec(shape_vec);
+    dimvec.insert(dimvec.end(), actual_shape_values_.begin(),
+                  actual_shape_values_.end());
+    VLOG(2) << "Initializing calibration optimization profile config with "
+            << "min=opt=max " << DebugString(dimvec);
+
+    OptimizationProfileConfig profConfig{dimvec, dimvec, dimvec};
+    calib_profiles_ = std::move(profConfig);
+  } else {
+    VLOG(2) << "Failed to initialize calibration optimization profile.";
+  }
+}
+
+Status TrtShapeOptimizationProfile::AddProfiles(
+    nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config,
+    const nvinfer1::INetworkDefinition* network) {
+  // Create optimization profile for calibration if necessary.
+  if (!calib_profiles_.min.empty()) {
+    VLOG(2) << "Setting up calibration profiles";
+    auto* calibProfile = builder->createOptimizationProfile();
+    Status status =
+        calib_profiles_.SetDimensions(network, calibProfile, input_mask_);
+    if (!status.ok()) {
+      return status;
+    }
+    bool result = false;
+    if (calibProfile->isValid()) {
+      result = config->setCalibrationProfile(calibProfile);
+    } else {
+      VLOG(2) << "Calibration profile is not valid";
+    }
+    if (result) {
+      VLOG(2) << "Added calibration optimization profile "
+              << calib_profiles_.DebugString() << " to builder config.";
+    } else {
+      VLOG(2) << "FAILED TO ADD PROFILE";
+      LOG(ERROR) << "Failed to add calibration optimization profile "
+                 << calib_profiles_.DebugString()
+                 << ". This usually happens when profile is invalid.";
+    }
+  }
+  // Create a vector of optimization profiles.
+  for (int i = 0; i < profiles_.size(); i++) {
+    auto* optProfile = builder->createOptimizationProfile();
+    Status status =
+        profiles_[i].SetDimensions(network, optProfile, input_mask_);
+    if (!status.ok()) {
+      return status;
+    }
+    int idx = -1;
+    if (optProfile->isValid()) {
+      idx = config->addOptimizationProfile(optProfile);
+    }
+    if (idx >= 0) {
+      if (i != idx) {
+        return errors::Internal(
+            "Profile index of engine config is different from source profile "
+            "index: ",
+            i, " != ", idx);
+      }
+      VLOG(1) << "Added optimization profile " << profiles_[i].DebugString()
+              << " with idx " << idx << " to builder config.";
+    } else {
+      LOG(ERROR) << "Failed to add optimization profile "
+                 << profiles_[i].DebugString()
+                 << ". This usually happens when profile is invalid.";
+    }
+  }
+  if (!profiles_.empty() && config->getNbOptimizationProfiles() == 0) {
+    return errors::Internal("Failure in adding an optimization profile.");
+  }
+  need_profiles_ = config->getNbOptimizationProfiles() > 0;
+  // Update the mask that flag shape tensors. The network is known now,
+  // the mask will be correct.
+  SetShapeTensorMask(network);
+  is_pruned_input_.resize(network->getNbInputs());
+  absl::c_fill(is_pruned_input_, false);
+  return OkStatus();
+}
+
+Status TrtShapeOptimizationProfile::ConfigureBuilder(
+    nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config,
+    const nvinfer1::INetworkDefinition* network) {
+  TF_RETURN_IF_ERROR(AddProfiles(builder, config, network));
+  return OkStatus();
+}
+
+// Sets the shape tensor mask from the TRT engine definition.
+void TrtShapeOptimizationProfile::SetShapeTensorMask(
+    const nvinfer1::ICudaEngine* engine, int n_inputs) {
+  is_shape_tensor_.resize(n_inputs, false);
+  for (int i = 0; i < n_inputs; i++) {
+    int binding_index;
+    Status status = GetTrtBindingIndex(i, 0, engine, &binding_index);
+    if (!status.ok()) {
+      continue;
+    }
+    is_shape_tensor_[i] = engine->isShapeBinding(binding_index);
+    if (is_shape_tensor_[i]) {
+      VLOG(2) << "Found shape tensor at " << i;
+    }
+  }
+  has_shape_tensor_ =
+      absl::c_any_of(is_shape_tensor_, [](bool b) { return b; });
+}
+
+// Sets the shape tensor mask using the network definition.
+void TrtShapeOptimizationProfile::SetShapeTensorMask(
+    const nvinfer1::INetworkDefinition* network) {
+  int n_inputs = network->getNbInputs();
+  is_shape_tensor_.resize(n_inputs, false);
+  for (int i = 0; i < n_inputs; i++) {
+    const ITensorProxyPtr input = network->getInput(i);
+    is_shape_tensor_[i] = input->isShapeTensor();
+    if (is_shape_tensor_[i]) {
+      VLOG(2) << "Found shape tensor " << input->getName() << " at " << i;
+    }
+  }
+  has_shape_tensor_ =
+      absl::c_any_of(is_shape_tensor_, [](bool b) { return b; });
+}
+
+// Sets the shape tensor mask using the input partial shapes. This only tells
+// whether the tensors are shape value compatible, only the final network
+// definition or the engine would give concrete answers.
+void TrtShapeOptimizationProfile::SetShapeTensorMask(
+    const std::vector<PartialTensorShape>& input_partial_shapes) {
+  if (is_shape_tensor_.size() == input_partial_shapes.size()) {
+    // Already initialized, e.g. by TRTEngineOp::ComputeAsync().
+    return;
+  }
+  is_shape_tensor_.resize(input_partial_shapes.size(), false);
+  for (int i = 0; i < input_partial_shapes.size(); i++) {
+    is_shape_tensor_[i] = IsTrtShapeTensorCompatible(input_partial_shapes[i]);
+    if (is_shape_tensor_[i]) {
+      VLOG(2) << "Found shape compatible tensor at " << i;
+    }
+  }
+  has_shape_tensor_ =
+      absl::c_any_of(is_shape_tensor_, [](bool b) { return b; });
+}
+
+int TrtShapeOptimizationProfile::GetProfileNumber(
+    const std::vector<TensorShape>& shapes) {
+  tensorflow::profiler::TraceMe activity(
+      "TrtShapeOptimizationProfile::GetProfileNumber",
+      tensorflow::profiler::TraceMeLevel::kInfo);
+  if (!need_profiles_) return 0;
+  // TODO(tfeher): Return the best profile not just the first compatible.
+  for (int i = 0; i < profiles_.size(); i++) {
+    if (profiles_[i].IncludesShapes(shapes, HasShapeTensor(),
+                                    actual_shape_values_, is_pruned_input_,
+                                    is_shape_tensor_)) {
+      return i;
+    }
+  }
+  VLOG(1) << "Profile not found for input shapes " << DebugString(shapes);
+  VLOG(2) << "  and shape values " << DebugString(actual_shape_values_);
+  return -1;
+}
+
+Status TrtShapeOptimizationProfile::CreateExecutionContexts(
+    nvinfer1::ICudaEngine* engine,
+    std::vector<ExecutionContext>* exec_contexts) {
+  int i = 0;
+  // The following loop runs once if we have static shapes, to create a single
+  // execution context without profiles. In dynamic mode we create one context
+  // for each profile and set the corresponding optimization profile.
+  do {
+    VLOG(1) << "Creating execution context " << i;
+    ExecutionContext context = ExecutionContext::Create(engine);
+    if (i > 0) {
+      // This condition is needed for two reasons:
+      // - using static shapes we do not have any profiles so we cannot call
+      //   set optimizationprofiles.
+      // - The 0th profile is set implicitly for the first execution context
+      //   therefore we do not need to set.
+      if (!context->setOptimizationProfile(i)) {
+        return errors::Internal("Could not set TRT optimization profile.");
+      }
+    }
+    exec_contexts->push_back(std::move(context));
+    i++;
+  } while (i < profiles_.size());
+
+  return OkStatus();
+}
+
+Status TrtShapeOptimizationProfile::SetInputShapeBinding(
+    int input_index, int binding_index, nvinfer1::ICudaEngine* cuda_engine,
+    nvinfer1::IExecutionContext* exec_context) const {
+  tensorflow::profiler::TraceMe activity(
+      "TrtShapeOptimizationProfile::SetInputShapeBinding",
+      tensorflow::profiler::TraceMeLevel::kInfo);
+  if (cuda_engine->isShapeBinding(binding_index)) {
+    // Input shape binding data has to be in host memory. That is the reason
+    // we can't use input_tensor.flat().data(). which contains the same
+    // values in device memory. Instead, we use data that was copied to host
+    // by CollectShapeValues.
+    VLOG(2) << "Setting input shape binding for idx " << binding_index
+            << ", with values "
+            << DebugString(actual_shape_values_.at(input_index));
+    bool ret = exec_context->setInputShapeBinding(
+        binding_index, actual_shape_values_.at(input_index).d);
+    if (!ret) {
+      return errors::Internal("Could not set input shape binding for idx ",
+                              binding_index);
+    }
+  }
+  return OkStatus();
+}
+
+// If binding_idx is a shape tensor, then returns the associated min/max/opt
+// shape values from prof_idx.
+nvinfer1::Dims GetDimsFromShapeVal(int prof_idx, int binding_idx,
+                                   nvinfer1::OptProfileSelector selector,
+                                   const nvinfer1::ICudaEngine* engine) {
+  if (engine->isShapeBinding(binding_idx)) {
+    const int32* shape_val_ptr =
+        engine->getProfileShapeValues(binding_idx, prof_idx, selector);
+    if (shape_val_ptr) {
+      VLOG(2) << "Found shape value in prof " << prof_idx << ", binding "
+              << binding_idx;
+      nvinfer1::Dims dims = engine->getBindingDimensions(binding_idx);
+      // nbDims == 0 represent scalar, -1 represents invalid dim
+      int n_values = (dims.nbDims == 0) ? 1 : dims.d[0];
+      if (n_values > 0) {
+        dims.nbDims = n_values;
+        std::copy(shape_val_ptr, shape_val_ptr + n_values, dims.d);
+      }
+      return dims;
+    }
+  }
+  return {0, {0}};
+}
+
+Status TrtShapeOptimizationProfile::SetPrunedMask(
+    const nvinfer1::ICudaEngine* engine, int n_network_inputs) {
+  is_pruned_input_.resize(n_network_inputs);
+  absl::c_fill(is_pruned_input_, false);
+  for (int j = 0; j < n_network_inputs; j++) {
+    int binding_idx;
+    Status status = GetTrtBindingIndex(j, 0, engine, &binding_idx);
+    if (!status.ok()) {
+      // Before TRT 8, an input tensor can be pruned (nvbugs/3153064)
+      // Resource inputs are also unknown by TRT, so we can treat them as
+      // pruned (the engine includes the variable as weights).
+      is_pruned_input_[j] = true;
+      VLOG(2) << "Skipping pruned input " << j;
+      continue;
+    }
+  }
+  return OkStatus();
+}
+
+Status TrtShapeOptimizationProfile::RestoreProfiles(
+    const nvinfer1::ICudaEngine* engine, int n_network_inputs) {
+  need_profiles_ = false;
+  if (!engine) {
+    // We do not need to restore profiles for an empty engine.
+    return OkStatus();
+  }
+  if (engine->hasImplicitBatchDimension()) {
+    // Nothing to do, we cannot have profiles in implicit batch mode.
+    return OkStatus();
+  }
+  int n_profiles = engine->getNbOptimizationProfiles();
+  need_profiles_ = n_profiles > 0;
+  int n_inputs = GetNumberOfEngineInputs(engine);
+  if (n_inputs > n_network_inputs) {
+    return errors::Internal("Incorrect number of engine inputs");
+  }
+  VLOG(2) << "Attempting to restore " << n_profiles << " profiles, each with "
+          << n_inputs << " inputs";
+  SetShapeTensorMask(engine, n_network_inputs);
+
+  TF_RETURN_IF_ERROR(SetPrunedMask(engine, n_network_inputs));
+
+  for (int prof_idx = 0; prof_idx < n_profiles; prof_idx++) {
+    OptimizationProfileConfig cfg;
+
+    cfg.min.resize(n_network_inputs * 2);
+    cfg.max.resize(n_network_inputs * 2);
+    cfg.opt.resize(n_network_inputs * 2);
+    // restore shape values
+    for (int j = 0; j < n_network_inputs; j++) {
+      if (is_pruned_input_[j]) continue;
+      int binding_idx;
+      TF_RETURN_IF_ERROR(GetTrtBindingIndex(j, 0, engine, &binding_idx));
+
+      nvinfer1::Dims min = engine->getProfileDimensions(
+          binding_idx, prof_idx, nvinfer1::OptProfileSelector::kMIN);
+      nvinfer1::Dims max = engine->getProfileDimensions(
+          binding_idx, prof_idx, nvinfer1::OptProfileSelector::kMAX);
+      nvinfer1::Dims opt = engine->getProfileDimensions(
+          binding_idx, prof_idx, nvinfer1::OptProfileSelector::kOPT);
+      cfg.min[j] = min;
+      cfg.max[j] = max;
+      cfg.opt[j] = opt;
+
+      cfg.min[j + n_inputs] = GetDimsFromShapeVal(
+          prof_idx, binding_idx, nvinfer1::OptProfileSelector::kMIN, engine);
+      cfg.max[j + n_inputs] = GetDimsFromShapeVal(
+          prof_idx, binding_idx, nvinfer1::OptProfileSelector::kMAX, engine);
+      cfg.opt[j + n_inputs] = GetDimsFromShapeVal(
+          prof_idx, binding_idx, nvinfer1::OptProfileSelector::kOPT, engine);
+    }
+    VLOG(2) << "Restored profile " << cfg.DebugString();
+    profiles_.push_back(std::move(cfg));
+  }
+  return OkStatus();
+}
+
+int TrtShapeOptimizationProfile::GetNumProfiles() const {
+  return profiles_.size();
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h

@@ -0,0 +1,351 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_SHAPE_OPTIMIZATION_PROFILES_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_SHAPE_OPTIMIZATION_PROFILES_H_
+
+#include <list>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/common/datavec.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/trt_parameters.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_execution_context.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// Stores optimization profile parameters (min/opt/max of each input shape).
+//
+// A TensorRT optimization profile describes the possible min/max values of
+// each dynamic input shape along with an optimum value. These values are used
+// by the TensorRT builder to select the best kernel for the optimum value among
+// those kernels that are valid for all input tensors in the [min, max] range.
+struct OptimizationProfileConfig {
+  // Length of vector == 2*num_inputs to engine. min[0:num_inputs-1] are the min
+  // input dimensions for execution tensors. If engine has shape input tensors,
+  // then min[num_inputs + i] store the shape value for input i. For inputs that
+  // are not shape tensors min = opt = max = {0, {}}.
+  //
+  // When the OptimizationProfileConfig is created from the network definition
+  // (AddProfiles), then each elements of the min, opt, max vectors are defined.
+  // When the OptimizationProfileConfig object is restored during engine
+  // deserialization (RestoreProfiles), then some inputs can be pruned
+  // (see TrtShapeOptimizationProfile::is_pruned_input_). In that case min[i]
+  // is not defined for pruned inputs (same is true for opt and max).
+  std::vector<nvinfer1::Dims> min;
+  std::vector<nvinfer1::Dims> opt;
+  std::vector<nvinfer1::Dims> max;
+
+  string DebugString() const {
+    using absl::StrCat;
+    return StrCat("[min: ", tensorflow::tensorrt::DebugString(min),
+                  ", opt: : ", tensorflow::tensorrt::DebugString(opt),
+                  ", max: ", tensorflow::tensorrt::DebugString(max), "]");
+  }
+
+  // Sets the min/opt/max dimensions for profile.
+  //
+  // The given min/opt/max dimensions should satisfy the condition
+  // min <= opt <= max. Additionally TRT requires that the min/opt/max values
+  // are compatible with the network input. Compatibility is defined the
+  // following way: let dim be the shape of an input binding and min/opt/max the
+  // corresponding profile dims. TRT requires that dim.d[k] must be -1 if
+  // (min.d[k] != dim.d[k] || opt.d[k] != dim.d[k] || max.d[k] != dim.d[k]).
+  //
+  // Parameters:
+  // network - TensorRT network, used to enumerate all the input tensors
+  // profile - on exit the profile information will be set for each input tensor
+  // input_mask - 1 for TRT inputs, 0 for TF inputs that are not TRT inputs
+  Status SetDimensions(const nvinfer1::INetworkDefinition* network,
+                       nvinfer1::IOptimizationProfile* profile,
+                       const std::vector<bool>& input_mask) const {
+    int n_inputs_trt = network->getNbInputs();
+    int n_inputs_tf = opt.size() / 2;
+    /// TODO(lsugy): check that the sum of the mask equals n_inputs.
+    if (input_mask.size() != n_inputs_tf) {
+      return errors::Internal("Incorrect input mask size: ", input_mask.size());
+    }
+    int n_mask_true = 0;
+    for (bool mask_val : input_mask) {
+      if (mask_val) {
+        n_mask_true++;
+      }
+    }
+    if (n_mask_true != n_inputs_trt) {
+      return errors::Internal(
+          "Number of true elements in input_mask (", n_mask_true,
+          ") doesn't match expected TRT inputs (", n_inputs_trt, ")");
+    }
+    int j = 0;
+    for (int i = 0; i < n_inputs_tf; i++) {
+      if (input_mask[i]) {
+        const ITensorProxyPtr input = network->getInput(j);
+        const char* name = input->getName();
+        if (input->isShapeTensor()) {
+          int idx = i + n_inputs_tf;
+          VLOG(2) << "Setting shape values for " << name << ", "
+                  << ::tensorflow::tensorrt::DebugString(opt[idx]);
+          profile->setShapeValues(name, nvinfer1::OptProfileSelector::kMIN,
+                                  min[idx].d, min[idx].nbDims);
+          profile->setShapeValues(name, nvinfer1::OptProfileSelector::kOPT,
+                                  opt[idx].d, opt[idx].nbDims);
+          profile->setShapeValues(name, nvinfer1::OptProfileSelector::kMAX,
+                                  max[idx].d, max[idx].nbDims);
+        }
+        VLOG(2) << "Setting input dimensions for " << name << ", "
+                << ::tensorflow::tensorrt::DebugString(opt[i]);
+        profile->setDimensions(name, nvinfer1::OptProfileSelector::kMIN,
+                               min[i]);
+        profile->setDimensions(name, nvinfer1::OptProfileSelector::kOPT,
+                               opt[i]);
+        profile->setDimensions(name, nvinfer1::OptProfileSelector::kMAX,
+                               max[i]);
+
+        j++;
+      }
+    }
+    return OkStatus();
+  }
+
+  // Returns true if profile range completely includes the given shapes.
+  bool IncludesShapes(const std::vector<TensorShape>& shapes,
+                      bool has_shape_tensor,
+                      const std::vector<nvinfer1::Dims>& shape_values,
+                      const std::vector<bool>& is_pruned_input,
+                      const std::vector<bool>& is_shape_tensor) const {
+    // min, max, and opt must have the same size which is already verified in
+    // SetDimensions.
+    if (min.size() != shapes.size() * 2 ||
+        (has_shape_tensor && min.size() != shape_values.size() * 2)) {
+      VLOG(2) << "Profile size mismatch min size " << min.size()
+              << " vs input shapes size " << shapes.size() << " "
+              << shape_values.size();
+      return false;
+    }
+    for (int i = 0; i < shapes.size(); i++) {
+      if (is_pruned_input[i]) {
+        continue;
+      }
+      auto current_shape = shapes[i];
+      // min, max, and opt must have the same nbDims, which is already verified
+      // in SetDimensions.
+      if (min[i].nbDims != current_shape.dims()) {
+        return false;
+      }
+      // Check if range [min, max] includes current_shape.
+      for (int dim = 0; dim < current_shape.dims(); dim++) {
+        if ((min[i].d[dim] > current_shape.dim_size(dim)) ||
+            (max[i].d[dim] < current_shape.dim_size(dim))) {
+          return false;
+        }
+      }
+    }
+    // Check shape values.
+    if (has_shape_tensor) {
+      int offset = shapes.size();
+      for (int i = 0; i < shape_values.size(); i++) {
+        if (is_pruned_input[i] || !is_shape_tensor[i]) {
+          continue;
+        }
+        auto shape_val = shape_values[i];
+        // min, max, and opt must have the same nbDims, which is already
+        // verified in SetDimensions.
+        if (min[i + offset].nbDims != shape_val.nbDims) {
+          return false;
+        }
+        // Check if range [min, max] includes shape_val.
+        for (int dim = 0; dim < shape_val.nbDims; dim++) {
+          if (min[i + offset].d[dim] > shape_val.d[dim] ||
+              max[i + offset].d[dim] < shape_val.d[dim]) {
+            return false;
+          }
+        }
+      }
+    }
+    return true;
+  }
+};
+
+// Manages Optimization profiles during TRT Engine construction.
+//
+// An optimization profile describes a range of dimensions for each TRT network
+// input, and the optimal dimensions that the auto-tuner should use for
+// optimization.
+//
+// This class stores the list of input shapes that were seen during the
+// build/profile_generation_mode phase, and using them it creates a set of
+// OptimizationProfileConfigs. These configs will be added to IBuilderConfig
+// before the engine is created.
+class TrtShapeOptimizationProfile {
+ public:
+  TrtShapeOptimizationProfile() {}
+
+  // Stores input shape information during profile_generation_mode.
+  void AddShape(const std::vector<TensorShape>& shapes) {
+    input_shapes_.push_back(shapes);
+    input_shape_values_.push_back(actual_shape_values_);
+    VLOG(1) << "Collected shape(s) " << DebugString(shapes) << " for profiles.";
+  }
+
+  // Stores the input mask.
+  void SetInputMask(const std::vector<bool>& input_mask) {
+    input_mask_ = input_mask;
+  }
+
+  // Collects ShapeTensorCompatible tensor values. This is needed both during
+  // profile_generation_mode and during normal inference calls.
+  Status CollectShapeValues(OpKernelContext* ctx);
+
+  // Collects ShapeTensorCompatible tensor values, used only for unit tests.
+  Status CollectShapeValues(const DataVec& input);
+
+  void clear() { profiles_.clear(); }
+
+  // Returns the profile number that should be used to execute the network with
+  // the given input shapes. Returns -1 if none of cached profiles are
+  // compatible with the given input shapes.
+  int GetProfileNumber(const std::vector<TensorShape>& shapes);
+
+  // Creates optimization profiles and add them to the builder config.
+  Status ConfigureBuilder(nvinfer1::IBuilder* builder,
+                          nvinfer1::IBuilderConfig* config,
+                          const nvinfer1::INetworkDefinition* network);
+
+  // Creates execution contexts for each optimization profile.
+  Status CreateExecutionContexts(nvinfer1::ICudaEngine* engine,
+                                 std::vector<ExecutionContext>* exec_contexts);
+
+  Status SetInputShapeBinding(int input_index, int binding_index,
+                              nvinfer1::ICudaEngine* cuda_engine,
+                              nvinfer1::IExecutionContext* exec_context) const;
+
+  // Creates optimization profiles profiles_ for the set of concrete input
+  // shapes collected in input_shapes_. The input_partial_shapes of the network
+  // is used to ensure that the created optimization profiles are compatible
+  // with the network.
+  void InitProfiles(const std::vector<PartialTensorShape>& input_partial_shapes,
+                    ProfileStrategy strategy);
+
+  void InitCalibProfile(const std::vector<TensorShape>& shapes);
+
+  // Returns number of created profiles.
+  int GetNumProfiles() const;
+
+  bool HasShape() const { return !input_shapes_.empty(); }
+  bool NeedProfiles() const { return need_profiles_; }
+
+  // Restores profiles from the engine (used after deserialization).
+  Status RestoreProfiles(const nvinfer1::ICudaEngine* engine,
+                         int n_network_inputs);
+
+  // Whether the network has any shape tensors.
+  bool HasShapeTensor() const { return has_shape_tensor_; }
+
+  void SetShapeTensorMask(const nvinfer1::INetworkDefinition* network);
+
+  // Whether the optimization profiles describe input that can be handled with
+  // a static engine (only 1 profile with min=max).
+  bool IsStaticCompatible() {
+    return strategy_ == ProfileStrategy::kOptimal && profiles_.size() == 1
+#if !IS_TRT_VERSION_GE(8, 0, 0, 0)
+           && !HasShapeTensor()
+#endif
+        ;
+    // TODO(tfeher): remove !HasShapeTensor() condition once the
+    // FixShapeValueProfile workaround is turned off.
+  }
+
+ private:
+  // Set of input shape vetors that we collect during profile_generation_mode.
+  std::vector<std::vector<TensorShape>> input_shapes_;
+
+  // Input shape values that we collect during profile_generation_mode. If the
+  // tensor is not compatible with a TRT shape tensor then an empty shape is
+  // stored.
+  std::vector<std::vector<nvinfer1::Dims>> input_shape_values_;
+
+  // Shape values present in the current inference call.
+  std::vector<nvinfer1::Dims> actual_shape_values_;
+
+  // The optimization profiles generated from input_shapes_.
+  std::vector<OptimizationProfileConfig> profiles_;
+
+  // The optimization profile for calibration.
+  OptimizationProfileConfig calib_profiles_;
+
+  // A TRTEngineOp can have resource inputs. These are treated as constants:
+  // their value is read during conversion and stored as weights in the TRT
+  // engine. This means that resource inputs have no corresponding TRT engine
+  // input, and we do not need to provide profile information for these. The
+  // input mask helps to identify the TRT inputs, where we need to define
+  // optimization profiles.
+  std::vector<bool> input_mask_;
+
+  // Whether the network has any shape tensors. Initially we assume that the
+  // network might have a shape value input. This will be updated when the
+  // network is created / engine is deserialized.
+  bool has_shape_tensor_ = true;
+
+  // Whether the network/engine requires optimization profiles.
+  bool need_profiles_ = false;
+
+  // Whether an input tensor is a shape tensor.
+  std::vector<bool> is_shape_tensor_;
+
+  // Whether a network input was pruned (only in TRT 7).
+  std::vector<bool> is_pruned_input_;
+
+  // Optimization profile generation strategy.
+  ProfileStrategy strategy_;
+
+  // Adds optimization profiles to the builder config.
+  Status AddProfiles(nvinfer1::IBuilder* builder,
+                     nvinfer1::IBuilderConfig* config,
+                     const nvinfer1::INetworkDefinition* network);
+
+  void SetShapeTensorMask(const nvinfer1::ICudaEngine* engine, int n_inputs);
+  void SetShapeTensorMask(
+      const std::vector<PartialTensorShape>& input_partial_shapes);
+
+  Status SetPrunedMask(const nvinfer1::ICudaEngine* engine,
+                       int n_network_inputs);
+
+  void ImplicitBatchModeCompatibleStrategy(
+      const std::vector<std::vector<nvinfer1::Dims>>& collected_shapes);
+  void OptimalStrategy(
+      const std::vector<std::vector<nvinfer1::Dims>>& collected_shapes);
+  Status RangeStrategy(
+      const std::vector<std::vector<nvinfer1::Dims>>& collected_shapes);
+};
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_SHAPE_OPTIMIZATION_PROFILES_H_

tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc

@@ -0,0 +1,256 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <string.h>
+
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/test.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+std::vector<TensorShape> DimVecToShapeVec(
+    std::vector<nvinfer1::Dims3> dimvec,
+    bool expand_with_empty_shape_values = false) {
+  std::vector<TensorShape> shapevec(dimvec.size());
+  for (int i = 0; i < dimvec.size(); i++) {
+    TensorShape shape;
+    TF_CHECK_OK(
+        TensorShapeUtils::MakeShape(dimvec[i].d, dimvec[i].nbDims, &shape));
+    shapevec[i] = shape;
+  }
+  if (expand_with_empty_shape_values) {
+    shapevec.resize(2 * dimvec.size());  // Append empty shape values
+  }
+  return shapevec;
+}
+
+bool DimsContained(const nvinfer1::Dims& dim, const nvinfer1::Dims& min,
+                   const nvinfer1::Dims& max) {
+  if (dim.nbDims != min.nbDims || dim.nbDims != max.nbDims) {
+    return false;
+  }
+  for (int i = 0; i < dim.nbDims; i++) {
+    if (dim.d[i] < min.d[i] || dim.d[i] > max.d[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool DimsEqual(const nvinfer1::Dims& a, const nvinfer1::Dims& b) {
+  if (a.nbDims != b.nbDims) {
+    return false;
+  }
+  for (int i = 0; i < a.nbDims; i++) {
+    if (a.d[i] != b.d[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+class TrtShapeOptimizationProfileTest
+    : public ::testing::TestWithParam<ProfileStrategy> {
+ protected:
+  TrtShapeOptimizationProfileTest() {
+    strategy_ = GetParam();
+    builder_ = TrtUniquePtrType<nvinfer1::IBuilder>(
+        nvinfer1::createInferBuilder(logger_));
+    network_ = TrtUniquePtrType<nvinfer1::INetworkDefinition>(
+        builder_->createNetworkV2(flags_));
+    builder_config_ = TrtUniquePtrType<nvinfer1::IBuilderConfig>(
+        builder_->createBuilderConfig());
+    builder_config_->setMaxWorkspaceSize(1 << 10);
+  }
+
+  // Defines a simple network: output = input1 + input2.
+  void DefineNetwork(nvinfer1::INetworkDefinition* network,
+                     nvinfer1::Dims3& dims) {
+    ITensorProxyPtr input1 =
+        network->addInput("input1", nvinfer1::DataType::kFLOAT, dims);
+    EXPECT_NE(nullptr, input1->trt_tensor());
+
+    ITensorProxyPtr input2 =
+        network->addInput("input2", nvinfer1::DataType::kFLOAT, dims);
+    EXPECT_NE(nullptr, input2->trt_tensor());
+
+    auto layer =
+        network->addElementWise(*input1->trt_tensor(), *input2->trt_tensor(),
+                                nvinfer1::ElementWiseOperation::kSUM);
+    EXPECT_NE(nullptr, layer);
+    // Mark the output.
+    ITensorProxyPtr output = layer->getOutput(0);
+    output->setName("output");
+    network->markOutput(*output->trt_tensor());
+  }
+
+  void CheckProfile(const std::vector<nvinfer1::Dims3>& dimvec,
+                    TrtShapeOptimizationProfile* profile, bool has_prof,
+                    bool test_optimality) {
+    std::vector<TensorShape> shape_vec = DimVecToShapeVec(dimvec);
+    int idx = profile->GetProfileNumber(shape_vec);
+    ASSERT_EQ(idx >= 0, has_prof);
+    if (idx < 0) return;
+    int prof_idx = exec_contexts_[idx]->getOptimizationProfile();
+    ASSERT_GE(prof_idx, 0);
+    for (int j = 0; j < dimvec.size(); j++) {
+      nvinfer1::Dims min = engine->getProfileDimensions(
+          j, prof_idx, nvinfer1::OptProfileSelector::kMIN);
+      nvinfer1::Dims max = engine->getProfileDimensions(
+          j, prof_idx, nvinfer1::OptProfileSelector::kMAX);
+      nvinfer1::Dims opt = engine->getProfileDimensions(
+          j, prof_idx, nvinfer1::OptProfileSelector::kOPT);
+
+      // This should always hold.
+      EXPECT_TRUE(DimsContained(dimvec[j], min, max));
+
+      if (test_optimality) {
+        // We shall have selected an optimal strategy.
+        EXPECT_TRUE(DimsEqual(dimvec[j], opt));
+      }
+    }
+  }
+
+  Logger& logger_ = *Logger::GetLogger();
+  TrtUniquePtrType<nvinfer1::IBuilder> builder_;
+  TrtUniquePtrType<nvinfer1::INetworkDefinition> network_;
+  TrtUniquePtrType<nvinfer1::IBuilderConfig> builder_config_;
+  TrtUniquePtrType<nvinfer1::ICudaEngine> engine;
+  std::vector<ExecutionContext> exec_contexts_;
+  // The order is important: exec_context_ must be destroyed first, and logger
+  // at last.
+  const uint32_t flags_ =
+      1U << static_cast<int>(
+          nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
+  ProfileStrategy strategy_;
+};
+
+INSTANTIATE_TEST_CASE_P(
+    OptProfilesTestInstantiation, TrtShapeOptimizationProfileTest,
+    ::testing::Values(ProfileStrategy::kRange, ProfileStrategy::kOptimal,
+                      ProfileStrategy::kRangeOptimal,
+                      ProfileStrategy::kImplicitBatchModeCompatible));
+
+TEST_P(TrtShapeOptimizationProfileTest, Static) {
+  // Static mode does not depend on strategies, we test only once.
+  if (strategy_ != ProfileStrategy::kRange) return;
+
+  // Network with static input shape.
+  nvinfer1::Dims3 dims(8, 8, 10);
+  DefineNetwork(network_.get(), dims);
+
+  TrtShapeOptimizationProfile profile;
+
+  // Configure and build engine - should be a no-op.
+  TF_CHECK_OK(profile.ConfigureBuilder(builder_.get(), builder_config_.get(),
+                                       network_.get()));
+
+  engine = TrtUniquePtrType<nvinfer1::ICudaEngine>(
+      builder_->buildEngineWithConfig(*network_, *builder_config_));
+  EXPECT_NE(nullptr, engine);
+  TF_CHECK_OK(profile.CreateExecutionContexts(engine.get(), &exec_contexts_));
+  // A single execution context should be created for a graph with static input.
+  ASSERT_EQ(exec_contexts_.size(), 1);
+  EXPECT_NE(nullptr, exec_contexts_[0]);
+
+  std::vector<nvinfer1::Dims3> dim_vec(2, dims);
+  std::vector<TensorShape> shape_vec = DimVecToShapeVec(dim_vec);
+  EXPECT_EQ(0, profile.GetProfileNumber(shape_vec));
+}
+
+TEST_P(TrtShapeOptimizationProfileTest, Dynamic) {
+  // Network with dynamic input shapes.
+  nvinfer1::Dims3 dims(-1, -1, 10);
+  DefineNetwork(network_.get(), dims);
+
+  TrtShapeOptimizationProfile profile;
+
+  // Set the input mask to true (no resource input)
+  std::vector<bool> input_mask(2, true);
+  profile.SetInputMask(input_mask);
+
+  std::vector<std::vector<nvinfer1::Dims3>> input_profiles{
+      {nvinfer1::Dims3(2, 2, 10), nvinfer1::Dims3(2, 2, 10)},
+      {nvinfer1::Dims3(3, 3, 10), nvinfer1::Dims3(3, 3, 10)},
+      {nvinfer1::Dims3(16, 16, 10), nvinfer1::Dims3(16, 16, 10)},
+  };
+
+  std::vector<nvinfer1::Dims3> unseen_shapes{nvinfer1::Dims3(5, 5, 10),
+                                             nvinfer1::Dims3(9, 9, 10)};
+
+  // Simulate a profile collection phase.
+  for (auto dim_vec : input_profiles) {
+    std::vector<TensorShape> shape_vec = DimVecToShapeVec(dim_vec, true);
+    profile.AddShape(shape_vec);
+  }
+  std::vector<PartialTensorShape> input_partial_shapes;
+  TF_CHECK_OK(GetNetworkInputShapes(network_.get(), &input_partial_shapes));
+  profile.InitProfiles(input_partial_shapes, strategy_);
+
+  // Configure and build engine.
+  TF_CHECK_OK(profile.ConfigureBuilder(builder_.get(), builder_config_.get(),
+                                       network_.get()));
+  engine = TrtUniquePtrType<nvinfer1::ICudaEngine>(
+      builder_->buildEngineWithConfig(*network_.get(), *builder_config_.get()));
+  ASSERT_NE(nullptr, engine);
+
+  TF_CHECK_OK(profile.CreateExecutionContexts(engine.get(), &exec_contexts_));
+
+  int n_profiles_exp;
+  switch (strategy_) {
+    case (ProfileStrategy::kImplicitBatchModeCompatible):
+    case (ProfileStrategy::kOptimal):
+      n_profiles_exp = input_profiles.size();
+      break;
+    case (ProfileStrategy::kRange):
+      n_profiles_exp = 1;
+      break;
+    case (ProfileStrategy::kRangeOptimal):
+      n_profiles_exp = 1 + input_profiles.size();
+      break;
+  }
+  // Each profile has an associated execution context.
+  EXPECT_EQ(exec_contexts_.size(), n_profiles_exp);
+
+  profile.SetShapeTensorMask(network_.get());
+
+  EXPECT_EQ(profile.HasShapeTensor(), false);
+
+  // Check if the profiles are assigned correctly.
+  for (auto dimvec : input_profiles) {
+    bool test_optimal_prof = strategy_ == ProfileStrategy::kOptimal ||
+                             strategy_ == ProfileStrategy::kRangeOptimal;
+    CheckProfile(dimvec, &profile, true, test_optimal_prof);
+  }
+  bool has_prof = (strategy_ == ProfileStrategy::kRange ||
+                   strategy_ == ProfileStrategy::kRangeOptimal);
+  CheckProfile(unseen_shapes, &profile, has_prof, false);
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_tensor_proxy.h

@@ -0,0 +1,458 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_TENSOR_PROXY_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_TENSOR_PROXY_H_
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+
+namespace tensorrt {
+
+// SimpleITensor implements part of the ITensor interfaces to support the TF-TRT
+// validator, as well as some TF-TRT tests. The former use case only utilizes
+// the interfaces related to shape and type information.
+class SimpleITensor {
+ public:
+  SimpleITensor(nvinfer1::DataType trt_dtype, const nvinfer1::Dims& trt_dims)
+      : trt_dtype_(trt_dtype), trt_dims_(trt_dims) {}
+
+  SimpleITensor() : dynamic_range_min_(0.0f), dynamic_range_max_(0.0f) {}
+  SimpleITensor(const nvinfer1::Dims& dims)
+      : trt_dims_(dims), dynamic_range_min_(0.0f), dynamic_range_max_(0.0f) {}
+
+  SimpleITensor(const std::vector<int>& dims) {
+    trt_dims_.nbDims = dims.size();
+    for (int i = 0; i < dims.size(); ++i) {
+      trt_dims_.d[i] = dims[i];
+    }
+    dynamic_range_min_ = 0.0f;
+    dynamic_range_max_ = 0.0f;
+  }
+
+  void setName(const char* name) {}
+
+  const char* getName() const { return ""; }
+
+  void setDimensions(nvinfer1::Dims dimensions) { trt_dims_ = dimensions; }
+
+  nvinfer1::Dims getDimensions() const { return trt_dims_; }
+
+  void setType(nvinfer1::DataType trt_dtype) { trt_dtype_ = trt_dtype; }
+
+  nvinfer1::DataType getType() const { return trt_dtype_; }
+
+  bool isNetworkInput() const { return false; }
+
+  bool isNetworkOutput() const { return false; }
+
+  void setBroadcastAcrossBatch(bool broadcastAcrossBatch) {}
+
+  bool getBroadcastAcrossBatch() const { return false; }
+
+  nvinfer1::TensorLocation getLocation() const { return location_; }
+
+  void setLocation(nvinfer1::TensorLocation location) { location_ = location; }
+  bool setDynamicRange(float min, float max) {
+    dynamic_range_max_ = max;
+    dynamic_range_min_ = min;
+    return true;
+  }
+
+  float getDynamicRange() const {
+    return (std::abs(dynamic_range_min_) + dynamic_range_max_) / 2.f;
+  }
+  bool dynamicRangeIsSet() const { return true; }
+
+  void resetDynamicRange() {
+    dynamic_range_min_ = 0.f;
+    dynamic_range_max_ = 0.f;
+  }
+  float getDynamicRangeMin() const { return dynamic_range_min_; }
+
+  float getDynamicRangeMax() const { return dynamic_range_max_; }
+
+  void setAllowedFormats(nvinfer1::TensorFormats formats) {}
+
+  nvinfer1::TensorFormats getAllowedFormats() const { return 1; }
+
+  bool isShapeTensor() const { return false; }
+  bool isExecutionTensor() const { return true; }
+
+ private:
+  nvinfer1::DataType trt_dtype_;
+  nvinfer1::Dims trt_dims_;
+  std::string name_;
+  nvinfer1::TensorLocation location_;
+  float dynamic_range_min_;
+  float dynamic_range_max_;
+};
+
+enum class TensorType : int { kTRT, kSIMPLE };
+
+class ITensorProxy {
+ public:
+  //! ITensor not owned
+  ITensorProxy(nvinfer1::ITensor* trt_tensor)
+      : trt_tensor_(trt_tensor), ttype_(TensorType::kTRT) {}
+
+  //! SimpleITensor owned
+  ITensorProxy(SimpleITensor* simple_itensor)
+      : simple_tensor_(simple_itensor), ttype_(TensorType::kSIMPLE) {}
+
+  //! SimpleITensor owned
+  explicit ITensorProxy(nvinfer1::DataType trt_dtype,
+                        const nvinfer1::Dims& trt_dims)
+      : simple_tensor_(std::unique_ptr<SimpleITensor>(
+            new SimpleITensor(trt_dtype, trt_dims))),
+        ttype_(TensorType::kSIMPLE) {}
+
+  //! Variants for testing purposes
+  ITensorProxy()
+      : simple_tensor_(std::unique_ptr<SimpleITensor>(new SimpleITensor())),
+        ttype_(TensorType::kSIMPLE) {}
+
+  explicit ITensorProxy(const nvinfer1::Dims& dims)
+      : simple_tensor_(std::unique_ptr<SimpleITensor>(new SimpleITensor(dims))),
+        ttype_(TensorType::kSIMPLE) {}
+
+  explicit ITensorProxy(const std::vector<int>& dims)
+      : simple_tensor_(std::unique_ptr<SimpleITensor>(new SimpleITensor(dims))),
+        ttype_(TensorType::kSIMPLE) {}
+
+  bool is_trt_tensor() const {
+    CHECK(validate());
+    return trt_tensor_ != nullptr;
+  }
+
+  bool is_simple_tensor() const {
+    CHECK(validate());
+    return simple_tensor_ != nullptr;
+  }
+
+  TensorType ttype() const { return ttype_; }
+
+  nvinfer1::ITensor* trt_tensor() const {
+    CHECK_NOTNULL(trt_tensor_);
+    CHECK(ttype_ == TensorType::kTRT);
+    return trt_tensor_;
+  }
+
+  SimpleITensor* simple_tensor() const {
+    CHECK_NOTNULL(simple_tensor_);
+    CHECK(ttype_ == TensorType::kSIMPLE);
+    return simple_tensor_.get();
+  }
+
+  void setName(const char* name) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setName(name);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setName(name);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  const char* getName() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getName();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getName();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  void setDimensions(nvinfer1::Dims dimensions) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setDimensions(dimensions);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setDimensions(dimensions);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  nvinfer1::Dims getDimensions() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getDimensions();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getDimensions();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  void setType(nvinfer1::DataType type) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setType(type);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setType(type);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  nvinfer1::DataType getType() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getType();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getType();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool isNetworkInput() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->isNetworkInput();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->isNetworkInput();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool isNetworkOutput() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->isNetworkOutput();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->isNetworkOutput();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  void setBroadcastAcrossBatch(bool broadcastAcrossBatch) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setBroadcastAcrossBatch(broadcastAcrossBatch);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setBroadcastAcrossBatch(broadcastAcrossBatch);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool getBroadcastAcrossBatch() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getBroadcastAcrossBatch();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getBroadcastAcrossBatch();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  nvinfer1::TensorLocation getLocation() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getLocation();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getLocation();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  void setLocation(nvinfer1::TensorLocation location) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setLocation(location);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setLocation(location);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool setDynamicRange(float min, float max) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setDynamicRange(min, max);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setDynamicRange(min, max);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool dynamicRangeIsSet() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->dynamicRangeIsSet();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->dynamicRangeIsSet();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  void resetDynamicRange() {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->resetDynamicRange();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->resetDynamicRange();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+  float getDynamicRangeMin() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getDynamicRangeMin();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getDynamicRangeMin();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  float getDynamicRangeMax() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getDynamicRangeMax();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getDynamicRangeMax();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+#if !IS_TRT_VERSION_GE(8, 0, 0, 0)
+  float getDynamicRange() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getDynamicRange();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getDynamicRange();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+#endif
+  void setAllowedFormats(nvinfer1::TensorFormats formats) {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->setAllowedFormats(formats);
+      case TensorType::kSIMPLE:
+        return simple_tensor_->setAllowedFormats(formats);
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  nvinfer1::TensorFormats getAllowedFormats() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->getAllowedFormats();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->getAllowedFormats();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool isShapeTensor() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->isShapeTensor();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->isShapeTensor();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+  bool isExecutionTensor() const {
+    switch (ttype_) {
+      case TensorType::kTRT:
+        return trt_tensor_->isExecutionTensor();
+      case TensorType::kSIMPLE:
+        return simple_tensor_->isExecutionTensor();
+    }
+    LOG(FATAL) << "Unsupported itensor_ type";
+  }
+
+ private:
+  bool validate() const {
+    return (trt_tensor_ && !simple_tensor_) || (!trt_tensor_ && simple_tensor_);
+  }
+
+  // When ITensorProxy represents an ITensor, the ITensor can be either passed
+  // by the caller via the constructor that takes an ITensor* as parameter, or
+  // be created as a SimpleITensor.
+  //
+  // In the first case, the ITensor pointer is stored in 'tensor_' below, and
+  // the ITensor itself is not owned by this class. This method is used by
+  // Converter (e.g. AddInputTensor) and op converters during TRT network
+  // construction, where the TRT network owns the ITensor.
+  //
+  nvinfer1::ITensor* trt_tensor_ = nullptr;  // Not owned.
+  // In the second case, the created SimpleITensor is stored in
+  // 'simple_itensor_' below and is owned by this class. SimpleITensor is a fake
+  // implementation of ITensor and is used for testing and by TrtNodeValidator
+  //  to validate the graph nodes.
+  std::shared_ptr<SimpleITensor> simple_tensor_ = nullptr;
+
+  TensorType ttype_;
+};
+
+class ITensorProxyPtr {
+ public:
+  ITensorProxyPtr(std::nullptr_t) : p_(nullptr) {}
+  ITensorProxyPtr(ITensorProxy* p) : p_(p) {}
+  ITensorProxyPtr(nvinfer1::ITensor* p) : p_(new ITensorProxy(p)) {}
+  ITensorProxyPtr(SimpleITensor* p) : p_(new ITensorProxy(p)) {}
+
+  ITensorProxyPtr() : p_(new ITensorProxy()) {}
+  ITensorProxyPtr(const nvinfer1::Dims& dims) : p_(new ITensorProxy(dims)) {}
+  ITensorProxyPtr(const std::vector<int>& dims) : p_(new ITensorProxy(dims)) {}
+
+  std::shared_ptr<ITensorProxy> p_{nullptr};
+  ITensorProxy* operator->() { return p_.get(); }
+  ITensorProxy* operator->() const { return p_.get(); }
+  ITensorProxy* operator*() { return p_.get(); }
+  ITensorProxy* operator*() const { return p_.get(); }
+};
+
+inline bool operator==(const ITensorProxyPtr& p1, const ITensorProxyPtr& p2) {
+  if (p1.p_ == nullptr) {
+    return p2.p_ == nullptr;
+  }
+  if (p2.p_ == nullptr) {
+    return p1.p_ == nullptr;
+  }
+  return (p1->ttype() == p2->ttype()) &&
+         ((p1->ttype() == TensorType::kTRT &&
+           p1->trt_tensor() == p2->trt_tensor()) ||
+          (p1->ttype() == TensorType::kSIMPLE &&
+           p1->simple_tensor() == p2->simple_tensor()));
+}
+
+inline bool operator!=(const ITensorProxyPtr& p1, const ITensorProxyPtr& p2) {
+  return !(p1 == p2);
+}
+
+struct ITensorProxyHash {
+  size_t operator()(const ITensorProxyPtr& tensor) const {
+    return reinterpret_cast<std::uintptr_t>(tensor.p_.get());
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_TENSOR_PROXY_H_

tensorflow/compiler/tf2tensorrt/utils/trt_testutils.cc

@@ -0,0 +1,76 @@
+/* Copyright 2021 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 "tensorflow/compiler/tf2tensorrt/utils/trt_testutils.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+
+namespace tensorflow {
+
+namespace tensorrt {
+namespace convert {
+
+::testing::Matcher<std::vector<float>> ArrayFloatNear(
+    const std::vector<float>& values, float max_abs_error, bool nan_sensitive) {
+  std::vector<::testing::Matcher<float>> matchers;
+  matchers.reserve(values.size());
+  for (const float& v : values) {
+    if (nan_sensitive) {
+      matchers.emplace_back(::testing::NanSensitiveFloatNear(v, max_abs_error));
+    } else if (max_abs_error == 0) {
+      matchers.emplace_back(::testing::FloatEq(v));
+    } else {
+      EXPECT_GE(max_abs_error, 0);
+      matchers.emplace_back(::testing::FloatNear(v, max_abs_error));
+    }
+  }
+  return ::testing::ElementsAreArray(matchers);
+}
+
+nvinfer1::Dims CreateDims(const std::vector<int>& d) {
+  nvinfer1::Dims dims;
+  dims.nbDims = d.size();
+  for (int i = 0; i < d.size(); ++i) {
+    dims.d[i] = d[i];
+  }
+  return dims;
+}
+
+NodeDef MakeNodeDef(const std::string& name, const std::string& op,
+                    const std::vector<std::string>& inputs,
+                    const std::map<std::string, AttrValue> attrs) {
+  NodeDef node_def;
+  node_def.set_name(name);
+  node_def.set_op(op);
+  for (const auto& input : inputs) {
+    node_def.add_input(input);
+  }
+  for (const auto& attr : attrs) {
+    (*node_def.mutable_attr())[attr.first] = attr.second;
+  }
+  return node_def;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/utils/trt_testutils.h

@@ -0,0 +1,183 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_TESTUTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_TESTUTILS_H_
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <algorithm>
+#include <map>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_format.h"
+#include "absl/types/span.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h"
+#include "tensorflow/core/framework/node_def.pb.h"  // NOLINT
+#include "tensorflow/core/framework/tensor.pb.h"    // NOLINT
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+// Creates a node with the given op, inputs, and attributes.
+NodeDef MakeNodeDef(const std::string& name, const std::string& op,
+                    const std::vector<std::string>& inputs,
+                    const std::map<std::string, AttrValue> attrs = {});
+
+// Creates a constant node with the given name and values arranged in the given
+// shape.
+template <typename T>
+NodeDef MakeConstNodeDef(const std::string& name, const std::vector<T>& vals,
+                         const TensorShape& shape) {
+  Scope s = Scope::NewRootScope();
+  Tensor t = test::AsTensor<T>(vals, shape);
+  auto const_op = ops::Const(s.WithOpName(name), t);
+  return const_op.node()->def();
+}
+
+// Creates a constant node with the given name and values, assuming a 1-D shape.
+template <typename T>
+NodeDef MakeConstNodeDef(const std::string& name, const std::vector<T>& vals) {
+  TensorShape shape;
+  const std::vector<int32> shape_dims = {static_cast<int32>(vals.size())};
+  TF_EXPECT_OK(TensorShapeUtils::MakeShape(shape_dims, &shape));
+  return MakeConstNodeDef(name, vals, shape);
+}
+
+// Creates an nvinfer1::Dims struct from the given vector.
+nvinfer1::Dims CreateDims(const std::vector<int>& d);
+
+// A gmock matcher that check that elements of a float vector match to a given
+// tolerance.
+::testing::Matcher<std::vector<float>> ArrayFloatNear(
+    const std::vector<float>& values, float max_abs_error = 1e-5,
+    bool nan_sensitive = false);
+
+// nvinfer1::Dims gMock matchers
+
+// matches nvinfer1::Dims to initializer list or vector of ints
+// Example: EXPECT_THAT(my_dims, DimsAreArray({1, 2, 3}))
+MATCHER_P(DimsAreArrayHelper, array_value,
+          absl::StrFormat("%s [%s]", negation ? "are" : "are not",
+                          ::testing::PrintToString(array_value))) {
+  if (arg.nbDims != array_value.size()) return false;
+  for (int i = 0; i < arg.nbDims; ++i) {
+    if (arg.d[i] != array_value[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+using DimsAreArray = DimsAreArrayHelperMatcherP<std::vector<int>>;
+
+// nvinfer1::INetworkDefinition gMock matchers
+
+// Checks that layer names are equal to initializer list or vector of strings.
+// Example: EXPECT_THAT(my_network, LayerNamesAreArray({"conv1", "conv2"}))
+MATCHER_P(LayerNamesAreArrayHelper, array_value,
+          absl::StrFormat("layer names %s [%s]", negation ? "are" : "are not",
+                          ::testing::PrintToString(array_value))) {
+  if (array_value.size() != arg->getNbLayers()) return false;
+  for (int i = 0; i < arg->getNbLayers(); ++i) {
+    if (arg->getLayer(i)->getName() == nullptr) {
+      return false;
+    }
+  }
+  return true;
+}
+using LayerNamesAreArray =
+    LayerNamesAreArrayHelperMatcherP<std::vector<std::string>>;
+
+// Checks layer names are all non-empty.
+MATCHER(LayerNamesNonEmpty, "") {
+  for (int i = 0; i < arg->getNbLayers(); ++i) {
+    if (arg->getLayer(i)->getName() == nullptr) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// TRT_ShapedWeights gMock matchers.
+
+// Checks that the weight dimensions are values are equal to the given values.
+// Example: EXPECT_THAT(my_weights,
+//                      ShapedWeightsHasDimsAndValues({1, 2},{1.0f, 2.0f}))
+MATCHER_P2(ShapedWeightsHasDimsAndValuesHelper, dims_vec, expected_values, "") {
+  DimsAdapter dims(dims_vec);
+  if (arg.Shape() != dims) {
+    return false;
+  }
+  if (arg.count() != expected_values.size()) {
+    return false;
+  }
+  using T = typename decltype(expected_values)::value_type;
+  const T* actual_values = arg.template GetPointer<T>();
+  for (int i = 0; i < expected_values.size(); ++i) {
+    if (expected_values[i] != actual_values[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+template <typename T>
+using ShapedWeightsHasDimsAndValues =
+    ShapedWeightsHasDimsAndValuesHelperMatcherP2<std::vector<int>,
+                                                 std::vector<T>>;
+
+// std::vector convenience utilities.
+
+// Creates a new vector by casting all values of the given InCType vector to
+// OutCType.
+template <typename InCType, typename OutCType>
+std::vector<OutCType> CastVector(
+    const gtl::ArraySlice<InCType>& vals) {  // non-absl ok
+  std::vector<OutCType> res(vals.size());
+  std::transform(vals.begin(), vals.end(), res.begin(),
+                 [](const InCType in_val) -> OutCType {
+                   return static_cast<OutCType>(in_val);
+                 });
+  return res;
+}
+
+// Creates a new vector of the given size and fills it with an increasing
+// sequence starting from the given start_value using std::iota.
+template <typename CType>
+std::vector<CType> CreateVectorIota(int size, CType start_value = CType(0)) {
+  std::vector<CType> res(size);
+  std::iota(res.begin(), res.end(), start_value);
+  return res;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_TESTUTILS_H_

tensorflow/compiler/tf2tensorrt/utils/trt_testutils_test.cc

@@ -0,0 +1,99 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_testutils.h"
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+using ::testing::AllOf;
+using ::testing::AnyOf;
+using ::testing::Eq;
+using ::testing::Not;
+
+TEST(TrtDimsMatcher, ParameterizedMatchers) {
+  EXPECT_THAT(nvinfer1::Dims4(1, 2, 3, 4), DimsAreArray({1, 2, 3, 4}));
+  // Check empty dims.
+  EXPECT_THAT(nvinfer1::Dims{}, Not(DimsAreArray({1, 2})));
+  std::vector<int> empty_dims;
+  EXPECT_THAT(nvinfer1::Dims{}, DimsAreArray(empty_dims));
+  // Check mismatching values.
+  EXPECT_THAT(nvinfer1::Dims4(1, 2, 3, 4), Not(DimsAreArray({1, 2, 3, 5})));
+  // Check mismatching number of arguments.
+  EXPECT_THAT(nvinfer1::Dims4(1, 2, 3, 4), Not(DimsAreArray({1, 2, 5})));
+}
+
+TEST(TrtDimsMatcher, EqualityMatcher) {
+  EXPECT_THAT(nvinfer1::Dims4(1, 2, 3, 4), Eq(nvinfer1::Dims4(1, 2, 3, 4)));
+  // Check empty dims.
+  EXPECT_THAT(nvinfer1::Dims{}, Eq(nvinfer1::Dims()));
+  // Check empty Dims is not equal to DimsHW, since their sizes differ.
+  EXPECT_THAT(nvinfer1::Dims{}, Not(Eq(nvinfer1::DimsHW())));
+  // Check mismatching values.
+  EXPECT_THAT(nvinfer1::Dims4(1, 2, 3, 4),
+              Not(Eq(nvinfer1::Dims4(1, 2, 3, 3))));
+  // Check mismatching number of arguments.
+  EXPECT_THAT(nvinfer1::Dims4(1, 2, 3, 4), Not(Eq(nvinfer1::Dims2(1, 2))));
+}
+
+TEST(INetworkDefinitionMatchers, CorrectlyMatch) {
+  Logger& logger = *Logger::GetLogger();
+  TrtUniquePtrType<nvinfer1::IBuilder> builder(
+      nvinfer1::createInferBuilder(logger));
+  TrtUniquePtrType<nvinfer1::INetworkDefinition> network(
+      builder->createNetworkV2(0L));
+
+  // Empty network checks.
+  EXPECT_THAT(network.get(), AllOf(Not(LayerNamesAreArray({"some layer"})),
+                                   LayerNamesNonEmpty()));
+
+  // Add the input and FC layers.
+  nvinfer1::Weights weights;
+  weights.type = nvinfer1::DataType::kFLOAT;
+  std::array<float, 1> vals;
+  weights.values = vals.data();
+  weights.count = 1;
+  auto input = network->addInput("input-tensor", nvinfer1::DataType::kFLOAT,
+                                 nvinfer1::Dims3{1, 1, 1});
+  ASSERT_NE(input, nullptr);
+
+  const char* fc_layer_name = "my-fc-layer";
+  auto layer = network->addFullyConnected(*input, 1, weights, weights);
+  ASSERT_NE(layer, nullptr);
+  layer->setName(fc_layer_name);
+
+  // Check layer names.
+  EXPECT_THAT(network.get(),
+              AllOf(LayerNamesNonEmpty(), LayerNamesAreArray({fc_layer_name})));
+
+  // Add layer with default name and check layer name.
+  layer = network->addFullyConnected(*input, 1, weights, weights);
+  EXPECT_THAT(network.get(), AllOf(LayerNamesNonEmpty(),
+                                   Not(LayerNamesAreArray({fc_layer_name}))));
+}
+
+}  // namespace convert
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT