chore: import upstream snapshot with attribution

tensorflow/compiler/tf2tensorrt/convert/algorithm_selector.cc

@@ -0,0 +1,271 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "tensorflow/compiler/tf2tensorrt/convert/algorithm_selector.h"
+
+#include <utility>
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/core/util/env_var.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+// getAlgorithmIOInfo is deprecated in TRT >= 8, replaced by
+// getAlgorithmIOInfoByIndex.
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+#define ALGORITHM_IO_INFO_BY_IDX(alg, idx) *(alg).getAlgorithmIOInfoByIndex(idx)
+#else
+#define ALGORITHM_IO_INFO_BY_IDX(alg, idx) (alg).getAlgorithmIOInfo(idx)
+#endif
+
+namespace nvinfer1 {
+
+std::ostream& operator<<(std::ostream& os,
+                         const nvinfer1::IAlgorithmContext& ctx) {
+  os << "AlgorithmContext(name=" << ctx.getName()
+     << ",nbInputs=" << ctx.getNbInputs() << ",nbOutputs=" << ctx.getNbOutputs()
+     << ")";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, const nvinfer1::IAlgorithm& alg) {
+  const nvinfer1::IAlgorithmVariant& variant = alg.getAlgorithmVariant();
+  os << "Algorithm(" << "variant.implementation=" << variant.getImplementation()
+     << ",variant.tactic=" << variant.getTactic()
+     << ",timingMSec=" << alg.getTimingMSec()
+     << ",workspaceSize=" << alg.getWorkspaceSize() << ")";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os,
+                         const nvinfer1::IAlgorithmIOInfo& info) {
+  os << "IOTensor(format=" << info.getTensorFormat()
+     << ",dtype=" << info.getDataType() << ",strides=" << info.getStrides()
+     << ")";
+  return os;
+}
+}  // namespace nvinfer1
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+bool operator>=(const AlgorithmSelectorImpl::TRTVersion& lhs,
+                const AlgorithmSelectorImpl::TRTVersion& rhs) {
+  if (lhs[0] > rhs[0]) return true;
+  if (lhs[0] == rhs[0] && lhs[1] > rhs[1]) return true;
+  if (lhs[0] == rhs[0] && lhs[1] == rhs[1] && lhs[2] > rhs[2]) return true;
+  if (lhs[0] == rhs[0] && lhs[1] == rhs[1] && lhs[2] == rhs[2] &&
+      lhs[3] >= rhs[3]) {
+    return true;
+  }
+  return false;
+}
+
+bool AlgorithmSelectorImpl::IsTrtVersionGE(const TRTVersion& version) const {
+  return version_ >= version;
+}
+
+bool AlgorithmSelectorImpl::IsShuffleLayer(ImplementationID id) const {
+  if (IsTrtVersionGE({8, 2, 0, 0})) {
+    return id == 0x80000000 + 13;
+  }
+  if (IsTrtVersionGE({8, 0, 0, 0})) {
+    return id == 0x80000000 + 14;
+  }
+  if (IsTrtVersionGE({7, 2, 0, 0})) {
+    return id == 0x80000000 + 16;
+  }
+  return id == 18;
+}
+
+std::set<AlgorithmSelectorImpl::TacticID>
+AlgorithmSelectorImpl::GetBannedTRT72TuringTactics() {
+  static const std::set<TacticID> banned_turing_72{
+      // turing_fp16_s1688cudnn_fp16_128x128_ldg8_relu_f2f_exp_medium_nhwc_gelu_tn_v1
+      -5927686925093575778,
+      // turing_fp16_s1688cudnn_fp16_128x128_ldg8_relu_f2f_exp_interior_nhwc_gelu_tn_v1
+      -3848538574386518527,
+      // turing_fp16_s1688cudnn_fp16_128x128_ldg8_relu_f2f_exp_small_nhwc_gelu_tn_v1
+      -959009792490796596};
+  return banned_turing_72;
+}
+
+bool AlgorithmSelectorImpl::IsBannedTactic(TacticID id) const {
+  // Disable problematic FP16-Turing tactics in TensorRT 7.2.
+  if (IsTrtVersionGE({7, 2, 0, 0}) && !IsTrtVersionGE({8, 0, 0, 0})) {
+    auto banned_turing_72 = GetBannedTRT72TuringTactics();
+    return banned_turing_72.find(id) != banned_turing_72.end();
+  }
+  return false;
+}
+
+bool AlgorithmSelectorImpl::AllowShuffleAlgorithm(
+    TacticID tactic, nvinfer1::DataType input_dtype,
+    nvinfer1::TensorFormat input_format) const {
+  if (IsTrtVersionGE({8, 0, 0, 0}) && !IsTrtVersionGE({8, 0, 3, 0})) {
+    // Reject shuffle node when input format is linear row major INT8
+    // format in TensorRT 8.0 GA.
+    return !(input_format == nvinfer1::TensorFormat::kLINEAR &&
+             input_dtype == nvinfer1::DataType::kINT8);
+  }
+
+  if (IsTrtVersionGE({7, 2, 0, 0}) && !IsTrtVersionGE({8, 0, 0, 0})) {
+    // For TRT 7.2, accept shuffle node when input format is not 32-wide
+    // channel vectorized row major FP32 format
+    return !(input_format == nvinfer1::TensorFormat::kCHW32 &&
+             input_dtype == nvinfer1::DataType::kFLOAT);
+  }
+  return true;
+}
+
+bool AlgorithmSelectorImpl::IsAlgorithmSelectorRequired() const {
+  // If we are in turing for TensorRT 7.2, we need the  selector for shuffle and
+  // avoiding specific Turing tactics.
+  if (IsTrtVersionGE({7, 2, 0, 0}) && !IsTrtVersionGE({8, 0, 0, 0})) {
+    return true;
+  }
+
+  // If we are in TensorRT 8.0 GA, we want to reject certain types of shuffles.
+  if (IsTrtVersionGE({8, 0, 0, 0}) && !IsTrtVersionGE({8, 0, 3, 0})) {
+    return true;
+  }
+
+  return false;
+}
+
+namespace {
+
+string FormatAlgorithmList(const nvinfer1::IAlgorithmContext& ctx,
+                           absl::Span<const nvinfer1::IAlgorithm* const> algs) {
+  return absl::StrFormat(
+      "%s:\n\t%s", absl::FormatStreamed(ctx),
+      absl::StrJoin(
+          algs, "\n\t",
+          [&ctx](std::string* out, const nvinfer1::IAlgorithm* const alg) {
+            absl::StrAppendFormat(out, "%s", absl::FormatStreamed(*alg));
+            for (int i = 0; i < ctx.getNbInputs() + ctx.getNbOutputs(); i++) {
+              absl::StrAppendFormat(
+                  out, "\n\t\t%s",
+                  absl::FormatStreamed(ALGORITHM_IO_INFO_BY_IDX(*alg, i)));
+            }
+          }));
+}
+
+}  // namespace
+
+TftrtAlgorithmSelector::TftrtAlgorithmSelector()
+    : fixed_algorithm_idx_(GetFixedAlgorithmID()),
+      selector_(AlgorithmSelectorImpl::CompileTimeTRTVersion()) {}
+
+std::optional<int64_t> TftrtAlgorithmSelector::GetFixedAlgorithmID() {
+  int64_t trt_algorithm_idx = 0;
+  constexpr auto null_idx =
+      std::numeric_limits<decltype(trt_algorithm_idx)>::min();
+  Status status = tensorflow::ReadInt64FromEnvVar("TF_TRT_FIXED_ALGORITHM_ID",
+                                                  /*default_val=*/null_idx,
+                                                  &trt_algorithm_idx);
+  if (!status.ok()) {
+    LOG(ERROR) << status;
+    return std::nullopt;
+  }
+  if (trt_algorithm_idx != null_idx) {
+    return std::max(static_cast<int32_t>(trt_algorithm_idx), 0);
+  }
+  return std::nullopt;
+}
+
+bool TftrtAlgorithmSelector::AlgorithmPolicy(
+    const nvinfer1::IAlgorithmContext& context,
+    const nvinfer1::IAlgorithm& alg) const {
+  const nvinfer1::IAlgorithmVariant& variant = alg.getAlgorithmVariant();
+
+  // Check if this tactic ID is banned.
+  TacticID tactic_id = variant.getTactic();
+  if (selector_.IsBannedTactic(tactic_id)) {
+    return false;
+  }
+
+  if (selector_.IsShuffleLayer(variant.getImplementation())) {
+    return selector_.AllowShuffleAlgorithm(
+        tactic_id, alg.getAlgorithmIOInfo(0).getDataType(),
+        alg.getAlgorithmIOInfo(0).getTensorFormat());
+  }
+  return true;
+}
+
+int32_t TftrtAlgorithmSelector::selectAlgorithms(
+    const nvinfer1::IAlgorithmContext& algoContext,
+    const nvinfer1::IAlgorithm* const* algoChoices, int32_t nbChoices,
+    int32_t* selection) noexcept {
+  if (fixed_algorithm_idx_) {
+    LOG(WARNING) << "Forcing TRT algorithm selection to: ID = "
+                 << *fixed_algorithm_idx_;
+    selection[0] = std::min(*fixed_algorithm_idx_, nbChoices - 1);
+    return 1;
+  }
+
+  int num_selections = 0;
+
+  VLOG(1) << "Algorithm selection choices: "
+          << FormatAlgorithmList(algoContext,
+                                 absl::MakeSpan(algoChoices, nbChoices));
+
+  for (int i = 0; i < nbChoices; i++) {
+    const nvinfer1::IAlgorithm& alg = *algoChoices[i];
+
+    // Check layer-specific issues.
+    if (!AlgorithmPolicy(algoContext, alg)) {
+      LOG(WARNING) << absl::StrFormat("Rejecting Algorithm: %s ",
+                                      absl::FormatStreamed(alg));
+      continue;
+    }
+    selection[num_selections++] = i;
+  }
+  return num_selections;
+}
+
+// Called by TensorRT to report choices it made.
+void TftrtAlgorithmSelector::reportAlgorithms(
+    const nvinfer1::IAlgorithmContext* const* algoContexts,
+    const nvinfer1::IAlgorithm* const* algoChoices,
+    int32_t nbAlgorithms) noexcept {
+  if (VLOG_IS_ON(1)) {
+    string selection_msg = "Algorithms selected:\n";
+    for (int i = 0; i < nbAlgorithms; i++) {
+      absl::StrAppend(&selection_msg,
+                      FormatAlgorithmList(*algoContexts[i],
+                                          absl::MakeSpan(algoChoices + i, 1)));
+    }
+    VLOG(1) << selection_msg;
+  }
+}
+
+std::unique_ptr<TftrtAlgorithmSelector> MaybeCreateAlgorithmSelector() {
+  auto selector = std::make_unique<TftrtAlgorithmSelector>();
+
+  if (selector->IsRequired()) {
+    return selector;
+  }
+
+  return nullptr;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/algorithm_selector.h

@@ -0,0 +1,121 @@
+/* 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_CONVERT_ALGORITHM_SELECTOR_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_ALGORITHM_SELECTOR_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include <array>
+#include <memory>
+#include <set>
+
+#include "absl/types/optional.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+// Implements core algorithm selection logic in a testable manner. The policy
+// implemented depends on the given TRT version. We have this class because TRT
+// interfaces make it difficult to directly test an IAlgorithmSelector
+// implementation.
+class AlgorithmSelectorImpl {
+ public:
+  using TRTVersion = std::array<int, 4>;
+  using ImplementationID = int64_t;
+  using TacticID = int64_t;
+
+  static constexpr TRTVersion CompileTimeTRTVersion() {
+    return TRTVersion{NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH,
+                      NV_TENSORRT_BUILD};
+  }
+
+  explicit AlgorithmSelectorImpl(
+      const TRTVersion& version = CompileTimeTRTVersion())
+      : version_(version) {}
+
+  bool IsShuffleLayer(ImplementationID id) const;
+
+  bool IsBannedTactic(TacticID id) const;
+
+  // Returns true if the algorithm implementing the IShuffleLayer is acceptable.
+  bool AllowShuffleAlgorithm(TacticID tactic, nvinfer1::DataType input_dtype,
+                             nvinfer1::TensorFormat input_format) const;
+
+  bool IsTrtVersionGE(const TRTVersion& version) const;
+
+  // Returns true if we know at compile time that the algorithm selector
+  // should be required. This is a conservative estimate.
+  bool IsAlgorithmSelectorRequired() const;
+
+  static std::set<TacticID> GetBannedTRT72TuringTactics();
+
+ private:
+  TRTVersion version_;
+};
+
+// Implements the TRT IAlgorithmSelector interface. The method
+// "selectAlgorithms" selects allowable algorithms for each layer, and
+// "reportAlgorithms" summarizes the algorithms selected by TensorRT.
+class TftrtAlgorithmSelector : public nvinfer1::IAlgorithmSelector {
+ private:
+  using TacticID = AlgorithmSelectorImpl::TacticID;
+
+  // An index we should choose for all algorithms. Used for debugging.
+  std::optional<int32_t> fixed_algorithm_idx_;
+
+  AlgorithmSelectorImpl selector_;
+
+ public:
+  TftrtAlgorithmSelector();
+
+  // If the environment variable TF_TRT_FIXED_ALGORITHM_ID is empty, this
+  // function returns nullopt. Otherwise, it returns the specified number.
+  static std::optional<int64_t> GetFixedAlgorithmID();
+
+  // Returns true if the algorithm associated with context is acceptable.
+  bool AlgorithmPolicy(const nvinfer1::IAlgorithmContext& context,
+                       const nvinfer1::IAlgorithm& alg) const;
+
+  // This function fills the array "selection" with the indices of selected
+  // algorithm candidates from "algoChoices", each of which is an implementation
+  // for the kernel described by the given IAlgorithmContext. It should return a
+  // number in [0, nbChoices] indicating the number of selected indices. If 0 is
+  // returned, TensorRT will use its default selection mechanism.
+  int32_t selectAlgorithms(const nvinfer1::IAlgorithmContext& algoContext,
+                           const nvinfer1::IAlgorithm* const* algoChoices,
+                           int32_t nbChoices,
+                           int32_t* selection) noexcept override;
+
+  // Called by TensorRT to report choices it made.
+  void reportAlgorithms(const nvinfer1::IAlgorithmContext* const* algoContexts,
+                        const nvinfer1::IAlgorithm* const* algoChoices,
+                        int32_t nbAlgorithms) noexcept override;
+
+  bool IsRequired() const {
+    return selector_.IsAlgorithmSelectorRequired() ||
+           fixed_algorithm_idx_ != std::nullopt;
+  }
+};
+
+// Returns an initialized AlgorithmSelector if an algorithm selector is required
+// for the current TRT version. Otherwise, returns nullptr.
+std::unique_ptr<TftrtAlgorithmSelector> MaybeCreateAlgorithmSelector();
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_ALGORITHM_SELECTOR_H_

tensorflow/compiler/tf2tensorrt/convert/algorithm_selector_test.cc

@@ -0,0 +1,97 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/algorithm_selector.h"
+
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+TEST(TestAlgorithmSelector, TensorRT7_1) {
+  // Verify that the algorithm selector for TRT 7.1 is not required.
+  AlgorithmSelectorImpl sel71({7, 1, 3, 4});
+  ASSERT_FALSE(sel71.IsAlgorithmSelectorRequired());
+}
+
+TEST(TestAlgorithmSelector, TensorRT7_2) {
+  // Verify that the algorithm selector for TRT 7.2 is required.
+  AlgorithmSelectorImpl sel72({7, 2, 0, 0});
+  ASSERT_TRUE(sel72.IsAlgorithmSelectorRequired());
+
+  // Check that the correct tactics are banned.
+  auto turing_tactics = AlgorithmSelectorImpl::GetBannedTRT72TuringTactics();
+
+  for (auto id : turing_tactics) {
+    EXPECT_TRUE(sel72.IsBannedTactic(id));
+  }
+
+  // Check that a bad shuffle format is banned.
+  EXPECT_FALSE(sel72.AllowShuffleAlgorithm(0, nvinfer1::DataType::kFLOAT,
+                                           nvinfer1::TensorFormat::kCHW32));
+
+  // Check that other formats are not banned.
+  EXPECT_TRUE(sel72.AllowShuffleAlgorithm(0, nvinfer1::DataType::kHALF,
+                                          nvinfer1::TensorFormat::kCHW32));
+  EXPECT_TRUE(sel72.AllowShuffleAlgorithm(0, nvinfer1::DataType::kINT32,
+                                          nvinfer1::TensorFormat::kCHW32));
+  EXPECT_TRUE(sel72.AllowShuffleAlgorithm(0, nvinfer1::DataType::kFLOAT,
+                                          nvinfer1::TensorFormat::kCHW16));
+}
+
+TEST(TestAlgorithmSelector, TensorRT8_0) {
+  // Verify that the algorithm selector for TRT 8.0 is required.
+  AlgorithmSelectorImpl sel80({8, 0, 1, 6});
+  ASSERT_TRUE(sel80.IsAlgorithmSelectorRequired());
+
+  // Check that the turing 7.2 tactics are not banned.
+  auto turing_tactics = AlgorithmSelectorImpl::GetBannedTRT72TuringTactics();
+  for (auto id : turing_tactics) {
+    EXPECT_FALSE(sel80.IsBannedTactic(id));
+  }
+
+  // Check that a bad shuffle format is banned.
+  EXPECT_FALSE(sel80.AllowShuffleAlgorithm(0, nvinfer1::DataType::kINT8,
+                                           nvinfer1::TensorFormat::kLINEAR));
+
+  // Check that other formats are not banned.
+  EXPECT_TRUE(sel80.AllowShuffleAlgorithm(0, nvinfer1::DataType::kHALF,
+                                          nvinfer1::TensorFormat::kLINEAR));
+  EXPECT_TRUE(sel80.AllowShuffleAlgorithm(0, nvinfer1::DataType::kINT32,
+                                          nvinfer1::TensorFormat::kLINEAR));
+  EXPECT_TRUE(sel80.AllowShuffleAlgorithm(0, nvinfer1::DataType::kFLOAT,
+                                          nvinfer1::TensorFormat::kLINEAR));
+  EXPECT_TRUE(sel80.AllowShuffleAlgorithm(0, nvinfer1::DataType::kINT8,
+                                          nvinfer1::TensorFormat::kCHW16));
+  EXPECT_TRUE(sel80.AllowShuffleAlgorithm(0, nvinfer1::DataType::kINT8,
+                                          nvinfer1::TensorFormat::kCHW32));
+}
+
+TEST(TestAlgorithmSelector, TensorRT8_2) {
+  // Verify that the algorithm selector for TRT 8.0 is required.
+  AlgorithmSelectorImpl sel({8, 2, 0, 0});
+  ASSERT_FALSE(sel.IsAlgorithmSelectorRequired());
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc

@@ -0,0 +1,944 @@
+/* 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/convert/convert_graph.h"
+
+#include <fstream>
+#include <list>
+#include <map>
+#include <set>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/segment/segment.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_process_state.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph_to_functiondef.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
+#include "tensorflow/core/grappler/costs/graph_properties.h"
+#include "tensorflow/core/grappler/devices.h"
+#include "tensorflow/core/grappler/optimizers/meta_optimizer.h"
+#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
+#include "tensorflow/core/protobuf/device_properties.pb.h"  // NOLINT
+#include "tensorflow/core/protobuf/rewriter_config.pb.h"  // NOLINT
+#include "tensorflow/core/util/device_name_utils.h"
+#include "tensorflow/tools/graph_transforms/transform_utils.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 {
+namespace convert {
+
+using absl::StrAppend;
+using absl::StrCat;
+using ::tensorflow::tensorrt::segment::ClusterProperty;
+using ::tensorflow::tensorrt::segment::NodePtrCompare;
+using ::tensorflow::tensorrt::segment::Segment;
+
+namespace {
+
+Status BuildNodeMap(const Graph& graph,
+                    std::unordered_map<string, Node*>* node_map) {
+  for (auto* node : graph.op_nodes()) {
+    if (!node_map->insert({node->name(), node}).second) {
+      return errors::AlreadyExists("Node name is not unique in graph: " +
+                                   node->name());
+    }
+  }
+  return OkStatus();
+}
+
+EngineInfo::EngineType GetEngineType(
+    const TRTOptimizationPass::ConversionParams& params) {
+  return (params.is_dynamic_op || params.use_calibration)
+             ? EngineInfo::EngineType::TRTDynamic
+             : EngineInfo::EngineType::TRTStatic;
+}
+
+// Returns true when use_implicit_batch is false or when we are building dynamic
+// engine, to allow unknown size for dimensions rather than dimension 0.
+bool AllowDynamicNonBatchDimension(
+    const TRTOptimizationPass::ConversionParams& params) {
+  return !params.use_implicit_batch ||
+         GetEngineType(params) == EngineInfo::EngineType::TRTDynamic;
+}
+
+struct EdgePtrCompare {
+  bool operator()(const Edge* lhs, const Edge* rhs) const {
+    return lhs->id() < rhs->id();
+  }
+};
+
+// TODO(laigd): instead of deciding the device here, the converter should accept
+// a device name as one of the conversion parameter so users can control on
+// which device they want to run the conversion.
+std::pair<TfDeviceId, PlatformDeviceId> GetFirstValidDeviceId() {
+  for (int tf_device_id_value = 0; tf_device_id_value < 100;
+       ++tf_device_id_value) {
+    TfDeviceId tf_device_id(tf_device_id_value);
+    PlatformDeviceId platform_device_id;
+    Status s =
+        GpuIdManager::TfToPlatformDeviceId(tf_device_id, &platform_device_id);
+    if (s.ok()) {
+      VLOG(1) << "Found TF GPU " << tf_device_id.value() << " at cuda device "
+              << platform_device_id.value();
+      return std::make_pair(tf_device_id, platform_device_id);
+    }
+  }
+  LOG(ERROR) << "Could not find any TF GPUs";
+  return std::make_pair(TfDeviceId(-1), PlatformDeviceId(-1));
+}
+
+// Returns false for const nodes (we intend to drop control edges from those).
+bool ShallKeepControlEdgeFrom(const Node* input_node) {
+  if (!input_node) {
+    LOG(ERROR) << "Node pointer is null, this should not happen";
+    return false;
+  }
+  return input_node->type_string() != "Const";
+}
+
+// Function to get subsegment information structure.
+Status GetEngineInfo(const Graph* g,
+                     const grappler::GraphProperties& graph_properties,
+                     const Segment& segment,
+                     const std::vector<Node*>& reverse_topo_order,
+                     EngineInfo* info) {
+  std::vector<const Node*> subgraph_nodes;  // Topologically sorted nodes.
+  std::set<const Node*> added_const_nodes;  // Used to prevent double insertion.
+
+  const ClusterProperty& segment_property = segment.property;
+  const std::set<const Node*, NodePtrCompare>& segment_nodes = segment.nodes;
+
+  // The device assignment accumulated from the compatible device assignments
+  // for the nodes in the segment.
+  const DeviceNameUtils::ParsedName segment_device =
+      segment_property.DeviceName();
+  info->max_batch_size = segment_property.BatchSize().GetOptionalMaxBatchSize();
+
+  // Map from src_node_name+port to the unique port numbers of the TRT op, where
+  // the src_node_name is the name of the source node of the input/output
+  // edge, thus there must not be any duplicates since source nodes of
+  // input/output edges must be in different split of the graph.
+  // TODO(aaroey): consider using node id and port instead.
+  // TODO(aaroey): using topo order instead of reverting reverse topo order.
+  std::unordered_map<string, int> input_to_engine_port, output_to_engine_port;
+  for (auto it = reverse_topo_order.rbegin(); it != reverse_topo_order.rend();
+       ++it) {
+    const Node* node = *it;
+    if (segment_nodes.count(node) == 0) continue;
+    subgraph_nodes.push_back(node);
+
+    const int node_id = node->id();
+    const string& node_name = node->name();
+
+    // Create input connections. Sort edges first to make deterministic since
+    // in_edges is a set of pointers.
+    std::vector<const Edge*> in_edges(node->in_edges().begin(),
+                                      node->in_edges().end());
+    std::sort(in_edges.begin(), in_edges.end(), EdgePtrCompare());
+    for (const auto edge : in_edges) {
+      auto input_node = edge->src();
+      if (input_node->IsSource() || segment_nodes.count(input_node)) {
+        continue;
+      }
+      if (edge->IsControlEdge()) {
+        if (ShallKeepControlEdgeFrom(input_node)) {
+          // Non-Const control input.
+          info->connections.emplace_back(input_node->name(), input_node->id(),
+                                         node_name, node_id,
+                                         /*input_edge=*/true);
+        }
+      } else if (input_node->type_string() == "Const") {
+        // Add constant data input nodes into the segment graphdef (thus also in
+        // the engine). We don't care if it has other output edges going into
+        // other engines or TF nodes. Since we add it only to the segment
+        // graphdef, not the segment itself, it won't be removed from the graph.
+        // If it doesn't have any edges, TF will prune it out.
+        //
+        // Note that the segmenter already ensure that the constant data input
+        // is valid and supported by the engine.
+        if (!added_const_nodes.insert(input_node).second) {
+          // Already added before.
+          continue;
+        }
+        VLOG(1) << "Adding const node " << input_node->name();
+      } else {
+        // Non-const data input.
+        int port = Graph::kControlSlot - 1;
+        // Use the source non-segment node name/port as key.
+        const string s = StrCat(input_node->name(), ":", edge->src_output());
+        VLOG(1) << "Input edge = " << s;
+        if (input_to_engine_port.count(s)) {
+          port = input_to_engine_port.at(s);
+        } else {
+          port = input_to_engine_port.size();
+          input_to_engine_port.insert({s, port});
+        }
+        info->connections.emplace_back(
+            input_node->name(), input_node->id(), edge->src_output(), node_name,
+            node_id, edge->dst_input(), /*input_edge=*/true, port);
+      }
+    }
+    // Create output connections. Sort edges first to make deterministic since
+    // out_edges is a set of pointers.
+    std::vector<const Edge*> out_edges(node->out_edges().begin(),
+                                       node->out_edges().end());
+    std::sort(out_edges.begin(), out_edges.end(), EdgePtrCompare());
+    for (const auto edge : out_edges) {
+      auto output_node = edge->dst();
+      if (output_node->IsSink() || segment_nodes.count(output_node)) {
+        continue;
+      }
+      if (edge->IsControlEdge()) {
+        // Control output.
+        if (ShallKeepControlEdgeFrom(node)) {
+          info->connections.emplace_back(output_node->name(), output_node->id(),
+                                         node_name, node_id,
+                                         /*input_edge=*/false);
+        }
+      } else {
+        // Data output.
+        int port = Graph::kControlSlot - 1;
+        // Use the source segment node name/port as key.
+        const string s = StrCat(node_name, ":", edge->src_output());
+        VLOG(1) << "Output edge = " << s;
+        if (output_to_engine_port.count(s)) {
+          port = output_to_engine_port.at(s);
+        } else {
+          port = output_to_engine_port.size();
+          output_to_engine_port.insert({s, port});
+        }
+        info->connections.emplace_back(
+            output_node->name(), output_node->id(), edge->dst_input(),
+            node_name, node_id, edge->src_output(), /*input_edge=*/false, port);
+      }
+    }
+  }  // For each segment node in topological order.
+
+  // Construct the const nodes first.
+  subgraph_nodes.insert(subgraph_nodes.begin(), added_const_nodes.begin(),
+                        added_const_nodes.end());
+  TF_RETURN_IF_ERROR(
+      ConvertSegmentToGraphDef(g, graph_properties, subgraph_nodes, info));
+  VLOG(1) << "Converted TensorRT candidate segment '" << info->engine_name
+          << "' to a GraphDef";
+  if (segment_device.has_type) {
+    // If the accumulated device assignment for the segment has a device type,
+    // the segmenter guarantees the device type is GPU. Use the device
+    // assignment in this case.
+    if (segment_device.type != "GPU") {
+      return errors::Internal(
+          "segment device is not GPU: ",
+          DeviceNameUtils::ParsedNameToString(segment_device));
+    }
+    info->device = DeviceNameUtils::ParsedNameToString(segment_device);
+  } else {
+    TfDeviceId tf_device_id;
+    PlatformDeviceId platform_device_id;
+    std::tie(tf_device_id, platform_device_id) = GetFirstValidDeviceId();
+    if (tf_device_id.value() >= 0) {
+      DeviceNameUtils::ParsedName parsed_name;
+      parsed_name.type = "GPU";
+      parsed_name.has_type = true;
+      parsed_name.id = tf_device_id.value();
+      parsed_name.has_id = true;
+      info->device = DeviceNameUtils::ParsedNameToString(parsed_name);
+    } else {
+      VLOG(1) << "No device is assigned to the segment. A device will be "
+                 "assigned during graph execution (inference).";
+    }
+  }
+  return OkStatus();
+}
+
+// Helper function to update edge connection from the removed node to the
+// engine node. If an outside node is gone, it must have been absorbed into
+// an engine node. Find the engine node.
+void UpdateToEngineNode(const std::vector<EngineInfo>& infos,
+                        const size_t my_engine_id,
+                        const std::vector<Node*>& engine_nodes,
+                        const bool is_input_edge, const string& node_name,
+                        Node** node, int* port) {
+  for (size_t t = 0; t < infos.size(); ++t) {
+    if (t == my_engine_id) {
+      continue;
+    }
+    const auto& info = infos.at(t);
+    for (const auto& eng_conn : info.connections) {
+      // If the connection being updated is an input connection, the source of
+      // the connection must be an output connection of another engine. And vise
+      // versa.
+      if (is_input_edge == eng_conn.is_input_edge) continue;
+      if (eng_conn.inside_node_name == node_name &&
+          eng_conn.inside_port == *port) {
+        *node = CHECK_NOTNULL(engine_nodes[t]);
+        QCHECK_EQ(info.engine_name, (**node).name())
+            << "Engine name mismatch: " << info.engine_name << " vs "
+            << (**node).name();
+        *port = eng_conn.port_number;
+        return;
+      }
+    }
+  }
+  LOG(FATAL) << "Node " << node_name << " not found in any engine.";
+}
+
+tensorflow::TensorShapeProto ComputeTRTNodeIOShape(
+    std::vector<PartialTensorShape>& partial_tensorshape_vect,
+    std::vector<tensorflow::TensorShapeProto>& shape_proto_vect,
+    const PartialTensorShape& conn_shape, int port_number) {
+  tensorflow::TensorShapeProto tmp_shape_proto;
+  conn_shape.AsProto(&tmp_shape_proto);
+
+  if (partial_tensorshape_vect.size() <= port_number) {
+    shape_proto_vect.resize(port_number + 1);
+    partial_tensorshape_vect.resize(port_number + 1);
+  }
+
+  return tmp_shape_proto;
+}
+
+// Function to insert a TRT engine node into the graph.
+// Create engine nodes in the following way:
+// 1. Each invocation of CreateTRTNode creates an engine node for infos[pos]
+// 2. When an engine node is created, add it into the graph with necessary
+//    re-wiring.
+//    2.1. If the outside connected node is existing, connect the engine
+//         node to it.
+//    2.2. If the outside connected node is gone, it must have been absorted
+//         into another engine node (which was processed before the processing
+//         one). Connect to the pre-existing engine node instead.
+// 3. In this way, we ensure the graph is topologically sort-able after each
+//    invocation of CreateTRTNode().
+Status CreateTRTNode(const TRTOptimizationPass::ConversionParams& params,
+                     const std::vector<EngineInfo>& infos, int pos,
+                     int default_max_batch_size, Graph* graph,
+                     std::vector<Node*>* engine_nodes,
+                     grappler::Cluster* cluster) {
+  const auto& info = infos.at(pos);
+  std::vector<tensorflow::TensorShapeProto> input_shape_protos;
+  std::vector<tensorflow::TensorShapeProto> output_shape_protos;
+  std::vector<PartialTensorShape> input_shapes;
+  std::vector<PartialTensorShape> output_shapes;
+  std::vector<NodeDefBuilder::NodeOut> inputs;
+  std::vector<Node*> input_nodes;
+  std::vector<Node*> control_input_nodes;
+  std::unordered_set<string> control_input_names;
+  std::vector<DataType> out_types;
+
+  VLOG(1) << "Processing " << info.engine_name;
+  // Collect needed info for creating the engine node in the graph
+  for (const auto& conn : info.connections) {
+    // Control edges
+    if (conn.is_control_edge()) {
+      // Skip control outputs for now. control output info are not needed for
+      // node creation and will be processed later.
+      if (!conn.is_input_edge) continue;
+
+      // Rewrire control input if it's not found in original graph.
+      Node* input_node = graph->FindNodeId(conn.outside_id);
+      int port = Graph::kControlSlot;
+      if (!input_node) {
+        UpdateToEngineNode(infos, pos, *engine_nodes, /*is_input_edge=*/true,
+                           conn.outside_node_name, &input_node, &port);
+        QCHECK_EQ(Graph::kControlSlot, port);
+      }
+      if (!control_input_names.insert(input_node->name()).second) {
+        continue;
+      }
+      control_input_nodes.push_back(input_node);
+      VLOG(1) << "Engine Control Input " << input_node->name() << " -> "
+              << info.engine_name;
+    } else {
+      // Data edges
+      if (!conn.is_input_edge) {
+        // Set the shapes and data types of the output edge.
+        tensorflow::TensorShapeProto out_shape = ComputeTRTNodeIOShape(
+            /*partial_tensorshape_vect=*/output_shapes,
+            /*shape_proto_vect=*/output_shape_protos,
+            /*conn_shape=*/conn.inside_shape,
+            /*port_number=*/conn.port_number);
+
+        output_shape_protos.at(conn.port_number) = out_shape;
+        output_shapes.at(conn.port_number) = conn.inside_shape;
+
+        if (out_types.size() <= conn.port_number) {
+          out_types.resize(conn.port_number + 1);
+        }
+        out_types.at(conn.port_number) = conn.connection_type;
+        VLOG(2) << "Collected output shape "
+                << output_shape_protos.at(conn.port_number).DebugString();
+      } else {
+        // Set the shapes of the input edge.
+        tensorflow::TensorShapeProto in_shape = ComputeTRTNodeIOShape(
+            /*partial_tensorshape_vect=*/input_shapes,
+            /*shape_proto_vect=*/input_shape_protos,
+            /*conn_shape=*/conn.outside_shape,
+            /*port_number=*/conn.port_number);
+
+        input_shape_protos.at(conn.port_number) = in_shape;
+        input_shapes.at(conn.port_number) = conn.outside_shape;
+
+        // Shape must be fully defined (excluding batch dimension) for static
+        // mode.
+        if (params.use_implicit_batch &&
+            info.engine_type == EngineInfo::EngineType::TRTStatic) {
+          for (int i = 1; i < conn.outside_shape.dims(); i++) {
+            if (conn.outside_shape.dim_size(i) <= 0) {
+              return errors::Internal(
+                  "Not fully defined input shape when in static mode which "
+                  "should have been excluded by the segmenter. ");
+            }
+          }
+        }
+
+        // Rewrire data input if it's not found in original graph.
+        Node* input_node = graph->FindNodeId(conn.outside_id);
+        int port = conn.outside_port;
+        if (!input_node) {
+          UpdateToEngineNode(infos, pos, *engine_nodes, /*is_input_edge=*/true,
+                             conn.outside_node_name, &input_node, &port);
+        }
+        if (std::find_if(
+                std::begin(inputs), std::end(inputs),
+                [input_node, &port](const NodeDefBuilder::NodeOut& inp) {
+                  return inp.node == input_node->name() && inp.index == port;
+                }) == std::end(inputs)) {
+          inputs.emplace_back(input_node->name(), port, conn.connection_type);
+          input_nodes.push_back(CHECK_NOTNULL(input_node));
+          VLOG(1) << "Engine Input " << input_node->name() << ":" << port
+                  << " -> " << info.engine_name << ":" << inputs.size() - 1;
+        }
+      }
+    }
+  }
+  // We don't support segments with no inputs. Fall back to native TF here to
+  // avoid crash later. Constant folding should've folded the ops that make up
+  // these segments.
+  if (inputs.empty()) {
+    return errors::Internal(
+        "Segment has no inputs (possible constfold failure)");
+  }
+
+  // Build the engine and get its serialized representation.
+  string segment_string;
+
+  int max_batch_size = info.max_batch_size.has_value()
+                           ? info.max_batch_size.value()
+                           : default_max_batch_size;
+
+  if (info.engine_type == EngineInfo::EngineType::TRTStatic) {
+    TF_RETURN_IF_ERROR(CreateStaticEngine(params, info, max_batch_size,
+                                          input_shapes, nullptr,
+                                          &segment_string, cluster));
+  }
+
+  string prec_string;
+  TF_RETURN_IF_ERROR(TrtPrecisionModeToName(info.precision_mode, &prec_string));
+  NodeDefBuilder node_builder(info.engine_name, "TRTEngineOp");
+  if (!info.device.empty()) node_builder.Device(info.device);
+  if (VLOG_IS_ON(1)) {
+    string ins = StrCat(info.engine_name, " inputs= ");
+    for (const auto& ii : inputs) {
+      StrAppend(&ins, ii.node, ":", ii.index, " ");
+    }
+    VLOG(1) << ins;
+  }
+  node_builder.Input(inputs);
+  for (const string& c : control_input_names) {
+    node_builder.ControlInput(c);
+  }
+
+  NodeDef trt_node;
+  NameAttrList function;
+  function.set_name(StrCat(info.engine_name, "_native_segment"));
+
+  node_builder.Attr("input_shapes", input_shape_protos)
+      .Attr("output_shapes", output_shape_protos)
+      .Attr("static_engine",
+            info.engine_type == EngineInfo::EngineType::TRTStatic)
+      .Attr("segment_func", function)
+      .Attr("serialized_segment", segment_string)
+      .Attr("calibration_data", "")
+      .Attr("max_cached_engines_count", info.maximum_cached_engines)
+      .Attr("workspace_size_bytes", info.max_workspace_size_bytes)
+      .Attr("max_batch_size", max_batch_size)
+      .Attr("precision_mode", prec_string)
+      .Attr("use_calibration", info.use_calibration)
+      .Attr("_use_implicit_batch", params.use_implicit_batch)
+      .Attr("use_explicit_precision", params.use_explicit_precision)
+      .Attr("_allow_build_at_runtime", info.allow_build_at_runtime)
+      .Attr("OutT", out_types);
+
+  if (!params.use_implicit_batch) {
+    node_builder.Attr("profile_strategy",
+                      ProfileStrategyToName(params.profile_strategy));
+  }
+
+  Status status = node_builder.Finalize(&trt_node);
+
+  if (!status.ok()) {
+    LOG(ERROR) << "Node construction failed with" << status;
+    return status;
+  }
+  VLOG(1) << "Adding TRTEngine " << info.engine_name << " to graph";
+
+  // Up until this point, graph is not modified. If we return !status.ok() from
+  // here, this segment will be skipped
+  // TODO(aaroey): let it return proper error status for the following logic
+  // instead of checking fail.
+  TF_ASSIGN_OR_RETURN(Node * engine_node, graph->AddNode(trt_node));
+  (*engine_nodes)[pos] = engine_node;
+  // Add control input and input edges to the engine node.
+  for (const auto in : control_input_nodes) {
+    VLOG(1) << "Connecting control edge from " << in->name() << " to "
+            << engine_node->name();
+    graph->AddControlEdge(in, engine_node);
+  }
+  VLOG(1) << "input_nodes size = " << input_nodes.size();
+  for (int i = 0; i < input_nodes.size(); ++i) {
+    Node* n = CHECK_NOTNULL(input_nodes[i]);
+    const auto& in = inputs[i];
+    VLOG(1) << "Connecting data edge from " << n->name() << ":" << in.index
+            << " to " << engine_node->name() << ":" << i;
+    graph->AddEdge(n, in.index, engine_node, i);
+  }
+
+  // Updates the inputs of output edges destination nodes, and point them to the
+  // engine node.
+  for (auto& conn : info.connections) {
+    if (conn.is_input_edge) {
+      continue;
+    }
+    Node* output_node = graph->FindNodeId(conn.outside_id);
+    int port = conn.outside_port;
+    if (!output_node) {
+      UpdateToEngineNode(infos, pos, *engine_nodes, /*is_input_edge=*/false,
+                         conn.outside_node_name, &output_node, &port);
+    }
+    if (conn.is_control_edge()) {
+      VLOG(1) << "Updating control edge from " << engine_node->name() << " to "
+              << output_node->name();
+      QCHECK_EQ(Graph::kControlSlot, port);
+      graph->AddControlEdge(engine_node, output_node);
+    } else {
+      VLOG(1) << "Updating data edge from " << engine_node->name() << ":"
+              << conn.port_number << " to " << output_node->name() << ":"
+              << port;
+      // Use UpdateEdge() to avoid adding the same edge multiple times.
+      TF_CHECK_OK(
+          graph->UpdateEdge(engine_node, conn.port_number, output_node, port));
+    }
+  }
+  return OkStatus();
+}
+
+int64 GetNextGraphSequenceNumber() {
+  static std::atomic<int64_t> graph_sequence_num;
+  return graph_sequence_num++;
+}
+
+constexpr char kCastInputTypeAttrName[] = "SrcT";
+
+// Transforms node = cast(x, fp32) where datatype(x) != fp16 to:
+//   castToFp16 = cast(x, fp16)
+//   node = cast(castToFp16, fp32)
+//
+Status MaybeRewriteCastToFp32(GraphDef* graph_def, NodeDef* node_def) {
+  if (node_def->op() != "Cast") {
+    return OkStatus();
+  }
+
+  DataTypeVector input_types;
+  DataTypeVector output_types;
+  TF_RETURN_IF_ERROR(
+      graph_transforms::GetInOutTypes(*node_def, &input_types, &output_types));
+
+  if (input_types.size() != 1 || output_types.size() != 1) {
+    return errors::Internal("Bad cast operation");
+  }
+
+  if (input_types[0] == DT_HALF || output_types[0] != DT_FLOAT) {
+    return OkStatus();
+  }
+
+  VLOG(2) << "Rewriting cast to FP32 " << node_def->DebugString();
+
+  NodeDef* castToFp16 = graph_def->add_node();
+  for (auto attr_value : node_def->attr()) {
+    (*castToFp16->mutable_attr())[attr_value.first] = attr_value.second;
+  }
+  castToFp16->set_name(node_def->name() + "_split");
+  castToFp16->set_op("Cast");
+  castToFp16->set_device(node_def->device());
+  castToFp16->add_input(node_def->input(0));
+  (*castToFp16->mutable_attr())[kCastOutputTypeAttrName].set_type(DT_HALF);
+
+  node_def->set_input(0, castToFp16->name() + ":0");
+  (*node_def->mutable_attr())[kCastInputTypeAttrName].set_type(DT_HALF);
+
+  VLOG(2) << castToFp16->DebugString();
+  VLOG(2) << node_def->DebugString();
+
+  return OkStatus();
+}
+
+}  // namespace
+
+Status RegisterGraphToFunctionLibrary(const GraphDef& segment_graph_def,
+                                      Graph* graph, const string& engine_name) {
+  Graph segment_graph(graph->flib_def());
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(GraphConstructorOptions(),
+                                            segment_graph_def, &segment_graph));
+  FunctionDefLibrary library;
+  auto segment_func = library.add_function();
+  TF_RETURN_IF_ERROR(GraphToFunctionDef(
+      segment_graph, StrCat(engine_name, "_native_segment"), segment_func));
+  if (VLOG_IS_ON(7)) {
+    VLOG(7) << engine_name << " Function_Def ";
+    VLOG(7) << segment_func->DebugString();
+  }
+  VLOG(1) << "Adding funcdef " << segment_func->signature().name()
+          << " to graphlib";
+  TF_RETURN_IF_ERROR(graph->AddFunctionLibrary(library));
+  return OkStatus();
+}
+
+std::pair<int, Allocator*> GetDeviceAndAllocator(
+    const grappler::Cluster* cluster, const EngineInfo& engine) {
+  int cuda_device_id = -1;
+  Allocator* dev_allocator = nullptr;
+  if (cluster == nullptr || cluster->GetDeviceSet() == nullptr ||
+      engine.device.empty()) {
+    // If device is not set, use the first found GPU device for the conversion.
+    TfDeviceId tf_device_id;
+    PlatformDeviceId platform_device_id;
+    std::tie(tf_device_id, platform_device_id) = GetFirstValidDeviceId();
+    cuda_device_id = platform_device_id.value();
+    if (cuda_device_id >= 0) {
+      GPUOptions gpu_options;
+      // If the TF to Cuda gpu id mapping exist, the device and corresponding
+      // allocator must have been initialized already, so the
+      // GetGPUAllocator() call won't create a new allocator.
+      dev_allocator = GPUProcessState::singleton()->GetGPUAllocator(
+          gpu_options, tf_device_id, /*total_bytes=*/1, /*peer_gpu_ids=*/{});
+    }
+    return std::make_pair(cuda_device_id, dev_allocator);
+  }
+
+  // Use the device requested by the engine.
+  auto device_set = cluster->GetDeviceSet();
+  std::vector<Device*> devices;
+  DeviceNameUtils::ParsedName parsed_name;
+  if (DeviceNameUtils::ParseFullName(engine.device, &parsed_name) &&
+      parsed_name.has_id) {
+    device_set->FindMatchingDevices(parsed_name, &devices);
+  }
+  if (!devices.empty()) {
+    if (devices.size() > 1) {
+      string msg = "Found multiple matching devices using name '";
+      StrAppend(&msg, engine.device, "': ");
+      for (auto d : devices) StrAppend(&msg, d->name(), ", ");
+      StrAppend(&msg, ". Will get the allocator from first one.");
+      LOG_WARNING_WITH_PREFIX << msg;
+    }
+    AllocatorAttributes alloc_attr;
+    cuda_device_id = devices[0]->tensorflow_accelerator_device_info()->gpu_id;
+    dev_allocator = devices[0]->GetAllocator(alloc_attr);
+    VLOG(1) << "Using allocator " << dev_allocator->Name()
+            << " and cuda_device_id " << cuda_device_id;
+  } else {
+    LOG_WARNING_WITH_PREFIX << "Cluster is set but device '" << engine.device
+                            << "' is not found in the cluster";
+  }
+  return std::make_pair(cuda_device_id, dev_allocator);
+}
+
+Status CreateStaticEngine(const TRTOptimizationPass::ConversionParams& params,
+                          const EngineInfo& info, int max_batch_size,
+                          const std::vector<PartialTensorShape>& input_shapes,
+                          TrtShapeOptimizationProfile* profile,
+                          string* segment_string, grappler::Cluster* cluster) {
+  std::pair<int, Allocator*> device_allocator =
+      GetDeviceAndAllocator(cluster, info);
+  int cuda_device_id = 0;
+  std::unique_ptr<TRTBaseAllocator> trt_allocator;
+  if (device_allocator.first >= 0) {
+    cuda_device_id = device_allocator.first;
+    trt_allocator.reset(new TRTDeviceAllocator(device_allocator.second));
+  } else {
+    // The value in trt_allocator is a nullptr and cudamalloc will be used.
+    LOG_WARNING_WITH_PREFIX << "Can't identify the cuda device. Running on "
+                               "device 0 and use cudamalloc as an allocator";
+  }
+  cudaSetDevice(cuda_device_id);
+
+  auto trt_logger = GetLoggerRegistry()->LookUp(params.trt_logger_name);
+  const bool calibrate_int8 =
+      (info.precision_mode == TrtPrecisionMode::INT8 && info.use_calibration);
+
+  // Create static engines with precision_mode fp32/fp16.
+  TrtUniquePtrType<nvinfer1::ICudaEngine> engine;
+  TF_RETURN_IF_ERROR(ConvertGraphDefToEngine(
+      info.segment_graph_def, nullptr,
+      calibrate_int8 ? TrtPrecisionMode::FP32 : info.precision_mode,
+      max_batch_size, info.max_workspace_size_bytes, input_shapes, trt_logger,
+      trt_allocator.get(), /*calibrator=*/nullptr, &engine,
+      info.use_calibration, params.use_implicit_batch,
+      /*convert_successfully=*/nullptr, profile, info.engine_name,
+      /*use_explicit_precision=*/params.use_explicit_precision, cluster));
+  TrtUniquePtrType<nvinfer1::IHostMemory> engine_data(engine->serialize());
+  *segment_string = string(static_cast<const char*>(engine_data->data()),
+                           engine_data->size());
+  return OkStatus();
+}
+
+Status ConvertGraph(const TRTOptimizationPass::ConversionParams& params,
+                    grappler::GrapplerItem& grappler_item,
+                    const std::vector<string>& input_output_names,
+                    grappler::Cluster* cluster, GraphDef* output) {
+  // Sanity checks.
+  TRT_ENSURE(output != nullptr)
+  if (params.precision_mode != TrtPrecisionMode::INT8 &&
+      params.use_calibration) {
+    return errors::InvalidArgument(
+        "Calibration with FP32 or FP16 is not supported.");
+  }
+
+  GraphDef& graph_def = grappler_item.graph;
+
+  // When precision_mode is FP16, transform cast(x, fp32) to
+  // cast(cast(x, fp16), fp32). This creates cast(fp16, f32) that can be
+  // included in the TRTEngineOp as an TensorRT Identity layer for performance:
+  //  . Avoid cast(fp32, fp16) in the TRT engine implementation for fp16
+  //    precision.
+  //  . Changing the input to the TRTEngine from fp32 to fp16 may reduce data
+  //    moving from the host to the GPU.
+  if (params.precision_mode == TrtPrecisionMode::FP16) {
+    for (int i = 0; i < graph_def.node_size(); i++) {
+      NodeDef* node_def = graph_def.mutable_node(i);
+      TF_RETURN_IF_ERROR(MaybeRewriteCastToFp32(&graph_def, node_def));
+    }
+  }
+
+  // Construct a GrapplerItem using the modified graph_def and the input
+  // grappler_item.
+  grappler::GraphProperties static_graph_properties(grappler_item);
+  TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true));
+
+  // Convert graphdef to graph.
+  FunctionLibraryDefinition flib(OpRegistry::Global(), graph_def.library());
+  Graph graph(flib);
+  TF_RETURN_IF_ERROR(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), graph_def, &graph));
+
+  // Segment the graph into subgraphs that can be converted to TensorRT
+  segment::SegmentOptions segment_options;
+  // TODO(ben,jie,sami): exclude output nodes (DISCUSS IT)
+  for (const auto& node : input_output_names) {
+    segment_options.exclude_node_list.insert(node);
+  }
+  segment_options.minimum_segment_size = params.minimum_segment_size;
+  segment_options.use_implicit_batch = params.use_implicit_batch;
+  if (segment_options.use_implicit_batch)
+    segment_options.maximum_batch_size = params.max_batch_size;
+  segment_options.allow_dynamic_non_batch_dim =
+      AllowDynamicNonBatchDimension(params);
+
+  segment::SegmentVector initial_segments;
+  TrtNodeValidator validator(static_graph_properties, params.precision_mode,
+                             params.use_calibration, params.use_implicit_batch,
+                             params.use_explicit_precision);
+  TF_RETURN_IF_ERROR(segment::SegmentGraph(
+      /*tf_graph=*/&graph,
+      /*graph_properties=*/&static_graph_properties,
+      /*candidate_fn=*/
+      std::bind(&TrtNodeValidator::IsTensorRTCandidate, &validator,
+                std::placeholders::_1),
+      // Input validation is already done by TrtNodeValidator, so we don't
+      // need to check the input edges.
+      /*input_candidate_fn=*/[](const Edge* edge) { return true; },
+      /*output_candidate_fn=*/OutputEdgeValidator(),
+      /*options=*/segment_options,
+      /*segments=*/&initial_segments));
+  LOG(INFO) << "Number of TensorRT candidate segments: "
+            << initial_segments.size();
+
+  // Get the EngineInfo for each segment.
+  std::unordered_map<string, Node*> node_map;
+  TF_RETURN_IF_ERROR(BuildNodeMap(graph, &node_map));
+  std::vector<EngineInfo> engine_segments;
+  engine_segments.reserve(initial_segments.size());
+  std::vector<Node*> reverse_topo_order;
+  GetPostOrder(graph, &reverse_topo_order);
+  segment::SegmentVector converted_segments;
+  converted_segments.reserve(initial_segments.size());
+  string engine_name_prefix =
+      StrCat("TRTEngineOp_",
+             absl::StrFormat("%0*d", 3, GetNextGraphSequenceNumber()), "_");
+  for (size_t t = 0; t < initial_segments.size(); t++) {
+    auto& curr_segment = initial_segments.at(t);
+    EngineInfo curr_engine;
+    curr_engine.engine_name =
+        StrCat(engine_name_prefix, absl::StrFormat("%0*d", 3, t));
+
+    bool int8_no_calib = (!params.use_calibration &&
+                          params.precision_mode == TrtPrecisionMode::INT8);
+    bool has_qdq = false;
+    if (int8_no_calib) {
+      has_qdq = absl::c_any_of(reverse_topo_order, IsQuantizeAndDequantizeOp);
+    }
+
+    Status status = GetEngineInfo(&graph, static_graph_properties, curr_segment,
+                                  reverse_topo_order, &curr_engine);
+    if (!status.ok()) {
+      LOG_WARNING_WITH_PREFIX << "Failed to get engine info for segment " << t
+                              << ": " << status;
+      continue;
+    }
+
+    curr_engine.engine_type = GetEngineType(params);
+    curr_engine.use_calibration = params.use_calibration;
+    // Building cuda engines for INT8 without calibration and without dynamic
+    // range info cause TRT failure. Avoid this situation by setting the
+    // precision to FP16.
+    if (int8_no_calib && !has_qdq) {
+      LOG(WARNING) << "Set engine precision to FP16 due to missing QDQ OP";
+      curr_engine.precision_mode = TrtPrecisionMode::FP16;
+    } else {
+      curr_engine.precision_mode = params.precision_mode;
+    }
+    curr_engine.maximum_cached_engines = params.max_cached_engines;
+    curr_engine.allow_build_at_runtime = params.allow_build_at_runtime;
+    if (!curr_engine.max_batch_size.has_value()) {
+      curr_engine.max_batch_size = params.max_batch_size;
+    }
+
+    status = RegisterGraphToFunctionLibrary(curr_engine.segment_graph_def,
+                                            &graph, curr_engine.engine_name);
+
+    if (!status.ok()) {
+      LOG_WARNING_WITH_PREFIX
+          << "Failed to register segment graphdef to the library " << t << ": "
+          << status;
+      continue;
+    }
+
+    engine_segments.push_back(std::move(curr_engine));
+    converted_segments.push_back(std::move(curr_segment));
+
+    if (VLOG_IS_ON(8)) {
+      string fname = engine_segments.back().engine_name;
+      StrAppend(&fname, ".pb");
+      std::fstream f;
+      f.open(fname.c_str(), std::fstream::out | std::fstream::binary);
+      f << engine_segments.at(t).segment_graph_def.SerializeAsString();
+      f.close();
+    }
+  }
+
+  // Save the cuda device since we may need to switch to another cuda device to
+  // build static engines.
+  std::optional<int> old_cuda_device = std::nullopt;
+  if (!params.is_dynamic_op) {
+    int cuda_device_id;
+    cudaError_t cuda_error = cudaGetDevice(&cuda_device_id);
+    if (cuda_error != cudaSuccess) {
+      LOG_WARNING_WITH_PREFIX << "Couldn't get current device: "
+                              << cudaGetErrorString(cuda_error);
+    } else {
+      VLOG(1) << "Current cuda device is " << cuda_device_id;
+      old_cuda_device = cuda_device_id;
+    }
+  }
+
+  auto restore_cuda_device = gtl::MakeCleanup([old_cuda_device] {
+    if (old_cuda_device.has_value()) {
+      cudaSetDevice(old_cuda_device.value());
+    }
+  });
+
+  std::vector<Node*> engine_nodes;
+  engine_nodes.resize(engine_segments.size());
+  for (int i = 0; i < engine_segments.size(); ++i) {
+    auto& engine = engine_segments.at(i);
+    // TODO(b/170762693): implement the heuristic to calculate
+    // max_workspace_size_bytes.
+    engine.max_workspace_size_bytes = params.max_workspace_size_bytes;
+    VLOG(1) << "Assigned " << engine.max_workspace_size_bytes << " bytes to "
+            << engine.engine_name;
+    auto status =
+        CreateTRTNode(params, engine_segments, i, params.max_batch_size, &graph,
+                      &engine_nodes, cluster);
+
+    string msg = StrCat("segment ", i, " consisting of ",
+                        converted_segments.at(i).nodes.size(), " nodes by ",
+                        engine.engine_name);
+    if (status.ok()) {
+      LOG(INFO) << "Replaced " << msg << ".";
+    } else {
+      // Graph is not modified.
+      LOG_WARNING_WITH_PREFIX << "Cannot replace " << msg
+                              << " reason: " << status.message()
+                              << " (keeping original segment).";
+    }
+    if (VLOG_IS_ON(1)) {
+      msg = "Segment consists of nodes: ";
+      for (const Node* node : converted_segments.at(i).nodes) {
+        StrAppend(&msg, node->name(), ", ");
+      }
+      VLOG(1) << msg;
+    }
+
+    // If status is ok, we successfully added the node to the graph and can
+    // remove segment ops. Otherwise graph is not modified.
+    if (status.ok()) {
+      for (const Node* node : converted_segments.at(i).nodes) {
+        graph.RemoveNode(const_cast<Node*>(node));
+      }
+    }
+  }
+  graph.ToGraphDef(output);
+  return OkStatus();
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/convert_graph.h

@@ -0,0 +1,70 @@
+/* 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_CONVERT_CONVERT_GRAPH_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_
+
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/types.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+// These functions are internal implementation functions for the
+// TRTOptimizationPass.
+
+// Performs segmentation and conversion on the given Grappler item. This method
+// contains the core logic of the TRTOptimizationPass.
+Status ConvertGraph(const TRTOptimizationPass::ConversionParams& params,
+                    grappler::GrapplerItem& grappler_item,
+                    const std::vector<string>& input_output_names,
+                    grappler::Cluster* cluster, GraphDef* output);
+
+// Helper method for the conversion, expose for testing.
+std::pair<int, Allocator*> GetDeviceAndAllocator(
+    const grappler::Cluster* cluster, const EngineInfo& engine);
+
+// Helper method that registers `segment_graph` as a function to the function
+// library in `graph`.
+Status RegisterGraphToFunctionLibrary(const GraphDef& segment_graph_def,
+                                      Graph* graph, const string& engine_name);
+
+// Creates and serializes an ICudaEngine. Used only in is_dynamic_op=false,
+// a.k.a. static engine mode.
+Status CreateStaticEngine(const TRTOptimizationPass::ConversionParams& params,
+                          const EngineInfo& info, int max_batch_size,
+                          const std::vector<PartialTensorShape>& input_shapes,
+                          TrtShapeOptimizationProfile* profile,
+                          string* segment_string, grappler::Cluster* cluster);
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_

tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc

@@ -0,0 +1,221 @@
+/* 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/convert/convert_graph.h"
+
+#include <regex>  // NOLINT
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_testutils.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/device_set.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
+#include "tensorflow/core/public/session.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class FakeCluster : public grappler::Cluster {
+ public:
+  FakeCluster() : Cluster(0) {}
+
+  void SetDeviceSet(const DeviceSet* device_set) { device_set_ = device_set; }
+
+  const DeviceSet* GetDeviceSet() const override { return device_set_; }
+
+  string type() const override { return ""; }
+  Status Provision() override { return OkStatus(); }
+  Status Initialize(const grappler::GrapplerItem& item) override {
+    return OkStatus();
+  }
+  Status Run(const GraphDef& graph_def,
+             const std::vector<std::pair<string, Tensor>>& feed,
+             const std::vector<string>& fetch, RunMetadata* metadata) override {
+    return OkStatus();
+  }
+
+ private:
+  const DeviceSet* device_set_ = nullptr;
+};
+
+TEST(GetDeviceAndAllocatorTest, GetDeviceAndAllocator) {
+  TRTOptimizationPass::ConversionParams params;
+  EngineInfo engine_info;
+  {
+    // cluster is not set, and no gpu device is available.
+    auto result = GetDeviceAndAllocator(nullptr, engine_info);
+    EXPECT_EQ(-1, result.first);
+    EXPECT_EQ(nullptr, result.second);
+  }
+
+  // Create a session with two (virtual) gpu device.
+  SessionOptions options;
+  ConfigProto* config = &options.config;
+  GPUOptions* gpu_options = config->mutable_gpu_options();
+  auto virtual_devices =
+      gpu_options->mutable_experimental()->add_virtual_devices();
+  virtual_devices->add_memory_limit_mb(200);
+  virtual_devices->add_memory_limit_mb(200);
+  std::unique_ptr<Session> session(NewSession(options));
+
+  {
+    // cluster is not set, should find and return first gpu id and
+    // corresponding allocator.
+    auto result = GetDeviceAndAllocator(nullptr, engine_info);
+    EXPECT_EQ(0, result.first);
+    EXPECT_NE(nullptr, result.second);
+    EXPECT_EQ("GPU_0_bfc", result.second->Name());
+  }
+
+  FakeCluster cluster;
+  {
+    // params.cluster->GetDeviceSet() returns null, should find and return first
+    // gpu id and corresponding allocator.
+    auto result = GetDeviceAndAllocator(&cluster, engine_info);
+    EXPECT_EQ(0, result.first);
+    EXPECT_NE(nullptr, result.second);
+    EXPECT_EQ("GPU_0_bfc", result.second->Name());
+  }
+
+  // Build the DeviceSet.
+  DeviceSet device_set;
+  const DeviceMgr* device_mgr = nullptr;
+  TF_ASSERT_OK(session->LocalDeviceManager(&device_mgr));
+  for (auto d : device_mgr->ListDevices()) {
+    device_set.AddDevice(d);
+  }
+  cluster.SetDeviceSet(&device_set);
+  {
+    // engine_info.device is not set, should find and return first gpu id and
+    // corresponding allocator.
+    auto result = GetDeviceAndAllocator(&cluster, engine_info);
+    EXPECT_EQ(0, result.first);
+    EXPECT_NE(nullptr, result.second);
+    EXPECT_EQ("GPU_0_bfc", result.second->Name());
+  }
+
+  engine_info.device = "/GPU:1";
+  {
+    // Set to use second device.
+    auto result = GetDeviceAndAllocator(&cluster, engine_info);
+    EXPECT_EQ(0, result.first);
+    EXPECT_NE(nullptr, result.second);
+    EXPECT_EQ("GPU_1_bfc", result.second->Name());
+  }
+
+  engine_info.device = "/GPU:3";
+  {
+    // Set to use nonexistent device.
+    auto result = GetDeviceAndAllocator(&cluster, engine_info);
+    EXPECT_EQ(-1, result.first);
+    EXPECT_EQ(nullptr, result.second);
+  }
+}
+
+class ConvertGraphTest : public ::testing::Test {
+ public:
+  Status RunConvertGraph(Scope s, GraphDef* output_graph_def,
+                         int maximum_batch_size = 1000) {
+    // Create GraphProperties.
+    grappler::GrapplerItem item;
+    TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+    grappler::GraphProperties graph_properties(item);
+    TF_EXPECT_OK(graph_properties.InferStatically(true));
+
+    // Construct ConversionParams.
+    const std::vector<string> input_output_names{"output"};
+    TRTOptimizationPass::ConversionParams params;
+    params.max_batch_size = maximum_batch_size;
+    params.max_workspace_size_bytes = 8 << 20;
+    params.minimum_segment_size = 1;
+    params.use_calibration = false;
+    params.trt_logger_name = "DefaultLogger";
+    return ConvertGraph(params, item, input_output_names, nullptr,
+                        output_graph_def);
+  }
+};
+
+TEST_F(ConvertGraphTest, DirectlyConnectedEngines) {
+  // Create the graph. There will be two TRTEngineOps after the conversion, and
+  // the upstream TRTEngineOp will have two output connections from the same
+  // node:port inside the op to the downstream TRTEngineOp. Then, if it adds the
+  // downstream TRTEngineOp first, when adding the upstream op it'll need to
+  // update the same output connection twice. This test ensures the correctness
+  // of the conversion under such condition.
+  Scope s = Scope::NewRootScope();
+  auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT,
+                                ops::Placeholder::Shape({2, 1}));
+  // We purposefully choose the name of the root node of each segment, so it'll
+  // process the segment in the downstream first, then, when it tries to update
+  // the edge between the two TRTEngineOps, it'll try to add the same edge
+  // multiple times.
+  auto segment_root_1 = ops::Identity(s.WithOpName("segment_root_b"), input);
+  auto add1 = ops::Add(s.WithOpName("add1"), segment_root_1, segment_root_1);
+  // Add incompatible reshapes that change the batch dimension.
+  auto incompatible =
+      ops::Reshape(s.WithOpName("reshape1"), add1, Input({1, 2}));
+  incompatible =
+      ops::Reshape(s.WithOpName("reshape2"), incompatible, Input({2, 1}));
+
+  auto add2 = ops::Add(s.WithOpName("add2"), incompatible, add1);
+  auto segment_root_2 = ops::Identity(s.WithOpName("segment_root_a"), add1);
+  auto add3 = ops::Add(s.WithOpName("add3"), add2, segment_root_2);
+  ops::Identity(s.WithOpName("output"), add3);
+
+  GraphDef output_graph_def;
+  TF_EXPECT_OK(RunConvertGraph(s, &output_graph_def));
+
+  auto remove_graph_sequence_number = [](std::string node_name) {
+    const std::regex pattern("TRTEngineOp_[0-9]+_");
+    return std::regex_replace(node_name, pattern, "TRTEngineOp_");
+  };
+  int num_trt_ops = 0;
+  for (const NodeDef& node : output_graph_def.node()) {
+    std::string node_name = node.name();
+    if (node.op() != "TRTEngineOp") continue;
+    node_name = remove_graph_sequence_number(node_name);
+    if (node_name == "TRTEngineOp_001") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("input", node.input(0));
+      ++num_trt_ops;
+    } else if (node_name == "TRTEngineOp_000") {
+      EXPECT_EQ(2, node.input_size());
+      EXPECT_EQ("TRTEngineOp_001", remove_graph_sequence_number(node.input(0)));
+      EXPECT_EQ("reshape2", node.input(1));
+      ++num_trt_ops;
+    }
+  }
+  EXPECT_EQ(2, num_trt_ops);
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h

@@ -0,0 +1,593 @@
+/* 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_CONVERT_CONVERT_NODES_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_
+
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/weights.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.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/compiler/tf2tensorrt/utils/trt_tensor_proxy.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/grappler/costs/graph_properties.h"
+#include "tensorflow/core/lib/core/status.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+namespace convert {
+using ::tsl::StatusOr;
+
+struct EngineConnection {
+  // Constructs a non-control edge.
+  EngineConnection(const string& outside, int out_id, int out_port,
+                   const string& inside, int in_id, int in_port,
+                   bool input_edge, int port)
+      : outside_node_name(outside),
+        outside_id(out_id),
+        outside_port(out_port),
+        inside_node_name(inside),
+        inside_id(in_id),
+        inside_port(in_port),
+        is_input_edge(input_edge),
+        port_number(port) {}
+
+  // Constructs a control edge.
+  EngineConnection(const string& outside, int out_id, const string& inside,
+                   int in_id, bool input_edge)
+      : outside_node_name(outside),
+        outside_id(out_id),
+        outside_port(Graph::kControlSlot),
+        inside_node_name(inside),
+        inside_id(in_id),
+        inside_port(Graph::kControlSlot),
+        is_input_edge(input_edge),
+        port_number(Graph::kControlSlot) {}
+
+  bool is_control_edge() const { return port_number == Graph::kControlSlot; }
+
+  const string outside_node_name;
+  const int outside_id;
+  const int outside_port;
+  PartialTensorShape outside_shape;  // Only set for input edge.
+
+  const string inside_node_name;
+  const int inside_id;
+  const int inside_port;
+  PartialTensorShape inside_shape;  // Only set for output edge.
+
+  DataType connection_type;
+  const bool is_input_edge;
+
+  // The port number of the TRT node connected with this edge.
+  const int port_number;
+};
+
+struct EngineInfo {
+  EngineInfo()
+      : engine_type(EngineType::TRTStatic),
+        max_workspace_size_bytes(0),
+        max_batch_size(std::nullopt),
+        maximum_cached_engines(0),
+        precision_mode(TrtPrecisionMode::FP32),
+        use_calibration(true),
+
+        allow_build_at_runtime(true),
+        use_explicit_precision(false) {}
+
+  string engine_name;
+  string device;
+  GraphDef segment_graph_def;
+
+  // Non-control input connections inside this vector are sorted in a way such
+  // that, the segment nodes connecting to them are topological sorted.
+  // In addition, for non-control connections, there must be no duplicates.
+  std::vector<EngineConnection> connections;
+
+  enum class EngineType { TRTStatic = 0, TRTDynamic = 1 };
+  EngineType engine_type;
+  int64 max_workspace_size_bytes;
+  std::optional<int> max_batch_size;
+  int maximum_cached_engines;
+  TrtPrecisionMode precision_mode;
+  bool use_calibration;
+  bool allow_build_at_runtime;
+  bool use_explicit_precision;
+};
+
+// Constructs a graphdef from the segment in the given graph and stores it to
+// the engine_info. Adds _Arg nodes for input edges (InputPH_*) and _Retval
+// nodes for output edges (OutputPH_*). Maintains the topological order of the
+// non-input/output nodes in the graphdef. This function needs to be called
+// before TensorRT layers are created because it prepares the original graph
+// for TensorRT conversion.
+//
+// - subgraph_node_names: the node names of the subgraph.
+// - subgraph_node_ids: the node ids of the subgraph, must be sorted in
+//   topological order.
+// - engine_info: a data structure that records the information about the
+//   engine containing the subgraph.
+//
+// TODO(aaroey): add tests to validate these properties.
+Status ConvertSegmentToGraphDef(
+    const Graph* graph, const grappler::GraphProperties& graph_properties,
+    const std::vector<const Node*>& subgraph_nodes, EngineInfo* engine_info);
+
+// Converts given subgraph to a TRT engine saved in 'engine'. Returns ok iff
+// 'builder' successfully build the engine. If the result is not ok, 'engine'
+// will be set to nullptr
+// Once returned, 'builder' is not needed any more and can be safely destroyed.
+//
+// - convert_successfully: indicates whether the conversion to TensorRT network
+//   is successful. This is different than successfully building the engine:
+//   building can still fail afterwards.
+// Note: When 'cluster' is not null, it contains the graph to be converted.
+//       We may perform additional optimizations to the graph before converting
+//       the graph.
+Status ConvertGraphDefToEngine(
+    const GraphDef& gdef, OpKernelContext* ctx, TrtPrecisionMode precision_mode,
+    int max_batch_size, size_t max_workspace_size_bytes,
+    const std::vector<PartialTensorShape>& input_shapes,
+    nvinfer1::ILogger* logger, nvinfer1::IGpuAllocator* allocator,
+    TRTInt8Calibrator* calibrator,
+    TrtUniquePtrType<nvinfer1::ICudaEngine>* engine, bool use_calibration,
+    const bool use_implicit_batch, bool* convert_successfully,
+    TrtShapeOptimizationProfile* profiles, absl::string_view engine_name,
+    bool use_explicit_precision,
+    tensorflow::grappler::Cluster* cluster = nullptr,
+    const string& device = "");
+
+// Helper class for the segmenter to determine whether an output edge from the
+// TRT segment is valid.
+class OutputEdgeValidator {
+ public:
+  // Return true if the specified edge is eligible to be an output edge of the
+  // TRT segment.
+  bool operator()(const Edge* out_edge) const;
+};
+
+// Class to verify if specific TF node is supported by TRT.
+class TrtNodeValidator {
+ public:
+  // 'graph_properties' is the GraphProperties of the graph whose nodes will be
+  // checked by IsTensorRTCandidate() later. It is used to get the shape and
+  // data type information of a tensor for validation purpose.
+  TrtNodeValidator(const grappler::GraphProperties& graph_properties,
+                   TrtPrecisionMode precision_mode, bool use_calibration,
+                   bool use_implicit_batch, bool use_explicit_precision);
+
+  // Returns OK iff 'node' is a TF-TRT conversion candidate, which will be added
+  // to TRT subgraph and later converted into TRT engine.
+  Status IsTensorRTCandidate(const Node* node);
+
+  static const std::set<string>* quantize_ops;
+
+  // Returns validator by op type. If no validator is registered for
+  // specific op, it means no validation is needed and ValidateNode() will
+  // return OK.
+  StatusOr<OpConverter> GetValidator(const std::string& op);
+
+ private:
+  // Convert a Const node to a TRT_TensorOrWeights.
+  Status ConvertConstToWeights(const NodeDef& const_node_def,
+                               const std::vector<TRT_TensorOrWeights>& inputs,
+                               TRT_TensorOrWeights* output);
+
+  // Convert a VariableV2 node to a TRT_TensorOrWeights.
+  Status ConvertVariableToWeights(
+      const NodeDef& const_node_def,
+      const std::vector<TRT_TensorOrWeights>& inputs,
+      TRT_TensorOrWeights* output);
+
+  // Convert the output tensor at 'output_port' of 'node_def' to a
+  // TRT_TensorOrWeights which will be later used as an input to other nodes and
+  // passed to ValidateNode() below.
+  Status ConvertToTensorOrWeights(const NodeDef& node_def, int output_port,
+                                  TRT_TensorOrWeights* tensor_or_weights);
+
+  // Store the weights added during validation. Some validations (e.g.
+  // validation for Const node) may produce weights.
+  TrtWeightStore weight_store_;
+
+  // GraphProperties of the graph whose nodes are to be validated by
+  // IsTensorRTCandidate().
+  const grappler::GraphProperties& graph_properties_;
+
+  // Quantization ops are only converted when using quantized precisions.
+  const TrtPrecisionMode precision_mode_;
+
+  const bool use_calibration_;
+
+  const bool use_implicit_batch_;
+
+  const bool use_explicit_precision_;
+
+  friend class ValidatorTest;
+  friend class OpConverterTest;
+};
+
+// Class to convert TF nodes to TRT network.
+class Converter {
+ public:
+  // Used for Converter::RenameAndMarkOutputTensors()
+  struct EngineOutputInfo {
+    // The TRT tensor name which produces the output.
+    string source_tensor_name;
+    // The TensorFlow node name which is receiving the output from the TRT
+    // engine. This should always be the Identity node created in
+    // ConvertSegmentToGraphDef.
+    string dest_node_name;
+    // Output type. TensorRT requires this to be explicitly set for engine
+    // outputs.
+    nvinfer1::DataType trt_dtype;
+  };
+
+  static StatusOr<std::unique_ptr<Converter>> Create(
+      TrtPrecisionMode precision_mode, bool use_calibration,
+      nvinfer1::ILogger* trt_logger, const bool use_implicit_batch,
+      absl::string_view engine_name, bool use_explicit_precision = false,
+      OpKernelContext* ctx = nullptr);
+
+  //////////////////////////////////////////////////////////////////////////////
+  // Methods used by the TRT engine builder to build a TRT network from a TF
+  // function/subgraph.
+
+  // Convert the node to TRT network.
+  Status ConvertNode(const NodeDef& node_def);
+
+  // Add input tensor to the TRT network with given 'name', 'dtype', 'dims' and
+  // 'batch_size'.
+  Status AddInputTensor(const string& name, nvinfer1::DataType dtype,
+                        const nvinfer1::Dims& dims, int batch_size);
+
+  // Store the ResourceHandle as a TRT_TensorOrWeights object. This can be
+  // later used as input to other nodes.
+  Status AddInputResource(const string& name, const ResourceHandle& resource);
+
+  // Mark the tensors with names specified by source_tensor_name as output of
+  // the TRT network, and set their names in the TRT network as dest_node_name.
+  Status RenameAndMarkOutputTensors(
+      const std::vector<EngineOutputInfo>& output_tensors);
+
+  // Build a TRT engine using the created network.
+  Status BuildCudaEngine(TrtUniquePtrType<nvinfer1::ICudaEngine>* engine,
+                         int max_batch_size, size_t max_workspace_size_bytes,
+                         nvinfer1::IGpuAllocator* allocator,
+                         TRTInt8Calibrator* calibrator,
+                         TrtShapeOptimizationProfile* profiles);
+
+  //////////////////////////////////////////////////////////////////////////////
+  // Methods used by op converters to convert individual TF node and add layers
+  // to the TRT network.
+
+  // Op converters (e.g. ConvertReshape) need to access the TRT network in order
+  // to add TRT layers.
+  nvinfer1::INetworkDefinition* network() { return trt_network_.get(); }
+
+  // What precision are we targeting?
+  TrtPrecisionMode precision_mode() const { return precision_mode_; }
+
+  // Variable converters need the context to read variable values.
+  OpKernelContext* context() { return ctx_; }
+
+  // Calibration will be or was previously performed on this network?
+  bool use_calibration() const { return use_calibration_; }
+
+  // Whether implicit batch mode is enabled
+  bool use_implicit_batch() const { return use_implicit_batch_; }
+
+  // This function should be called when we know the quantization range of a
+  // tensor from a quantize/dequantize node.
+  void ProvideQuantizationRange(ITensorProxyPtr* tensor, float min_range,
+                                float max_range);
+
+  // Should be called when full TRT network has been constructed and before
+  // building the engine.
+  void MaybeApplyQuantizationRanges();
+
+  // Below are helper methods for op converters to add different layers to the
+  // TRT network.
+
+  // Transpose 'input_tensor' with given permutation 'order_with_batch_dim' to
+  // 'output_tensor'. The permutation 'order_with_batch_dim' contains the batch
+  // dimension which should always be 0. If this is for adding a transpose layer
+  // to support the conversion of 'node_def', callers need to provide a
+  // non-empty 'sub_op_name' appended to the name of 'node_def' to avoid layer
+  // name conflicts.
+  Status TransposeTensor(ITensorProxyPtr input_tensor,
+                         const std::vector<int>& order_with_batch_dim,
+                         ITensorProxyPtr* output_tensor,
+                         const NodeDef& node_def,
+                         absl::string_view sub_op_name = "");
+
+  // Reshapes a dynamic shape tensor by removing or adding dimensions of size 1,
+  // and/or permuting the dimensions. The new shape is derived from the shape of
+  // the input tensor according to the slices and size_for_added_dims arguments.
+  //
+  // If there would be at most one unknown dimension, we could set the new shape
+  // using IShuffleLayer::setReshapeDimensions, which treats -1 as a special
+  // value (the same way as TF). In general, we can have more than one unknown
+  // dimensions, and we have to manipulate the shape tensors during runtime to
+  // define the new shape. This helper function defines the necessary shape
+  // inference layers and calls reshape using the calculated new shape.
+  //
+  // Example:
+  //
+  // Assume that we want to reshape a tensor from shape {A,B,C,D} to {C,D,A,B}
+  // (no transpose, just change the shape). In dynamic shape mode, the A,B,C,D
+  // values are not necessarily known at conversion time, they can be all -1. We
+  // can only define the new shape at runtime, when the actual shape is already
+  // known. To define the new shape:
+  // - We use an IShapeLayer to retrieve a shape tensor with the {A,B,C,D}
+  //   values.
+  // - Create two slices {C,D} and {A,B} of the shape tensor.
+  // - Concatenate these slices {C,D,A,B},
+  // - Set the {C,D,A,B} shape tensor as an input shape tensor for
+  // IShuffleLayer.
+  //
+  // This can be achieved by calling DynamicReshape(input, {{2,4},{0,2}},
+  // params).
+  //
+  // Before each slice we can insert new dims if the corresponding
+  // size_for_added_dims element is not negative. The size_for_added_dims array
+  // can have more than slices.size() elements, in order to insert a dimension
+  // after the last slice. For example, to add two leading 1 dimensions, and
+  // three trailing 1 dimensions, call DynamicReshape(input, {{0,nbDims}},
+  // {2, 3}).
+  //
+  // Parameters:
+  // input - input tensor
+  // slices - [start, end) pairs of slices
+  // params - conversion parameters
+  // output - reshaped tensor
+  // size_for_added_dims - size of dimension inserted right before slice[i]. We
+  //   only insert a new dim if size_for_added_dims[i] >= 0.
+  Status DynamicReshape(ITensorProxyPtr input,
+                        std::vector<std::pair<int, int>> slices,
+                        const OpConverterParams* params,
+                        ITensorProxyPtr* output,
+                        std::vector<int> size_for_added_dims = {},
+                        std::optional<int> op_instance = std::nullopt);
+
+  // Inserts a singleton dimension at axis for a dynamic shape tensor.
+  Status DynamicExpandDims(ITensorProxyPtr input, const nvinfer1::Dims& dims,
+                           int axis, const OpConverterParams* params,
+                           ITensorProxyPtr* output,
+                           std::optional<int> op_instance = std::nullopt);
+
+  // Helper function to add a squeeze op to the network.
+  //
+  // The input_dims argument stores the TRT dimensions of the input tensor,
+  // where the dimensions to be squeezed are replaced by 0.
+  Status SqueezeTensor(ITensorProxyPtr input, std::vector<int>* input_dims,
+                       const OpConverterParams* params, ITensorProxyPtr* output,
+                       std::optional<int> op_instance = std::nullopt);
+
+  // Creates an IConstantLayer using 'weights' whose dimensions are specified by
+  // 'dims', and returns the output ITensor.
+  ITensorProxyPtr CreateConstantLayer(const TRT_ShapedWeights& weights,
+                                      const nvinfer1::Dims& dims);
+
+  // Gets the min and max value in a TRT_ShapedWeights
+  Status GetWeightRange(const TRT_ShapedWeights& weights, float* out_min,
+                        float* out_max) const;
+
+  // Constructs a name and passed it to the TensorRT layer to support xprof.
+  void SetLayerName(nvinfer1::ILayer* layer, const NodeDef& node_def,
+                    absl::string_view sub_op_name = "",
+                    std::optional<int> sub_op_instance = std::nullopt,
+                    std::optional<std::string> origin_node_name = std::nullopt);
+
+  void SetLayerName(nvinfer1::ILayer* layer, absl::string_view main_op_name,
+                    absl::string_view sub_op_name,
+                    std::optional<int> sub_op_instance = std::nullopt);
+
+  std::unordered_map<string, TRT_TensorOrWeights>& TensorsMap() {
+    return trt_tensors_;
+  }
+
+  bool UseExplicitPrecision() const { return use_explicit_precision_; }
+
+ private:
+  Converter(TrtPrecisionMode precision_mode, bool use_calibration,
+            nvinfer1::ILogger* trt_logger, const bool use_implicit_batch,
+            absl::string_view engine_name, bool use_explicit_precision,
+            OpKernelContext* ctx);
+
+  Status Init(nvinfer1::ILogger* trt_logger);
+
+  // Verify the provided batch_size is consistent with batch_size_ and update it
+  // if necessary.
+  Status MaybeUpdateBatchSize(int batch_size);
+
+  // Add the provided tensor/weights to the map trt_tensors_.
+  Status AddTensorOrWeights(const string& name, TRT_TensorOrWeights input);
+
+  // Get the tensor/weights from trt_tensors_ by 'name'.
+  Status GetTensorOrWeights(const string& name, TRT_TensorOrWeights* output);
+
+  // Get the inputs of 'node_def' from trt_tensors_.
+  Status GetInputs(const NodeDef& node_def,
+                   std::vector<TRT_TensorOrWeights>* inputs) const;
+
+  // Tensors/weights added during construction of trt_network_.
+  std::unordered_map<string, TRT_TensorOrWeights> trt_tensors_;
+
+  // The TRT builder used to create the network and build the engine. Not owned.
+  TrtUniquePtrType<nvinfer1::IBuilder> trt_builder_;
+
+  // The TRT network being built.
+  TrtUniquePtrType<nvinfer1::INetworkDefinition> trt_network_;
+
+  // Store the weights added during construction of trt_network_.
+  TrtWeightStore weight_store_;
+
+  // Store the context.
+  OpKernelContext* ctx_;
+
+  // During conversion, this table is populated with quantization ranges per
+  // tensor. MaybeApplyQuantizationRanges() will use this table to set the TRT
+  // quantization ranges. Since TRT only supports symmetric ranges, we will
+  // store the range as a single float = max(abs(min_range), abs(max_range)).
+  // Range refers to the floating point values, e.g. min_range = 0.0f, max_range
+  // = 6.0f for Relu6.
+  std::unordered_map<ITensorProxyPtr*, float> quantization_ranges_proxy_;
+  std::unordered_map<nvinfer1::ITensor*, float> quantization_ranges_;
+
+  const TrtPrecisionMode precision_mode_;
+
+  const bool use_calibration_;
+
+  // If this is false, all dimensions including the batch dimension are
+  // set explicitly.
+  const bool use_implicit_batch_;
+
+  // Batch size of inputs to trt_network_ added by AddInputTensor(). During
+  // network construction it will update this, use it to verify the batch
+  // size of all inputs are compatible, and make sure individual TF node is
+  // acceptable by TRT.
+  int batch_size_ = -1;
+
+  // Assign a ID to each constant layer we create, so that we can assign a
+  // unique name to the layer.
+  int next_constant_layer_id_ = 0;
+
+  // The name of the TRTEngineOp node.
+  absl::string_view engine_name_;
+
+  // Indicates whether to use explicit precision in TensorRT (Q/DQ support).
+  bool use_explicit_precision_;
+
+  friend class ConverterTest;
+  friend class OpConverterTest;
+};
+
+// Converts a TensorFlow tensor to TRT shaped weights.
+Status TfTensorToTrtWeights(const Tensor& tensor, TrtWeightStore* weight_store,
+                            TRT_ShapedWeights* weights);
+
+// Converts 'input' of 'node_def' into 'tensor' with shape specified by 'dims'
+// (which doesn't contain the batch dimension).
+//
+// If validation_only is true, it doesn't do the conversion but only do some
+// minimum validation for the eligibility of the conversion, and *tensor will
+// be set to nullptr.
+// If validation_only is false converter must not be nullptr.
+Status PrepareTensorForShape(
+    Converter* converter, const TRT_TensorOrWeights& input,
+    const DimsAdapter& dims, const bool validation_only,
+    ITensorProxyPtr* tensor, const NodeDef& node_def,
+    std::optional<int> op_instance = std::nullopt,
+    std::optional<std::string> origin_node_name = std::nullopt);
+
+// Return OK if the broadcast scheme is supported and compute the shapes after
+// broadcasting. check_feasibility can be set to false in cases where dimensions
+// do not need to match exactly (as in the case of BatchMatMulV2).
+Status GetTrtBroadcastShape(const TRT_TensorOrWeights& operand_l,
+                            const TRT_TensorOrWeights& operand_r,
+                            const bool check_feasibility,
+                            const bool use_implicit_batch,
+                            nvinfer1::Dims* operand_l_new_dims,
+                            nvinfer1::Dims* operand_r_new_dims);
+
+template <typename T>
+using OperationMap = std::unordered_map<std::string, T>;
+
+// Map from Tensorflow operation names to TensorRT unary operations.
+using UnaryOperationMapType = OperationMap<nvinfer1::UnaryOperation>;
+const UnaryOperationMapType* UnaryOperationMap();
+
+// Map from Tensorflow boolean operation names to TensorRT unary operations.
+const UnaryOperationMapType* UnaryBooleanOperationMap();
+
+// Map of all supported ActivationTypes.
+using ActivationTypeMapType = OperationMap<nvinfer1::ActivationType>;
+const ActivationTypeMapType* ActivationTypeMap();
+
+// Map from Tensorflow binary operation names to TensorRT binary operations
+// types.
+using BinaryOperationMapType = OperationMap<nvinfer1::ElementWiseOperation>;
+const BinaryOperationMapType* BinaryOperationMap();
+
+// Map from Tensorflow boolean binary operation names to TensorRT binary
+// operations types.
+const BinaryOperationMapType* BinaryBooleanOperationMap();
+
+template <typename T>
+absl::InlinedVector<std::string, 10> GetOperationNames(const T& set) {
+  absl::InlinedVector<std::string, 10> result;
+  absl::c_transform(set, std::back_inserter(result),
+                    [](const auto x) { return x.first; });
+  return result;
+}
+
+// Adds a matrix multiplication operation to the TensorRT graph. The "params"
+// pointer is only used to access the TRT network builder. The inputs and
+// parameters for the op are fully specified by input_[a|b] and transpose_[a|b].
+StatusOr<ITensorProxyPtr> ConvertMatMulImpl(const OpConverterParams* params,
+                                            TRT_TensorOrWeights input_a,
+                                            TRT_TensorOrWeights input_b,
+                                            bool transpose_a, bool transpose_b);
+
+Status ApplyBroadcast(std::unique_ptr<TRT_TensorOrWeights>& operand,
+                      const DimsAdapter& broadcasted_dims,
+                      const OpConverterParams* params,
+                      std::optional<int> op_instance);
+
+std::string convert_range_error_msg(float start, float limit, float delta);
+std::string convert_range_expected_msg(const NodeDef& node_def);
+std::string bool_weight_error_msg(const NodeDef& node_def);
+std::string unexpected_type_error_msg(nvinfer1::DataType type_being_checked,
+                                      nvinfer1::DataType type_expected,
+                                      const NodeDef& node_def, int idx = 0);
+std::string then_else_dtypes_error_msg(nvinfer1::DataType type_then,
+                                       nvinfer1::DataType type_else,
+                                       const NodeDef& node_def);
+std::string input_shapes_error_msg(const nvinfer1::Dims& shape1,
+                                   const nvinfer1::Dims& shape2,
+                                   const NodeDef& node,
+                                   bool then_vs_else = false);
+std::string batch_size_error(absl::string_view name, absl::string_view comment);
+
+inline bool find_name(const string& name, const std::vector<string> names) {
+  return std::find(names.begin(), names.end(), name) != names.end();
+}
+
+Status check_type(nvinfer1::DataType type_being_checked,
+                  nvinfer1::DataType type_expected, const NodeDef& node_def,
+                  int idx = 0);
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_

tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc

@@ -0,0 +1,60 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
+
+#include <unordered_map>
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/mutex.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+class LoggerRegistryImpl : public LoggerRegistry {
+  Status Register(const string& name, nvinfer1::ILogger* logger) override {
+    mutex_lock lock(mu_);
+    if (!registry_.emplace(name, std::unique_ptr<nvinfer1::ILogger>(logger))
+             .second) {
+      return errors::AlreadyExists("Logger ", name, " already registered");
+    }
+    return OkStatus();
+  }
+
+  nvinfer1::ILogger* LookUp(const string& name) override {
+    mutex_lock lock(mu_);
+    const auto found = registry_.find(name);
+    if (found == registry_.end()) {
+      return nullptr;
+    }
+    return found->second.get();
+  }
+
+ private:
+  mutable mutex mu_;
+  mutable std::unordered_map<string, std::unique_ptr<nvinfer1::ILogger>>
+      registry_ TF_GUARDED_BY(mu_);
+};
+
+LoggerRegistry* GetLoggerRegistry() {
+  static LoggerRegistryImpl* registry = new LoggerRegistryImpl;
+  return registry;
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/logger_registry.h

@@ -0,0 +1,58 @@
+/* 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_CONVERT_LOGGER_REGISTRY_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_LOGGER_REGISTRY_H_
+
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+class LoggerRegistry {
+ public:
+  virtual Status Register(const string& name, nvinfer1::ILogger* logger) = 0;
+  virtual nvinfer1::ILogger* LookUp(const string& name) = 0;
+  virtual ~LoggerRegistry() {}
+};
+
+LoggerRegistry* GetLoggerRegistry();
+
+class RegisterLogger {
+ public:
+  RegisterLogger(const string& name, nvinfer1::ILogger* logger) {
+    TF_CHECK_OK(GetLoggerRegistry()->Register(name, logger));
+  }
+};
+
+#define REGISTER_TENSORRT_LOGGER(name, logger) \
+  REGISTER_TENSORRT_LOGGER_UNIQ_HELPER(__COUNTER__, name, logger)
+#define REGISTER_TENSORRT_LOGGER_UNIQ_HELPER(ctr, name, logger) \
+  REGISTER_TENSORRT_LOGGER_UNIQ(ctr, name, logger)
+#define REGISTER_TENSORRT_LOGGER_UNIQ(ctr, name, logger)                 \
+  static ::tensorflow::tensorrt::RegisterLogger register_trt_logger##ctr \
+      TF_ATTRIBUTE_UNUSED =                                              \
+          ::tensorflow::tensorrt::RegisterLogger(name, logger)
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_LOGGER_REGISTRY_H_

tensorflow/compiler/tf2tensorrt/convert/logger_registry_test.cc

@@ -0,0 +1,34 @@
+/* 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 <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace {
+
+class TestLogger : public nvinfer1::ILogger {
+  void log(nvinfer1::ILogger::Severity severity, const char* msg) override {}
+};
+
+TestLogger test_logger;
+
+REGISTER_TENSORRT_LOGGER("test_logger", &test_logger);
+
+TEST(LoggerRegistryTest, RegistersCorrectly) {
+  auto registered_logger = GetLoggerRegistry()->LookUp("test_logger");
+  EXPECT_THAT(registered_logger, Eq(&test_logger));
+}
+
+}  // namespace

tensorflow/compiler/tf2tensorrt/convert/op_converter.h

@@ -0,0 +1,224 @@
+/* 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_CONVERT_OP_CONVERTER_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OP_CONVERTER_H_
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <memory>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/trt_parameters.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/weights.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class Converter;
+
+// Specifies the expected type taken by a TRT_TensorOrWeights input during op
+// conversion.
+// kResource is only used for resource variable ops. For an operation like
+// Add(tensor, ReadVariableOp(...)), the second operand of Add is the result of
+// the ReadVariableOp, which is a kWeight.
+enum class TrtInputArg { kTensor = 1, kWeight = 2, kBoth = 3, kResource = 4 };
+
+// Parameters for each op converter.
+struct OpConverterParams {
+  // Constructor used for validation only.
+  OpConverterParams(const NodeDef& node_def,
+                    const std::vector<TRT_TensorOrWeights>& inputs,
+                    std::vector<TRT_TensorOrWeights>* outputs,
+                    TrtWeightStore* weight_store,
+                    TrtPrecisionMode precision_mode, bool use_calibration,
+                    bool use_implicit_batch, bool use_explicit_precision);
+
+  // Constructor used for conversion.
+  OpConverterParams(Converter* converter, const NodeDef& node_def,
+                    const std::vector<TRT_TensorOrWeights>& inputs,
+                    std::vector<TRT_TensorOrWeights>* outputs,
+                    TrtWeightStore* weight_store);
+
+  Converter* converter = nullptr;
+  const NodeDef& node_def;
+  const std::vector<TRT_TensorOrWeights>& inputs;
+  std::vector<TRT_TensorOrWeights>* outputs;
+  const bool validation_only;
+  TrtWeightStore* weight_store;
+  const TrtPrecisionMode precision_mode;
+  const bool use_calibration;
+  const bool use_implicit_batch;
+  const bool use_explicit_precision;
+};
+
+// Operation converter function specification.
+using OpConverter = std::function<Status(const OpConverterParams*)>;
+
+struct InputArgSpec {
+  absl::string_view name;
+  TrtInputArg allowed_roles;
+
+  static constexpr InputArgSpec Create(absl::string_view n, TrtInputArg role) {
+    return InputArgSpec{n, role};
+  }
+};
+
+template <typename T>
+std::string convert_not_supported_dtype_msg(const T& allowed_types,
+                                            DataType tf_type,
+                                            const NodeDef& node) {
+  string allowed_types_string =
+      absl::StrJoin(allowed_types, ", ", [](string* out, const DataType& type) {
+        absl::StrAppendFormat(out, "%s", DataTypeString(type));
+      });
+
+  return absl::StrCat("Data type ", DataTypeString(tf_type),
+                      " is not supported for ", node.op(), ", must be one of [",
+                      allowed_types_string, "]");
+}
+
+std::string convert_not_supported_implicit(const std::string& pOpName,
+                                           const std::string& pNodeName,
+                                           const char* pOpType = NULL);
+
+// A Curiously recurring template pattern (CRTP) template class for operation
+// converters.
+template <typename Impl>
+class OpConverterBase {
+ public:
+  explicit OpConverterBase(const OpConverterParams* params,
+                           const std::vector<DataType>& data_types =
+                               {DataType::DT_FLOAT, DataType::DT_HALF})
+      : params_(params),
+        node_def_attrs_(params->node_def),
+        allowed_dtypes_(data_types) {}
+
+  // Default NodeDef attribute name to inspect in order to determine node data
+  // type. The Impl class can override this by implementing the same function.
+  static constexpr const char* NodeDefDataTypeAttributeName() { return "T"; }
+
+  // Validate data type of the given NodeDef against allowed types.
+  Status ValidateNodeDefDataType() {
+    // If the attribute name is empty, we should skip this check.
+    if (absl::string_view(Impl::NodeDefDataTypeAttributeName()).empty()) {
+      return OkStatus();
+    }
+
+    // Get the NodeDef data type.
+    auto dtype = GetAttrValue<DataType>(Impl::NodeDefDataTypeAttributeName());
+    if (!dtype.ok()) {
+      return errors::InvalidArgument("Attribute with name ",
+                                     Impl::NodeDefDataTypeAttributeName(),
+                                     " not found.");
+    }
+
+    // Check allowed data types.;
+    if (std::find(allowed_dtypes_.begin(), allowed_dtypes_.end(), *dtype) ==
+        allowed_dtypes_.end()) {
+      return errors::Unimplemented(convert_not_supported_dtype_msg(
+          allowed_dtypes_, *dtype, params_->node_def));
+    }
+    return OkStatus();
+  }
+
+  static constexpr bool HasFixNumberOfInputs() { return true; }
+
+  // Validates input argument roles and data types.
+  Status ValidateInputs() {
+    const NodeDef& node_def = params_->node_def;
+    const auto& inputs = params_->inputs;
+    if (Impl::HasFixNumberOfInputs()) {
+      TRT_ENSURE(inputs.size() == Impl::InputSpec().size());
+    } else {
+      TRT_ENSURE(inputs.size() <= Impl::InputSpec().size());
+    }
+    for (int i = 0; i < inputs.size(); i++) {
+      const InputArgSpec arg_spec = Impl::InputSpec()[i];
+      if (arg_spec.allowed_roles == TrtInputArg::kWeight &&
+          inputs.at(i).is_tensor()) {
+        return errors::Unimplemented("The input \"", arg_spec.name, "\" for ",
+                                     node_def.op(), " must be a constant, at ",
+                                     node_def.name());
+      }
+      if (arg_spec.allowed_roles == TrtInputArg::kTensor &&
+          inputs.at(i).is_weights()) {
+        return errors::Unimplemented("The input \"", arg_spec.name, "\" for ",
+                                     node_def.op(), " must be a tensor, at ",
+                                     node_def.name());
+      }
+    }
+    return OkStatus();
+  }
+
+  Status operator()() {
+    // Validate data type and inputs.
+    TF_RETURN_IF_ERROR(this->ValidateNodeDefDataType());
+    TF_RETURN_IF_ERROR(this->ValidateInputs());
+
+    // Perform op-level validation.
+    TF_RETURN_IF_ERROR(reinterpret_cast<Impl*>(this)->Validate());
+    if (params_->validation_only) {
+      return OkStatus();
+    }
+
+    // Perform conversion.
+    return reinterpret_cast<Impl*>(this)->Convert();
+  }
+
+ protected:
+  Status NotSupportedInImplicitBatch(const char* pOpType = nullptr) {
+    if (params_->use_implicit_batch) {
+      const auto& op = params_->node_def.op();
+      const auto& nodeName = params_->node_def.name();
+      const auto& error = convert_not_supported_implicit(op, nodeName, pOpType);
+      return errors::Unimplemented(error);
+    }
+    return OkStatus();
+  }
+
+  void AddOutput(const TRT_TensorOrWeights& out) {
+    params_->outputs->push_back(out);
+  }
+
+  template <typename T>
+  StatusOr<T> GetAttrValue(absl::string_view key) const {
+    T result;
+    TF_RETURN_IF_ERROR(GetNodeAttr(node_def_attrs_, key, &result));
+    return result;
+  }
+
+  const OpConverterParams* const params_;
+  const AttrSlice node_def_attrs_;
+  const std::vector<DataType> allowed_dtypes_;
+};
+
+// Constructs and returns a converter function for a given operation converter
+// class T. This requires T to be a derived class of StructuredOpConverter.
+template <typename T>
+OpConverter MakeConverterFunction() {
+  return [](const OpConverterParams* params) -> Status {
+    T converter(params);
+    return converter();
+  };
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OP_CONVERTER_H_

tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.cc

@@ -0,0 +1,157 @@
+/* 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/convert/op_converter_registry.h"
+
+#include <set>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/util/env_var.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+struct OpConverterRegistration {
+  OpConverter converter;
+  int priority;
+};
+class OpConverterRegistry::Impl {
+ public:
+  ~Impl() = default;
+
+  InitOnStartupMarker Register(const string& name, const int priority,
+                               OpConverter converter) {
+    mutex_lock lock(mu_);
+    auto item = registry_.find(name);
+    if (item != registry_.end()) {
+      const int existing_priority = item->second.priority;
+      if (priority <= existing_priority) {
+        LOG(WARNING) << absl::StrCat(
+            "Ignoring TF->TRT ", name, " op converter with priority ",
+            existing_priority, " due to another converter with priority ",
+            priority);
+        return {};
+      } else {
+        LOG(WARNING) << absl::StrCat(
+            "Overwriting TF->TRT ", name, " op converter with priority ",
+            existing_priority, " using another converter with priority ",
+            priority);
+        registry_.erase(item);
+      }
+    }
+    registry_.insert({name, OpConverterRegistration{converter, priority}});
+    return {};
+  }
+
+  StatusOr<OpConverter> LookUp(string name) {
+    // Fetch the user-provide TF operations denylisted for conversion by TF-TRT.
+    static const absl::flat_hash_set<string> tftrt_op_fakelist = [] {
+      string tftrt_op_fakelist_str;
+      TF_CHECK_OK(ReadStringFromEnvVar("TF_TRT_OP_FAKELIST",
+                                       /*default_value=*/"",
+                                       &tftrt_op_fakelist_str));
+      absl::flat_hash_set<string> tftrt_op_fakelist{};
+      for (const auto& x : str_util::Split(tftrt_op_fakelist_str, ",")) {
+        tftrt_op_fakelist.insert(x);
+      }
+      // Force a rehash of the flat hash set
+      tftrt_op_fakelist.rehash(0);
+      return tftrt_op_fakelist;
+    }();
+
+    // In case the TensorFlow OP `name` matches any of the names passed to
+    // TF_TRT_OP_FAKELIST environment variable, force ::LookUp to resolves to
+    // ConvertFake OP converter.
+    if (tftrt_op_fakelist.contains(name)) {
+      LOG_FIRST_N(INFO, 2) << "Emulating OP Converter: `" << name << "`. It "
+                           << "will cause TRT engine building to fail. This "
+                           << "feature is only intended to be used for "
+                           << "TF-TRT graph segmentation experiments. This "
+                           << "feature is controlled using: "
+                           << "`TF_TRT_OP_FAKELIST=OpName1,OpName2`.";
+      // Forces ::LookUp to resolve to `ConvertFake` registered to `FakeOp`.
+      mutex_lock lock(mu_);
+      return registry_.find("FakeOp")->second.converter;
+    }
+
+    mutex_lock lock(mu_);
+    auto found = registry_.find(name);
+    if (found != registry_.end()) {
+      return found->second.converter;
+    }
+    return errors::NotFound("No converter for op ", name);
+  }
+
+  void Clear(const std::string& name) {
+    mutex_lock lock(mu_);
+    auto itr = registry_.find(name);
+    if (itr == registry_.end()) {
+      return;
+    }
+    registry_.erase(itr);
+  }
+
+  std::vector<std::string> ListRegisteredOps() const {
+    mutex_lock lock(mu_);
+    std::vector<std::string> result;
+    result.reserve(registry_.size());
+    for (const auto& item : registry_) {
+      result.push_back(item.first);
+    }
+    return result;
+  }
+
+ private:
+  mutable mutex mu_;
+  mutable std::unordered_map<std::string, OpConverterRegistration> registry_
+      TF_GUARDED_BY(mu_);
+};
+
+OpConverterRegistry::OpConverterRegistry() : impl_(std::make_unique<Impl>()) {}
+
+StatusOr<OpConverter> OpConverterRegistry::LookUp(const string& name) {
+  return impl_->LookUp(name);
+}
+
+InitOnStartupMarker OpConverterRegistry::Register(const string& name,
+                                                  const int priority,
+                                                  OpConverter converter) {
+  return impl_->Register(name, priority, converter);
+}
+
+std::vector<std::string> OpConverterRegistry::ListRegisteredOps() const {
+  return impl_->ListRegisteredOps();
+}
+
+void OpConverterRegistry::Clear(const std::string& name) { impl_->Clear(name); }
+
+OpConverterRegistry* GetOpConverterRegistry() {
+  static OpConverterRegistry* registry = new OpConverterRegistry();
+  return registry;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h

@@ -0,0 +1,104 @@
+/* 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_CONVERT_OP_CONVERTER_REGISTRY_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OP_CONVERTER_REGISTRY_H_
+
+#include <initializer_list>
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <array>
+#include <type_traits>
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class OpConverterRegistry {
+ public:
+  OpConverterRegistry();
+  ~OpConverterRegistry() = default;
+
+  InitOnStartupMarker Register(const string& name, const int priority,
+                               OpConverter converter);
+
+  InitOnStartupMarker Register(const std::initializer_list<std::string>& names,
+                               const int priority, OpConverter converter) {
+    for (const auto& name : names) {
+      Register(name, priority, converter);
+    }
+    return {};
+  }
+
+  template <typename T,
+            typename std::enable_if<std::is_convertible<
+                typename T::value_type, std::string>::value>::type* = nullptr>
+  InitOnStartupMarker Register(const T& names, const int priority,
+                               OpConverter converter) {
+    for (const auto& name : names) {
+      Register(name, priority, converter);
+    }
+    return {};
+  }
+
+  // Clear all registered converters for the given Tensorflow operation name.
+  void Clear(const std::string& name);
+
+  StatusOr<OpConverter> LookUp(const string& name);
+
+  std::vector<std::string> ListRegisteredOps() const;
+
+ private:
+  class Impl;
+  std::unique_ptr<Impl> impl_;
+};
+
+OpConverterRegistry* GetOpConverterRegistry();
+
+class RegisterOpConverter {
+ public:
+  RegisterOpConverter(const string& name, const int priority,
+                      OpConverter converter) {
+    GetOpConverterRegistry()->Register(name, priority, converter);
+  }
+};
+
+constexpr int kDefaultConverterPriority = 1;
+
+}  // namespace convert
+}  // namespace tensorrt
+
+#define REGISTER_TRT_OP_CONVERTER_IMPL(ctr, func, priority, ...)    \
+  static ::tensorflow::InitOnStartupMarker const                    \
+      register_trt_op_converter##ctr TF_ATTRIBUTE_UNUSED =          \
+          TF_INIT_ON_STARTUP_IF(true)                               \
+          << tensorrt::convert::GetOpConverterRegistry()->Register( \
+                 __VA_ARGS__, priority, func)
+
+#define REGISTER_TRT_OP_CONVERTER(func, priority, ...)               \
+  TF_NEW_ID_FOR_INIT(REGISTER_TRT_OP_CONVERTER_IMPL, func, priority, \
+                     __VA_ARGS__)
+
+#define REGISTER_DEFAULT_TRT_OP_CONVERTER(func, ...) \
+  REGISTER_TRT_OP_CONVERTER(                         \
+      func, tensorrt::convert::kDefaultConverterPriority, __VA_ARGS__)
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OP_CONVERTER_REGISTRY_H_

tensorflow/compiler/tf2tensorrt/convert/op_converter_registry_test.cc

@@ -0,0 +1,67 @@
+/* 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/convert/op_converter_registry.h"
+
+#include <gtest/gtest.h>
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+TEST(TestOpConverterRegistry, TestOpConverterRegistry) {
+  bool flag{false};
+
+  auto set_true_func = [&flag](const OpConverterParams*) -> Status {
+    flag = true;
+    return OkStatus();
+  };
+
+  auto set_false_func = [&flag](const OpConverterParams*) -> Status {
+    flag = false;
+    return OkStatus();
+  };
+
+  GetOpConverterRegistry()->Register("FakeFunc", kDefaultConverterPriority,
+                                     set_true_func);
+
+  // Lower priority fails to override.
+  GetOpConverterRegistry()->Register("FakeFunc", kDefaultConverterPriority - 1,
+                                     set_false_func);
+
+  // The lookup should return set_true_func (default).
+  auto func = GetOpConverterRegistry()->LookUp("FakeFunc");
+  EXPECT_TRUE(func.ok());
+  EXPECT_TRUE(((*func)(nullptr)).ok());
+  EXPECT_TRUE(flag);
+
+  // Override with higher priority.
+  GetOpConverterRegistry()->Register("FakeFunc", kDefaultConverterPriority + 1,
+                                     set_false_func);
+  func = GetOpConverterRegistry()->LookUp("FakeFunc");
+  EXPECT_TRUE(func.ok());
+  EXPECT_TRUE((*func)(nullptr).ok());
+  EXPECT_FALSE(flag);
+
+  // After clearing the op, lookup should return an error.
+  GetOpConverterRegistry()->Clear("FakeFunc");
+  EXPECT_FALSE(GetOpConverterRegistry()->LookUp("FakeFunc").ok());
+}
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif

tensorflow/compiler/tf2tensorrt/convert/op_converter_test.cc

@@ -0,0 +1,123 @@
+/* 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/convert/op_converter.h"
+
+#include <gtest/gtest.h>
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+using ::tensorflow::testing::IsOk;
+using ::tensorflow::testing::StatusIs;
+using ::testing::HasSubstr;
+
+class ExampleOpConverter : public OpConverterBase<ExampleOpConverter> {
+ public:
+  explicit ExampleOpConverter(const OpConverterParams* params)
+      : OpConverterBase<ExampleOpConverter>(params, {DataType::DT_FLOAT}) {}
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    return "data_type";
+  }
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("input_tensor", TrtInputArg::kTensor),
+        InputArgSpec::Create("weight", TrtInputArg::kWeight)};
+  }
+
+  Status Validate() { return OkStatus(); }
+
+  Status Convert() {
+    AddOutput(TRT_TensorOrWeights(nvinfer1::DataType::kFLOAT,
+                                  nvinfer1::Dims{1, {1, 1, 1}}, 1));
+    return OkStatus();
+  }
+};
+
+TEST(TestOpConverterBase, TestOpConverterBase) {
+  // Register a converter which uses the base converter class.
+  GetOpConverterRegistry()->Register(
+      "FakeFunc", 1, MakeConverterFunction<ExampleOpConverter>());
+
+  NodeDef def;
+  def.set_op("FakeFunc");
+  auto converter = Converter::Create(TrtPrecisionMode::FP32, false,
+                                     Logger::GetLogger(), false, "test_engine");
+  EXPECT_THAT(converter, IsOk());
+
+  // Base class should check attribute with key given by
+  // Impl::NodeDefDataTypeAttributeName().
+  Status conversion_status = (*converter)->ConvertNode(def);
+  EXPECT_THAT(conversion_status,
+              StatusIs(error::INVALID_ARGUMENT,
+                       HasSubstr("Attribute with name data_type not found")));
+
+  // Add partial inputs to the node and make the converter aware.
+  def.mutable_input()->Add("input1");
+  conversion_status = (*converter)
+                          ->AddInputTensor("input1", nvinfer1::DataType::kFLOAT,
+                                           nvinfer1::Dims{4, {1, 1, 1, 1}}, 1);
+  EXPECT_THAT(conversion_status, IsOk());
+
+  // Base class method should check number of inputs.
+  AddNodeAttr("data_type", DT_FLOAT, &def);
+  conversion_status = (*converter)->ConvertNode(def);
+  EXPECT_THAT(conversion_status, StatusIs(error::INTERNAL));
+
+  // Add second input to the node and make the converter aware.
+  def.mutable_input()->Add("input2");
+  conversion_status = (*converter)
+                          ->AddInputTensor("input2", nvinfer1::DataType::kFLOAT,
+                                           nvinfer1::Dims{4, {1, 1, 1, 1}}, 1);
+  EXPECT_THAT(conversion_status, IsOk());
+
+  // Base class validation should check the type (Constant or Tensor) of the
+  // inputs.
+  conversion_status = (*converter)->ConvertNode(def);
+  EXPECT_THAT(
+      conversion_status,
+      StatusIs(error::UNIMPLEMENTED,
+               HasSubstr("input \"weight\" for FakeFunc must be a constant")));
+
+  // Correct input2 so that it is a weight.
+  (*converter)->TensorsMap().erase("input2");
+  (*converter)
+      ->TensorsMap()
+      .insert(std::make_pair("input2", TRT_TensorOrWeights(TRT_ShapedWeights(
+                                           nvinfer1::DataType::kFLOAT))));
+
+  // With the correct input types, check that the converter is called and sets
+  // one output tensor.
+  conversion_status = (*converter)->ConvertNode(def);
+  EXPECT_THAT(conversion_status, IsOk());
+  EXPECT_EQ((*converter)->TensorsMap().size(), 3U);
+
+  GetOpConverterRegistry()->Clear("FakeFunc");
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif

tensorflow/compiler/tf2tensorrt/convert/ops/binary_ops.cc

@@ -0,0 +1,235 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+const BinaryOperationMapType* BinaryOperationMap() {
+  static const auto* map = new BinaryOperationMapType({
+    {"Add", nvinfer1::ElementWiseOperation::kSUM},
+        {"AddV2", nvinfer1::ElementWiseOperation::kSUM},
+        {"Mul", nvinfer1::ElementWiseOperation::kPROD},
+        {"Sub", nvinfer1::ElementWiseOperation::kSUB},
+        {"Div", nvinfer1::ElementWiseOperation::kDIV},
+        {"FloorDiv", nvinfer1::ElementWiseOperation::kFLOOR_DIV},
+        {"RealDiv", nvinfer1::ElementWiseOperation::kDIV},
+        {"Minimum", nvinfer1::ElementWiseOperation::kMIN},
+        {"Maximum", nvinfer1::ElementWiseOperation::kMAX},
+        {"Pow", nvinfer1::ElementWiseOperation::kPOW},
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+        {"Greater", nvinfer1::ElementWiseOperation::kGREATER},
+        {"Less", nvinfer1::ElementWiseOperation::kLESS},
+        {"Equal", nvinfer1::ElementWiseOperation::kEQUAL},
+        // Operators are implemented as NOT Less and NOT Greater, respectively.
+        {"GreaterEqual", nvinfer1::ElementWiseOperation::kLESS},
+        {"LessEqual", nvinfer1::ElementWiseOperation::kGREATER},
+#endif
+  });
+  return map;
+}
+
+const BinaryOperationMapType* BinaryBooleanOperationMap() {
+  static const auto* map = new BinaryOperationMapType({
+      {"LogicalOr", nvinfer1::ElementWiseOperation::kOR},
+      {"LogicalAnd", nvinfer1::ElementWiseOperation::kAND},
+  });
+  return map;
+}
+
+namespace {
+class ConvertBinaryImpl {
+ protected:
+  ConvertBinaryImpl(const BinaryOperationMapType* pOperMap)
+      : pOperMap_(pOperMap) {}
+
+  Status ValidateImpl(
+      const OpConverterParams& params,
+      const std::vector<string>& implicit_batch_not_supported_ops = {},
+      bool both_tensors = false) {
+    const auto& node_def = params.node_def;
+    const auto& op = node_def.op();
+    const auto op_pair = pOperMap_->find(op);
+    if (op_pair == pOperMap_->end()) {
+      return errors::Unimplemented("Binary op: ", op, " not supported");
+    }
+
+    // Constant folding should have been done by TensorFlow.
+    const auto& inputs = params.inputs;
+    if (inputs.at(0).is_weights() && inputs.at(1).is_weights()) {
+      return errors::Unimplemented(
+          "Constant folding is falled back to TensorFlow, binary op '", op,
+          "' received both input as constant");
+    }
+
+    if ((convertToBool_ = find_name(op, implicit_batch_not_supported_ops))) {
+      if (params.use_implicit_batch) {
+        return errors::Unimplemented(
+            convert_not_supported_implicit(op, node_def.name(), "Binary"));
+      }
+    }
+
+    if (both_tensors) {
+      if (inputs.at(0).is_weights() || inputs.at(1).is_weights()) {
+        return errors::InvalidArgument("Both inputs  of '", op,
+                                       "' are expected to be tensors");
+      }
+      // No need to convert the output of "LogicalOr" and "LogicalAnd"
+      convertToBool_ = false;
+    }
+
+    nvinfer1::Dims broadcasted_dims[2];
+    TF_RETURN_IF_ERROR(GetTrtBroadcastShape(
+        inputs.at(0), inputs.at(1), true, params.use_implicit_batch,
+        broadcasted_dims, broadcasted_dims + 1));
+
+    for (int i = 0; i < tensor_.size(); i++) {
+      // This will also convert constants to tensors.
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params.converter, inputs.at(i), broadcasted_dims[i],
+          params.validation_only, &tensor_[i], node_def, i));
+    }
+    operation_ = op_pair->second;
+    return OkStatus();
+  }
+
+  Status ConvertImpl(const OpConverterParams& params,
+                     const std::vector<string>& revert_bool_ops = {}) {
+    const auto& node_def = params.node_def;
+    // Add ElementWise layer.
+    auto* network = params.converter->network();
+    nvinfer1::ILayer* layer = network->addElementWise(
+        *tensor_[0]->trt_tensor(), *tensor_[1]->trt_tensor(), operation_);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+
+    if (params.use_explicit_precision) {
+      layer->setPrecision(nvinfer1::DataType::kFLOAT);
+    }
+
+    params.converter->SetLayerName(layer, node_def);
+    const auto& output = layer->getOutput(0);
+    if (convertToBool_) {
+      output->setType(nvinfer1::DataType::kBOOL);
+      if (find_name(node_def.op(), revert_bool_ops)) {
+        nvinfer1::IUnaryLayer* unary_layer =
+            network->addUnary(*output, nvinfer1::UnaryOperation::kNOT);
+        TFTRT_RETURN_ERROR_IF_NULLPTR(unary_layer, node_def.name());
+        params.outputs->push_back(
+            TRT_TensorOrWeights(unary_layer->getOutput(0)));
+        return OkStatus();
+      }
+    }
+
+    params.outputs->push_back(TRT_TensorOrWeights(output));
+    return OkStatus();
+  }
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("x", TrtInputArg::kBoth),
+        InputArgSpec::Create("y", TrtInputArg::kBoth)};
+  }
+
+ private:
+  const BinaryOperationMapType* pOperMap_;
+  std::array<ITensorProxyPtr, 2> tensor_{nullptr, nullptr};
+  nvinfer1::ElementWiseOperation operation_;
+  bool convertToBool_;
+};
+
+class ConvertBinary : public OpConverterBase<ConvertBinary>,
+                      protected ConvertBinaryImpl {
+ public:
+  explicit ConvertBinary(const OpConverterParams* params)
+      : OpConverterBase<ConvertBinary>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}),
+        ConvertBinaryImpl(BinaryOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return ConvertBinaryImpl::InputSpec();
+  }
+
+  Status Validate() {
+    const std::vector<string> implicit_batch_not_supported_ops {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+      "Greater", "Less", "Equal", "GreaterEqual", "LessEqual"
+#endif
+    };
+    return ValidateImpl(*params_, implicit_batch_not_supported_ops);
+  }
+  Status Convert() {
+    const std::vector<string> implemented_with_reverted_ops {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+      "GreaterEqual", "LessEqual"
+#endif
+    };
+    return ConvertImpl(*params_, implemented_with_reverted_ops);
+  }
+};
+
+class ConvertBooleanBinary : public OpConverterBase<ConvertBooleanBinary>,
+                             public ConvertBinaryImpl {
+ public:
+  explicit ConvertBooleanBinary(const OpConverterParams* params)
+      : OpConverterBase<ConvertBooleanBinary>(params, {DataType::DT_BOOL}),
+        ConvertBinaryImpl(BinaryBooleanOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return ConvertBinaryImpl::InputSpec();
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "LogicalOr"
+      input: "..."
+      input: "..."
+      attr {
+        key: "_output_shapes"
+        ...
+      }
+    }
+    */
+    return "";
+  }
+  Status Validate() {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    return ValidateImpl(*params_, {"LogicalOr", "LogicalAnd"}, true);
+#else
+    return errors::Unimplemented("Boolean op: ", params_->node_def.op(),
+                                 " is not supported in TRT version < 8.2");
+#endif
+  }
+  Status Convert() { return ConvertImpl(*params_); }
+};
+}  // namespace
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertBinary>(),
+                                  GetOperationNames(*BinaryOperationMap()));
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertBooleanBinary>(),
+    GetOperationNames(*BinaryBooleanOperationMap()));
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/data_format_vec_permute.cc

@@ -0,0 +1,176 @@
+/* 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/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "third_party/tensorrt/NvInfer.h"
+#include "third_party/tensorrt/NvInferRuntimeCommon.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+int get_spatial_dim_count(string format) {
+  // Spatial dimensions are the dimensions besides NC, and here we assume NC
+  // always appear in the format string.
+  return format.size() - 2;
+}
+
+class ConvertDataFormatVecPermute
+    : public OpConverterBase<ConvertDataFormatVecPermute> {
+ public:
+  ConvertDataFormatVecPermute(const OpConverterParams* params)
+      : OpConverterBase<ConvertDataFormatVecPermute>(params,
+                                                     {DataType::DT_INT32}) {}
+
+  struct DataFormatVecPermuteAttributes {
+    string dst_format;
+    string src_format;
+    int x_dim_count;
+  };
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return {InputArgSpec::Create("x", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    const auto& inputs = params_->inputs;
+    const auto& nodeName = params_->node_def.name();
+
+    x_input_ = inputs.at(0);
+
+    // Check input rank.
+    const auto x_dims = x_input_.GetTrtDims();
+    int input_rank = x_dims.nbDims;
+    if (input_rank != 1 && input_rank != 2) {
+      return errors::InvalidArgument(
+          "Input must be a vector or matrix, but got rank ", input_rank,
+          ", at ", nodeName);
+    }
+
+    // Verify and consume node attributes.
+    StatusOr<string> dst_format = GetAttrValue<string>("dst_format");
+    StatusOr<string> src_format = GetAttrValue<string>("src_format");
+    TRT_ENSURE_OK(dst_format);
+    TRT_ENSURE_OK(src_format);
+
+    // Check input dims.
+    const int full_dim_count = src_format->size();
+    const int spatial_dim_count = get_spatial_dim_count(*src_format);
+    if (input_rank == 1) {
+      if (x_dims.d[0] != spatial_dim_count && x_dims.d[0] != full_dim_count) {
+        return errors::InvalidArgument(
+            "1D input must be of size ", spatial_dim_count, " or ",
+            full_dim_count, ", but got size ", x_dims.d[0], ", at ", nodeName);
+      }
+    } else if (input_rank == 2) {
+      if (x_dims.d[0] != spatial_dim_count && x_dims.d[0] != full_dim_count) {
+        return errors::InvalidArgument(
+            "First dimension of 2D input must be of size ", spatial_dim_count,
+            " or ", full_dim_count, ", but got shape (", x_dims.d[0], ", ",
+            x_dims.d[1], "), at ", nodeName);
+      }
+      if (x_dims.d[1] != 2) {
+        return errors::InvalidArgument(
+            "Second dimension of 2D input must be of size 2, but got shape (",
+            x_dims.d[0], ", ", x_dims.d[1], "), at ", nodeName);
+      }
+    }
+
+    // Set custom attributes.
+    attrs_.x_dim_count = x_dims.d[0];
+    attrs_.dst_format = *dst_format;
+    attrs_.src_format = *src_format;
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    // Copy format strings in case they need to be modified.
+    string dst_format = attrs_.dst_format;
+    string src_format = attrs_.src_format;
+    const int& spatial_dim_count = get_spatial_dim_count(src_format);
+
+    // If the input is a vector of size spatial_dim_count, treat the elements
+    // as spatial dimensions.
+    if (attrs_.x_dim_count == spatial_dim_count) {
+      auto keep_only_spatial_dimensions =
+          [spatial_dim_count](string* format_str) -> void {
+        auto new_end = std::remove_if(format_str->begin(), format_str->end(),
+                                      [spatial_dim_count](const char dim) {
+                                        return dim == 'N' || dim == 'C';
+                                      });
+        format_str->erase(new_end, format_str->end());
+      };
+      keep_only_spatial_dimensions(&src_format);
+      keep_only_spatial_dimensions(&dst_format);
+    }
+
+    // Create indices for the gather layer and make weights out of them.
+    std::vector<int32> dst_indices(attrs_.x_dim_count);
+    for (int i = 0; i < attrs_.x_dim_count; ++i) {
+      for (int j = 0; j < attrs_.x_dim_count; ++j) {
+        if (src_format[i] == dst_format[j]) {
+          dst_indices[j] = i;
+          break;
+        }
+      }
+    }
+    nvinfer1::Dims indices_dims = {1, {attrs_.x_dim_count}};
+    StatusOr<TRT_ShapedWeights> indices_weights =
+        params_->weight_store->GetTempWeights(nvinfer1::DataType::kINT32,
+                                              indices_dims);
+    TRT_ENSURE_OK(indices_weights);
+    int32* indices_ptr = indices_weights->GetPointer<int32>();
+    std::copy(dst_indices.data(), dst_indices.data() + attrs_.x_dim_count,
+              indices_ptr);
+    ITensorProxyPtr x_tensor =
+        x_input_.is_weights() ? params_->converter->CreateConstantLayer(
+                                    x_input_.weights(), x_input_.GetTrtDims())
+                              : x_input_.tensor();
+    ITensorProxyPtr indices_tensor =
+        params_->converter->CreateConstantLayer(*indices_weights, indices_dims);
+
+    // Gather layer with 1D indices on axis 0, conserves shape.
+    nvinfer1::IGatherLayer* layer = params_->converter->network()->addGather(
+        *x_tensor->trt_tensor(), *indices_tensor->trt_tensor(), 0);
+    TRT_ENSURE(layer);
+    params_->converter->SetLayerName(layer, params_->node_def);
+
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+
+    params_->outputs->push_back(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+
+ private:
+  TRT_TensorOrWeights x_input_;
+  DataFormatVecPermuteAttributes attrs_{};
+};
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertDataFormatVecPermute>(),
+    {"DataFormatVecPermute"});
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/einsum.cc

@@ -0,0 +1,908 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include <iterator>
+#include <limits>
+#include <memory>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/util/einsum_op_util.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+#if IS_TRT_VERSION_GE(7, 1, 3, 0)
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+namespace {
+
+#if !IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+// Finds the indices of elements in [begin, end) in array
+// [array_begin, array_end), and appends the indices to permute. This is used to
+// construct the permutation sequence for the operand with input labels
+// [array_begin, array_end) to the desired permuted labels [begin, end).
+template <typename T>
+Status FindIndicesoOfAllValuesInSrc(absl::Span<const T> values,
+                                    absl::Span<const T> src,
+                                    std::vector<int>* indices) {
+  if (src.size() < values.size()) {
+    return errors::Internal(
+        "Span 'src' cannot contain all elements of 'values'");
+  }
+  for (auto i = 0; i < values.size(); i++) {
+    auto iter = absl::c_find(src, values[i]);
+    if (iter == src.end()) {
+      return errors::Internal("Label ", values[i], " not found");
+    }
+    int idx = std::distance(src.begin(), iter);
+    indices->push_back(idx);
+  }
+  return OkStatus();
+}
+
+// Layout of the einsum dimensions: Batch, Free and Contraction indices.
+// Example: adbc,adce -> adbe. The first tensor has layout BFC, the second BCF.
+enum class EinsumLayout { BFC, BCF, MIX };
+
+using DimType = EinsumDimensionType;
+constexpr auto kBatch = DimType::kBatch;
+constexpr auto kFree = DimType::kFree;
+constexpr auto kContract = DimType::kContract;
+
+// Describes an operand: input shape, number of batch, free and contract
+// dimensions, and the permutation that is needed to bring it to a matmul
+// compatible form.
+class EinsumDescriptor {
+ private:
+  // Checks whether input_labels[offset:offset+m] matches labels from other.
+  static bool OrderMatches(const Labels& input_labels, int offset, int m,
+                           EinsumDimensionType dim_type,
+                           const std::unique_ptr<EinsumDescriptor>& other) {
+    if (other == nullptr) {
+      return true;
+    }
+    int offset_other = 0;
+    if (dim_type == kFree) {
+      offset = other->offset_f;
+    } else if (dim_type == kContract) {
+      offset = other->offset_c;
+    }
+    return std::equal(input_labels.begin() + offset,
+                      input_labels.begin() + offset + m,
+                      other->permuted_labels.begin() + offset_other);
+  }
+
+  using label_t_iterator = std::vector<EinsumDimensionType>::const_iterator;
+  static int32_t CountLabels(label_t_iterator begin, label_t_iterator end,
+                             EinsumDimensionType val) {
+    return static_cast<int32_t>(std::count_if(
+        begin, end, [val](EinsumDimensionType t) { return t == val; }));
+  }
+
+  // Appends indices to the "permute" vector where types matches value.
+  void AppendMatchingIndicesToPermute(
+      const std::vector<EinsumDimensionType>& types, EinsumDimensionType val) {
+    for (int i = 0; i < types.size(); i++) {
+      if (types[i] == val) {
+        permute.push_back(i);
+      }
+    }
+  }
+
+  Status DetermineLayout(const Labels& input_labels,
+                         const std::vector<EinsumDimensionType>& types,
+                         const std::unique_ptr<EinsumDescriptor>& other) {
+    // Check if the current layout is BFC or BCF. In that case we could avoid
+    // transpose.
+    layout = EinsumLayout::MIX;
+    if (CountLabels(types.begin(), types.begin() + b, kBatch) == b &&
+        OrderMatches(input_labels, 0, b, kBatch, other)) {
+      // Batch dims are the leading dims. They have the same order as other.
+      if (CountLabels(types.begin() + b, types.begin() + b + f, kFree) == f) {
+        // All the free dims are placed consecutively after the batch dims.
+        // Their order is arbitrary. The final transpose will ensure that the
+        // output has correct order. We still have to check that the contract
+        // indices have correct order.
+        if (OrderMatches(input_labels, b + f, c, kContract, other)) {
+          layout = EinsumLayout::BFC;
+        }
+      } else if (CountLabels(types.begin() + b, types.begin() + b + c,
+                             kContract) == c) {
+        // All the contract dims are placed consecutively after the batch
+        // dims. Check whether the contract dims have the same order as the
+        // contract dims in other.
+        if (OrderMatches(input_labels, b, c, kContract, other)) {
+          layout = EinsumLayout::BCF;
+        }
+      }
+    }
+    return OkStatus();
+  }
+
+  Status CalculateMixedLayoutPermutation(
+      const EinsumLayout preferred_layout, const Labels& input_labels,
+      const std::vector<EinsumDimensionType>& types,
+      const std::unique_ptr<EinsumDescriptor>& other) {
+    // Input label types are mixed. Calculate a permutation that maps them
+    // to the preferred layout (BCF or BFC).
+    layout = preferred_layout;
+    if (other == nullptr) {
+      AppendMatchingIndicesToPermute(types, kBatch);
+    } else {
+      TF_RETURN_IF_ERROR(
+          FindIndicesoOfAllValuesInSrc(/*values=*/
+                                       absl::MakeConstSpan(
+                                           other->permuted_labels.begin(),
+                                           other->b),
+                                       /*src=*/
+                                       absl::MakeConstSpan(input_labels.begin(),
+                                                           input_labels.size()),
+                                       /*indices=*/&permute));
+    }
+    if (layout == EinsumLayout::BFC) {
+      AppendMatchingIndicesToPermute(types, kFree);
+      if (other == nullptr) {
+        AppendMatchingIndicesToPermute(types, kContract);
+      } else {
+        TF_RETURN_IF_ERROR(FindIndicesoOfAllValuesInSrc(
+            /*values=*/absl::MakeConstSpan(
+                other->permuted_labels.begin() + other->offset_c, other->c),
+            /*src=*/
+            absl::MakeConstSpan(input_labels.begin(), input_labels.size()),
+            /*indices=*/&permute));
+      }
+      return OkStatus();
+    }
+    if (other == nullptr) {
+      AppendMatchingIndicesToPermute(types, kContract);
+    } else {
+      TF_RETURN_IF_ERROR(FindIndicesoOfAllValuesInSrc(
+          /*values=*/absl::MakeConstSpan(
+              other->permuted_labels.begin() + other->offset_c, other->c),
+          /*src=*/absl::MakeConstSpan(input_labels.begin(), input_labels.end()),
+          /*indices=*/&permute));
+    }
+    AppendMatchingIndicesToPermute(types, kFree);
+    return OkStatus();
+  }
+
+  Status Initialize(const TRT_TensorOrWeights& operand, Labels input_labels,
+                    std::vector<EinsumDimensionType>& label_types,
+                    EinsumLayout preferred_layout,
+                    const std::unique_ptr<EinsumDescriptor>& other = nullptr) {
+    if (preferred_layout == EinsumLayout::MIX) {
+      return errors::Internal("Preferred einsum layout cannot be MIX");
+    }
+    // Map label indices to label types.
+    std::vector<EinsumDimensionType> types;  // Input label types.
+    std::transform(input_labels.begin(), input_labels.end(),
+                   std::back_inserter(types),
+                   [&label_types](int i) { return label_types.at(i); });
+
+    b = CountLabels(types.begin(), types.end(), kBatch);
+    f = CountLabels(types.begin(), types.end(), kFree);
+    c = CountLabels(types.begin(), types.end(), kContract);
+
+    if (c == 0 || f == 0) {
+      VLOG(2) << "Einsum equation needs to have at least one free and one "
+                 "contract dimension";
+      return errors::Unimplemented("No conversion for einsum equation.");
+    }
+
+    TF_RETURN_IF_ERROR(DetermineLayout(input_labels, types, other));
+    if (layout == EinsumLayout::MIX) {
+      TF_RETURN_IF_ERROR(CalculateMixedLayoutPermutation(
+          preferred_layout, input_labels, types, other));
+    }
+
+    if (layout == EinsumLayout::BFC) {
+      offset_f = b;
+      offset_c = f + b;
+    } else {
+      offset_f = b + c;
+      offset_c = b;
+    }
+
+    dims = operand.GetTrtDims();
+    for (int i = 0; i < b; i++) {
+      // Set unknown batch dims to zero. These dims will be used in reshape op,
+      // where zero is a special value for retaining the original dim size.
+      if (dims.d[i] == -1) {
+        dims.d[i] = 0;
+      }
+    }
+    permuted_labels = input_labels;
+    if (!permute.empty()) {
+      // Apply the permutation on the dimension array.
+      nvinfer1::Dims orig_dims = dims;
+      for (int i = 0; i < permute.size(); i++) {
+        dims.d[i] = orig_dims.d[permute[i]];
+        permuted_labels[i] = input_labels[permute[i]];
+      }
+    }
+    size_tensors.resize(dims.nbDims, nullptr);
+    return OkStatus();
+  }
+
+ public:
+  EinsumDescriptor() : b(0), f(0), c(0) {}
+
+  // Deduces the number of batch, free, contract dimensions from the input
+  // labels, decides what layout to use, and determines permutation indices for
+  // that layout.
+  static StatusOr<std::unique_ptr<EinsumDescriptor>> Create(
+      const TRT_TensorOrWeights& operand, Labels input_labels,
+      std::vector<EinsumDimensionType>& label_types,
+      EinsumLayout preferred_layout,
+      const std::unique_ptr<EinsumDescriptor>& other = nullptr) {
+    auto desc = std::make_unique<EinsumDescriptor>();
+    TF_RETURN_IF_ERROR(desc->Initialize(operand, input_labels, label_types,
+                                        preferred_layout, other));
+    VLOG(2) << desc->DebugString();
+    return desc;
+  }
+
+  int NumBatchDims() const { return b; }
+  int NumContractDims() const { return c; }
+  int NumFreeDims() const { return f; }
+  int ContractDimOffset() const { return offset_c; }
+  const Labels& PermutedLabels() const { return permuted_labels; }
+
+  std::string DebugString() const {
+    return absl::StrCat("Descriptor with ",
+                        (layout == EinsumLayout::BFC ? "BFC" : "BCF"),
+                        " layout, b=", b, ", f=", f, ", c=", c);
+  }
+
+  // Returns whether the free and contract dimension have static shape.
+  bool HasStaticShape() const {
+    return !std::any_of(dims.d + b, dims.d + dims.nbDims,
+                        [](int k) { return k == -1; });
+  }
+
+  nvinfer1::Permutation GetPermutation() const {
+    nvinfer1::Permutation p;
+    std::copy(permute.begin(), permute.end(), p.order);
+    return p;
+  }
+
+  std::vector<int> PermuteVector() const { return permute; }
+
+  // Sets the "size_tensors" vector to be filled with scalar constant tensors
+  // representing the shape of the operand.
+  Status SetDynamicSize(TRTNetworkBuilder* builder,
+                        const TRT_TensorOrWeights& operand) {
+    TRT_ENSURE(operand.GetTrtDims().nbDims == dims.nbDims);
+    if (operand.is_weights()) {
+      // Generate constants for each dimension of the constant weight tensor's
+      // shape.
+      for (int i = 0; i < operand.GetTrtDims().nbDims; i++) {
+        StatusOr<nvinfer1::IConstantLayer*> size_tensor =
+            builder->Constant<int32_t>(dims.d[i], 1);
+        TRT_ENSURE_PTR_OK(size_tensor);
+        size_tensors[i] = (*size_tensor)->getOutput(0);
+      }
+      return OkStatus();
+    }
+
+    // If the operand is a dynamic tensor, compute the shape value dynamically.
+    StatusOr<nvinfer1::IShapeLayer*> shape_layer =
+        builder->Shape(operand.tensor()->trt_tensor());
+    TRT_ENSURE_PTR_OK(shape_layer);
+    nvinfer1::ITensor* shape = (*shape_layer)->getOutput(0);
+    for (int i = 0; i < operand.GetTrtDims().nbDims; i++) {
+      int idx = permute.empty() ? i : permute.at(i);
+      StatusOr<nvinfer1::ISliceLayer*> slice_layer =
+          builder->Slice(shape, {1, {idx}}, {1, {1}}, {1, {1}});
+      TRT_ENSURE_PTR_OK(slice_layer);
+      size_tensors[i] = (*slice_layer)->getOutput(0);
+    }
+    return OkStatus();
+  }
+
+  EinsumLayout layout;
+  int b;  // number of batch dims
+  int f;  // number of free dims
+  int c;  // number of conraction dims
+  int offset_f;
+  int offset_c;
+  nvinfer1::Dims dims;
+  std::vector<int> permute;
+  std::vector<ITensorProxyPtr> size_tensors;
+  Labels permuted_labels;
+};
+
+// Reshapes operand so that the free dimensions are combined into a single dim,
+// and the contract dimensions are combined into another single dim.
+Status GetEinsumNewDynamicShape(TRTNetworkBuilder* builder,
+                                const EinsumDescriptor& desc,
+                                ITensorProxyPtr* new_shape) {
+  std::vector<nvinfer1::ITensor*> size;
+  size.reserve(desc.b + 2);
+  absl::c_transform(absl::MakeSpan(desc.size_tensors).subspan(0, desc.b + 2),
+                    std::back_inserter(size),
+                    [](const ITensorProxyPtr x) { return x->trt_tensor(); });
+
+  int idx_f = desc.layout == EinsumLayout::BFC ? desc.b : desc.b + 1;
+  int idx_c = desc.layout == EinsumLayout::BFC ? desc.b + 1 : desc.b;
+
+  std::vector<nvinfer1::ITensor*> size_tensors;
+  size_tensors.reserve(desc.size_tensors.size());
+  absl::c_transform(desc.size_tensors, std::back_inserter(size_tensors),
+                    [](const ITensorProxyPtr x) -> nvinfer1::ITensor* {
+                      return x->trt_tensor();
+                    });
+
+  StatusOr<nvinfer1::ILayer*> shape_vol = builder->CumulativeProd(
+      absl::MakeSpan(size_tensors).subspan(desc.offset_f, desc.f));
+  TRT_ENSURE_PTR_OK(shape_vol);
+  size[idx_f] = (*shape_vol)->getOutput(0);
+
+  shape_vol = builder->CumulativeProd(
+      absl::MakeSpan(size_tensors).subspan(desc.offset_c, desc.c));
+  TRT_ENSURE_PTR_OK(shape_vol);
+  size[idx_c] = (*shape_vol)->getOutput(0);
+  StatusOr<nvinfer1::IConcatenationLayer*> layer =
+      builder->Concat(size, /*axis=*/0);
+  TRT_ENSURE_PTR_OK(layer);
+  *new_shape = (*layer)->getOutput(0);
+  return OkStatus();
+}
+
+// Reshapes operand so that the free dimensions are combined into a single dim,
+// and the contract dimensions are combined into another single dim.
+Status GetEinsumNewStaticShape(const EinsumDescriptor& desc,
+                               nvinfer1::Dims* new_dims) {
+  // Copy the batch dims and append two additional dimensions.
+  DimsAdapter adap(
+      absl::MakeSpan(static_cast<const int32_t*>(desc.dims.d), desc.b));
+  adap.Append(1).Append(1);
+
+  // Combine free dims and contract dims.
+  int idx_f = desc.layout == EinsumLayout::BFC ? desc.b : desc.b + 1;
+  int idx_c = desc.layout == EinsumLayout::BFC ? desc.b + 1 : desc.b;
+
+  // Find the volume of the free dimensions.
+  int64_t vol_f =
+      DimsAdapter(
+          absl::MakeSpan(
+              static_cast<const int32_t*>(desc.dims.d) + desc.offset_f, desc.f))
+          .Volume();
+
+  // Find the volume of the contracted dimensions.
+  int64_t vol_c =
+      DimsAdapter(
+          absl::MakeSpan(
+              static_cast<const int32_t*>(desc.dims.d) + desc.offset_c, desc.c))
+          .Volume();
+
+  adap.dim(idx_f) = vol_f;
+  adap.dim(idx_c) = vol_c;
+  *new_dims = adap.AsTrtDims();
+  return OkStatus();
+}
+
+StatusOr<TRT_TensorOrWeights> ConditionEinsumWeights(
+    TRTNetworkBuilder* builder, const TRT_TensorOrWeights& operand,
+    const EinsumDescriptor& desc, const bool need_transpose) {
+  TRT_ENSURE(operand.is_weights());
+  if (!need_transpose) {
+    // If we don't need to transpose, then the operand remains as a weights
+    // constant. In this case we also don't need a reshape.
+    TRT_ShapedWeights weights(operand.weights());
+    nvinfer1::Dims new_dims;
+    TF_RETURN_IF_ERROR(GetEinsumNewStaticShape(desc, &new_dims));
+    TF_RETURN_IF_ERROR(weights.SetShape(new_dims));
+    return TRT_TensorOrWeights(weights);
+  }
+
+  // Let TensorRT handle constant folding where possible.
+  StatusOr<nvinfer1::IConstantLayer*> tensor = builder->WeightsToConstant(
+      operand.weights().GetTrtWeights(), operand.GetTrtDims());
+  TRT_ENSURE_PTR_OK(tensor);
+  return TRT_TensorOrWeights((*tensor)->getOutput(0));
+}
+
+// Builds a TRT shuffle operation for the given operand. Replaces operand with a
+// pointer to the shuffle output.
+Status ConditionEinsumTensor(TRTNetworkBuilder* builder,
+                             std::unique_ptr<TRT_TensorOrWeights>* operand,
+                             const EinsumDescriptor& desc,
+                             const bool need_transpose,
+                             const bool need_reshape) {
+  StatusOr<ShuffleBuilder> shuffle =
+      ShuffleBuilder::Create(builder, (*operand)->tensor()->trt_tensor());
+  TRT_ENSURE_OK(shuffle);
+
+  // Set new shape.
+  if (need_reshape) {
+    if (desc.HasStaticShape()) {
+      nvinfer1::Dims new_dims;
+      TF_RETURN_IF_ERROR(GetEinsumNewStaticShape(desc, &new_dims));
+      shuffle->SetReshape(new_dims);
+    } else {
+      ITensorProxyPtr new_shape;
+      TF_RETURN_IF_ERROR(GetEinsumNewDynamicShape(&*builder, desc, &new_shape));
+      shuffle->SetReshape(new_shape->trt_tensor());
+    }
+  }
+
+  if (need_transpose) {
+    shuffle->SetFirstTranspose(desc.GetPermutation());
+  }
+
+  StatusOr<nvinfer1::ITensor*> shuffle_out = shuffle->Output();
+  TRT_ENSURE_PTR_OK(shuffle_out);
+  *operand = std::make_unique<TRT_TensorOrWeights>(*shuffle_out);
+  return OkStatus();
+}
+
+// Handles einsum operand conditioning for both constant and non-constant
+// inputs. This is supported using the ShuffleEinsumWeights and
+// ShuffleEinsumTensor routines.
+Status ConditionEinsumOperand(TRTNetworkBuilder* builder,
+                              std::unique_ptr<TRT_TensorOrWeights>* operand,
+                              const EinsumDescriptor& desc) {
+  bool need_reshape = (desc.f != 1 || desc.c != 1);
+  bool need_transpose = !desc.permute.empty();
+
+  VLOG(2) << "Condition operand. Need reshape: " << need_reshape
+          << ". Need transpose: " << need_transpose;
+
+  if ((*operand)->is_weights()) {
+    StatusOr<TRT_TensorOrWeights> result =
+        ConditionEinsumWeights(builder, **operand, desc, need_transpose);
+    TRT_ENSURE_OK(result);
+    *operand = std::make_unique<TRT_TensorOrWeights>(std::move(result).value());
+  }
+
+  // If we didn't convert the operand to a tensor, we can return here.
+  if ((*operand)->is_weights()) {
+    return OkStatus();
+  }
+
+  TF_RETURN_IF_ERROR(ConditionEinsumTensor(builder, operand, desc,
+                                           need_transpose, need_reshape));
+
+  return OkStatus();
+}
+
+// Combines output dims/labels by copying batch and free dims/labels from input
+// A, and concatenating free values from input B.
+template <typename InputIterator, typename OutputIterator>
+void AssembleOutput(InputIterator begin_a, InputIterator begin_b,
+                    const EinsumDescriptor& desc_a,
+                    const EinsumDescriptor& desc_b, OutputIterator out) {
+  std::copy(begin_a, begin_a + desc_a.b, out);
+  begin_a += desc_a.offset_f;
+  std::copy(begin_a, begin_a + desc_a.f, out + desc_a.b);
+  begin_b += desc_b.offset_f;
+  std::copy(begin_b, begin_b + desc_b.f, out + desc_a.b + desc_a.f);
+}
+
+// Restores free dimensions and sets final index order. Consider C = A * B,
+// batched MatMul op, where A.shape = [B, x, k] and B.shape = [B, k, y]. Then
+// C.shape = [B, x, y]. Here B can denote multiple batch indices while x, y, k
+// are single indices. The original inputs to Einsum can have multiple free
+// indices. These were combined into a singe free dimension x and y, for example
+// x = f_a1 * f_a2 * f_a3, y = f_b1 * f_b2. This routine creates a shuffle layer
+// to expand x into and y the original free dims, e.g. C is reshaped to
+// [B, f_a1, f_a2, f_a3, f_b1, f_b2]. Finally, a permutation is applied to
+// transform the shape to the shape of the original Einsum output.
+Status ShuffleEinsumOutput(const OpConverterParams* params,
+                           EinsumDescriptor desc_a, EinsumDescriptor desc_b,
+                           const std::vector<int>& permutation,
+                           ITensorProxyPtr* output) {
+  if (permutation.empty() && (desc_a.f == 1 && desc_b.f == 1)) {
+    return OkStatus();
+  }
+
+  nvinfer1::IShuffleLayer* layer =
+      params->converter->network()->addShuffle(*(*output)->trt_tensor());
+  TFTRT_RETURN_ERROR_IF_NULLPTR(layer, params->node_def.name());
+  params->converter->SetLayerName(layer, params->node_def, "shuffle",
+                                  /*sub_op_instance=*/2);
+
+  int output_rank = desc_a.b + desc_a.f + desc_b.f;
+  if (desc_a.f != 1 || desc_b.f != 1) {
+    if (desc_a.HasStaticShape() && desc_b.HasStaticShape()) {
+      nvinfer1::Dims dims_out = {output_rank, {}};
+      AssembleOutput(desc_a.dims.d, desc_b.dims.d, desc_a, desc_b, dims_out.d);
+      layer->setReshapeDimensions(dims_out);
+    } else {
+      std::vector<ITensorProxyPtr> size_tensors(output_rank);
+      AssembleOutput(desc_a.size_tensors.begin(), desc_b.size_tensors.begin(),
+                     desc_a, desc_b, size_tensors.begin());
+      ITensorProxyPtr new_shape;
+      auto builder = TRTNetworkBuilder::Create(params->converter->network(),
+                                               params->weight_store);
+      TRT_ENSURE_OK(builder);
+      std::vector<nvinfer1::ITensor*> size_itensors;
+      absl::c_transform(size_tensors, std::back_inserter(size_itensors),
+                        [](auto x) { return x->trt_tensor(); });
+      StatusOr<nvinfer1::IConcatenationLayer*> concat =
+          builder->Concat(size_itensors, /*axis=*/0);
+      TRT_ENSURE_PTR_OK(concat);
+      new_shape = (*concat)->getOutput(0);
+      layer->setInput(1, *new_shape->trt_tensor());
+    }
+  }
+
+  if (!permutation.empty()) {
+    nvinfer1::Permutation p;
+    std::copy(permutation.begin(), permutation.end(), p.order);
+    layer->setSecondTranspose(p);
+  }
+  *output = layer->getOutput(0);
+  return OkStatus();
+}
+
+// Updates "final_transpose" according to the given descriptors and output
+// labels.
+StatusOr<std::vector<int>> GetOutputTranspose(
+    const EinsumDescriptor& descriptor_a, const EinsumDescriptor& descriptor_b,
+    Labels output_labels) {
+  // Get final transpose.
+  std::vector<int> final_transpose;
+  final_transpose.reserve(descriptor_a.b + descriptor_a.f + descriptor_b.f);
+  Labels matmul_output_labels(descriptor_a.b + descriptor_a.f + descriptor_b.f);
+  AssembleOutput(descriptor_a.permuted_labels.begin(),
+                 descriptor_b.permuted_labels.begin(), descriptor_a,
+                 descriptor_b, matmul_output_labels.begin());
+  TF_RETURN_IF_ERROR(
+      FindIndicesoOfAllValuesInSrc(/*values=*/
+                                   absl::MakeConstSpan(output_labels.begin(),
+                                                       output_labels.end()),
+                                   /*src=*/
+                                   absl::MakeConstSpan(
+                                       matmul_output_labels.begin(),
+                                       matmul_output_labels.end()),
+                                   /*indices=*/&final_transpose));
+  // Clear identity transpose.
+  bool identity_transpose = true;
+  for (int i = 0; i < final_transpose.size() && identity_transpose; i++) {
+    identity_transpose &= final_transpose.at(i) == i;
+  }
+  if (identity_transpose) {
+    final_transpose.clear();
+  }
+  return final_transpose;
+}
+
+// Prepares EinsumDescriptors after parsing the equation and determines the
+// final transpose.
+Status ParseEquation(const std::string& equation,
+                     std::unique_ptr<TRT_TensorOrWeights>* input_a,
+                     std::unique_ptr<TRT_TensorOrWeights>* input_b,
+                     std::unique_ptr<EinsumDescriptor>* descriptor_a,
+                     std::unique_ptr<EinsumDescriptor>* descriptor_b,
+                     std::vector<int>* final_transpose) {
+  VLOG(2) << "Einsum equation " << equation;
+  OperandLabels input_labels;
+  Labels output_labels;
+  std::vector<EinsumDimensionType> label_types;
+  OperandLabelCounts input_label_counts;
+  LabelCounts output_label_counts;
+  absl::InlinedVector<bool, 2> input_has_ellipsis;
+  bool output_has_ellipsis;
+  TF_RETURN_IF_ERROR(
+      ParseEinsumEquation(equation, &input_labels, &output_labels, &label_types,
+                          &input_label_counts, &output_label_counts,
+                          &input_has_ellipsis, &output_has_ellipsis));
+
+  if (input_has_ellipsis[0] || input_has_ellipsis[1] || output_has_ellipsis) {
+    // TODO(tfeher): Handle ellipsis like EinsumHelper::ProcessDimensions.
+    // Note: ProcessDimensions would introduce kBroadcasting labels, which we
+    // need to replace with kBatch before we call InitDescriptor.
+    VLOG(2) << "Ellipsis not yet supported";
+    return errors::Unimplemented("No conversion for einsum equation.");
+  }
+
+  if (absl::c_any_of(label_types, [](auto l) {
+        return l == EinsumDimensionType::kReduce ||
+               l == EinsumDimensionType::kBroadcasting;
+      })) {
+    VLOG(2) << "Einsum reductions not implemented";
+    return errors::Unimplemented("No conversion for einsum equation.");
+  }
+
+  auto no_duplicated_labels = [](const LabelCounts& label_counts) {
+    return absl::c_any_of(label_counts, [](int i) { return i > 1; });
+  };
+  if (no_duplicated_labels(input_label_counts[0]) ||
+      no_duplicated_labels(input_label_counts[1]) ||
+      no_duplicated_labels(output_label_counts)) {
+    VLOG(2) << "Einsum invalid label count";
+    return errors::Unimplemented("No conversion for einsum equation.");
+  }
+
+  if ((*input_a)->is_weights() && (*input_b)->is_tensor()) {
+    // We prefer to use FC layer, needs A as tensor and B as weight.
+    std::swap(*input_a, *input_b);
+    std::swap(input_labels[0], input_labels[1]);
+    std::swap(input_label_counts[0], input_label_counts[1]);
+  }
+
+  auto desc = EinsumDescriptor::Create(**input_a, input_labels[0], label_types,
+                                       EinsumLayout::BFC);
+  TF_RETURN_IF_ERROR(desc.status());
+  *descriptor_a = std::move(desc).value();
+
+  desc = EinsumDescriptor::Create(**input_b, input_labels[1], label_types,
+                                  EinsumLayout::BCF, *descriptor_a);
+  TF_RETURN_IF_ERROR(desc.status());
+  *descriptor_b = std::move(desc).value();
+
+  auto out_transpose =
+      GetOutputTranspose(**descriptor_a, **descriptor_b, output_labels);
+
+  TRT_ENSURE_OK(out_transpose)
+  *final_transpose = std::move(out_transpose).value();
+  return OkStatus();
+}
+
+class ConvertEinsum : public OpConverterBase<ConvertEinsum> {
+ public:
+  explicit ConvertEinsum(const OpConverterParams* params)
+      : OpConverterBase<ConvertEinsum>(params) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return {InputArgSpec::Create("input_a", TrtInputArg::kBoth),
+            InputArgSpec::Create("input_b", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    const auto& inputs = params_->inputs;
+    input_a = std::make_unique<TRT_TensorOrWeights>(inputs.at(0));
+    input_b = std::make_unique<TRT_TensorOrWeights>(inputs.at(1));
+
+    StatusOr<std::string> eq = GetAttrValue<std::string>("equation");
+    TRT_ENSURE_OK(eq);
+    TF_RETURN_IF_ERROR(ParseEquation(*eq, &input_a, &input_b, &descriptor_a,
+                                     &descriptor_b, &final_transpose));
+    return OkStatus();
+  }
+
+  Status Convert() {
+    auto builder = TRTNetworkBuilder::Create(params_->converter->network(),
+                                             params_->weight_store);
+    TRT_ENSURE_OK(builder);
+    TRT_ENSURE(input_a && input_b);
+    TRT_ENSURE(descriptor_a && descriptor_b);
+
+    // Populate the size_tensor vector in the descriptor.
+    TF_RETURN_IF_ERROR(descriptor_a->SetDynamicSize(&*builder, *input_a));
+    TF_RETURN_IF_ERROR(descriptor_b->SetDynamicSize(&*builder, *input_b));
+
+    // Condition the operands for lowering to matmul.
+    TF_RETURN_IF_ERROR(
+        ConditionEinsumOperand(&*builder, &input_a, *descriptor_a));
+    TF_RETURN_IF_ERROR(
+        ConditionEinsumOperand(&*builder, &input_b, *descriptor_b));
+
+    // Build the matmul implementation.
+    StatusOr<ITensorProxyPtr> result = ConvertMatMulImpl(
+        params_, *input_a, *input_b, descriptor_a->layout == EinsumLayout::BCF,
+        descriptor_b->layout == EinsumLayout::BFC);
+    TF_RETURN_IF_ERROR(result.status());
+    ITensorProxyPtr output = result.value();
+
+    // Reshape and permute the output.
+    TF_RETURN_IF_ERROR(ShuffleEinsumOutput(
+        params_, *descriptor_a, *descriptor_b, final_transpose, &output));
+    this->AddOutput(output);
+    return OkStatus();
+  }
+
+ private:
+  std::unique_ptr<TRT_TensorOrWeights> input_a{nullptr};
+  std::unique_ptr<TRT_TensorOrWeights> input_b{nullptr};
+  std::vector<int> final_transpose;
+  std::unique_ptr<EinsumDescriptor> descriptor_a{nullptr};
+  std::unique_ptr<EinsumDescriptor> descriptor_b{nullptr};
+};
+#else
+
+// Helper class to reindex equations to contain only lowercase characters. We
+// simply define a mapping from the old character set to a new set.
+// - The input is assumed to be a valid TF equation.
+// - The input is TRT compatible, therefore it has max 8 dims. (Thus we have
+//   enough lowercase English characters to represent the equation.)
+// How do we reindex/map equations:
+// - Only uppercase letters are changed, if possible we just lowercase them.
+// - If the equation contains both upper and lowercase variant of a letter, say
+//   X and x, then we map X to the first unused lowercase letter.
+class ReIndexer {
+ public:
+  // Initializes the index map with existing lowercase labels.
+  ReIndexer(std::string eq) {
+    for (char c : eq) {
+      if (absl::ascii_islower(c)) {
+        idx_map_[c] = c;
+      }
+    }
+  }
+  // Finds new character for uppercase character c.
+  char operator()(char c) {
+    if (!absl::ascii_isupper(c)) return c;
+    if (idx_map_.count(c) > 0) return idx_map_[c];
+    char new_idx = absl::ascii_tolower(c);
+
+    // If lower(c) is not used in the equation, use it to replace c.
+    if (idx_map_.count(new_idx) == 0) {
+      idx_map_[c] = new_idx;
+      idx_map_[new_idx] = new_idx;  // mark that new_idx is taken
+      return new_idx;
+    }
+
+    // Otherwise, find the first available lower case to replace c.
+    for (char k = 'a'; k <= 'z'; k++) {
+      if (idx_map_.count(k) == 0) {
+        new_idx = k;
+        idx_map_[c] = new_idx;
+        idx_map_[new_idx] = new_idx;  // mark that new_idx is taken
+        break;
+      }
+    }
+    return new_idx;
+  }
+
+ private:
+  // Each key is an index used in the original or in the reindexed equation.
+  // The values are the corresponding new lowercase indices.
+  std::map<char, char> idx_map_;
+};
+
+class ConvertEinsum : public OpConverterBase<ConvertEinsum> {
+ public:
+  explicit ConvertEinsum(const OpConverterParams* params)
+      : OpConverterBase<ConvertEinsum>(params) {}
+
+  Status ValidateInputs() {
+    TRT_ENSURE(params_->inputs.size() <= 2);
+    return OkStatus();
+  }
+  static constexpr bool HasFixNumberOfInputs() { return false; }
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return {InputArgSpec::Create("input_a", TrtInputArg::kBoth),
+            InputArgSpec::Create("input_b", TrtInputArg::kBoth)};
+  }
+
+  std::string MakeLowerCase(const std::string& eq) {
+    std::string res = eq;
+    ReIndexer reindexer(eq);
+    std::transform(eq.begin(), eq.end(), res.begin(), reindexer);
+    return res;
+  }
+
+  // Checks if the equation is supported by TRT.
+  Status ValidateEinsumEquation(const std::string& eq) {
+    const auto& inputs = params_->inputs;
+    OperandLabels input_labels;
+    Labels output_labels;
+    std::vector<EinsumDimensionType> label_types;
+    OperandLabelCounts input_label_counts;
+    LabelCounts output_label_counts;
+    absl::InlinedVector<bool, 2> input_has_ellipsis;
+    bool output_has_ellipsis;
+    VLOG(2) << "Parsing equation " << eq;
+    TF_RETURN_IF_ERROR(ParseEinsumEquation(
+        eq, &input_labels, &output_labels, &label_types, &input_label_counts,
+        &output_label_counts, &input_has_ellipsis, &output_has_ellipsis));
+
+    Status unimplemented =
+        errors::Unimplemented("No conversion for einsum equation.");
+    if (input_has_ellipsis[0] || (inputs.size() > 1 && input_has_ellipsis[1]) ||
+        output_has_ellipsis) {
+      VLOG(2) << "Ellipsis not yet supported";
+      return unimplemented;
+    }
+    for (int i = 0; i < input_label_counts.size(); i++) {
+      for (int k = 0; k < input_label_counts[i].size(); k++) {
+        if (input_label_counts[i][k] > 1) {
+          VLOG(2) << "Diagonal operation or reduction not yet supported";
+          return unimplemented;
+        }
+      }
+    }
+    bool has_out_idx =
+        std::reduce(output_label_counts.begin(), output_label_counts.end(),
+                    false, std::logical_or<int>());
+    if (!has_out_idx) {
+      VLOG(2) << "Scalar output not allowed in dynamic shape mode";
+      return unimplemented;
+    }
+    // Check for outer product
+    if (input_label_counts.size() == 2 && output_label_counts.size() == 2 &&
+        output_label_counts[0] == 1 && output_label_counts[1] == 1) {
+      VLOG(2) << "Outer product not supported";
+      return unimplemented;
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    VLOG(2) << "Running validation using the new einsum "
+               "converter";
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    StatusOr<std::string> eq = GetAttrValue<std::string>("equation");
+    TRT_ENSURE_OK(eq);
+
+    TF_RETURN_IF_ERROR(ValidateEinsumEquation(*eq));
+
+    // While TF has case sensitive equations, TensorRT expects lowercase eq (as
+    // of version 8.4). See
+    // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#einsum-layer
+    equation = MakeLowerCase(*eq);
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    auto builder = TRTNetworkBuilder::Create(params_->converter->network(),
+                                             params_->weight_store);
+    TRT_ENSURE_OK(builder);
+
+    std::vector<nvinfer1::ITensor*> trt_input;
+    for (const TRT_TensorOrWeights& input_arg : params_->inputs) {
+      ITensorProxyPtr ptr = nullptr;
+      if (input_arg.is_tensor()) {
+        ptr = input_arg.tensor();
+      } else {
+        StatusOr<nvinfer1::IConstantLayer*> const_layer =
+            builder->WeightsToConstant(input_arg.weights().GetTrtWeights(),
+                                       input_arg.GetTrtDims());
+        TRT_ENSURE_PTR_OK(const_layer);
+        ptr = (*const_layer)->getOutput(0);
+      }
+      trt_input.push_back(ptr->trt_tensor());
+    }
+    nvinfer1::IEinsumLayer* layer = params_->converter->network()->addEinsum(
+        trt_input.data(), trt_input.size(), equation.c_str());
+    TRT_ENSURE(layer);
+
+    ITensorProxyPtr output = layer->getOutput(0);
+    this->AddOutput(output);
+    return OkStatus();
+  }
+
+ private:
+  std::string equation;
+};
+
+#endif
+
+}  // namespace
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertEinsum>(),
+                                  "Einsum");
+#endif
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/fill_ops.cc

@@ -0,0 +1,313 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+template <typename Impl>
+class ConvertFillBase : public OpConverterBase<Impl> {
+ public:
+  explicit ConvertFillBase(const OpConverterParams* params)
+      : OpConverterBase<Impl>(params, {DataType::DT_FLOAT, DataType::DT_HALF,
+                                       DataType::DT_INT32}) {}
+};
+
+class ConvertFill : public ConvertFillBase<ConvertFill> {
+ public:
+  explicit ConvertFill(const OpConverterParams* params)
+      : ConvertFillBase<ConvertFill>(params) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("dims", TrtInputArg::kBoth),
+        InputArgSpec::Create("value", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    const auto& params = *this->params_;
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+
+    const auto& inputs = params.inputs;
+    const auto& node_def = params.node_def;
+    const TRT_TensorOrWeights& dims_input = inputs.at(0);
+
+    const auto dims_type = dims_input.TrtDType();
+    if (dims_type != nvinfer1::DataType::kINT32) {
+      return errors::InvalidArgument("The dims parameter of ", node_def.op(),
+                                     " operation in ", node_def.name(),
+                                     " is expected to be of type ",
+                                     DebugString(nvinfer1::DataType::kINT32),
+                                     " type, got ", DebugString(dims_type));
+    }
+
+    const auto nbDims = dims_input.GetTrtDims().nbDims;
+    if (nbDims < 0) {
+      return errors::InvalidArgument("The shape of parameter ", node_def.op(),
+                                     " operation in ", node_def.name(),
+                                     " cannot be partial.");
+    }
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& params = *this->params_;
+    auto* network = params.converter->network();
+    const auto& inputs = params.inputs;
+
+    const bool is_dims_static = inputs[0].is_weights();
+    const bool is_value_static = inputs[1].is_weights();
+
+    const TRT_TensorOrWeights& dims_input = inputs.at(0);
+    const TRT_TensorOrWeights& value_input = inputs.at(1);
+
+    int nbDims = dims_input.GetTrtDims().d[0];
+
+    nvinfer1::Dims trt_dims{0};
+    if (is_dims_static) {
+      const auto dims_weights = dims_input.weights();
+      DimsAdapter dims_adapter(dims_weights.GetSpan<int32>());
+      dims_adapter.TrtDims(&trt_dims);
+    }
+
+    auto builder = TRTNetworkBuilder::Create(network, params.weight_store);
+    StatusOr<nvinfer1::ILayer*> layer =
+        builder->AddFill(value_input, dims_input, is_value_static,
+                         is_dims_static, nbDims, trt_dims);
+    ITensorProxyPtr output_tensor = (*layer)->getOutput(0);
+    this->AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+class ConvertRange : public ConvertFillBase<ConvertRange> {
+ public:
+  explicit ConvertRange(const OpConverterParams* params)
+      : ConvertFillBase<ConvertRange>(params) {}
+
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return std::array<InputArgSpec, 3>{
+        InputArgSpec::Create("start", TrtInputArg::kBoth),
+        InputArgSpec::Create("limit", TrtInputArg::kBoth),
+        InputArgSpec::Create("delta", TrtInputArg::kBoth)};
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "Range"
+      ...
+      attr {
+        key: "Tidx"
+        value {
+          type: DT_INT32
+        }
+      }
+    }
+    */
+    return "Tidx";
+  }
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    const auto& params = *this->params_;
+    const auto& inputs = params.inputs;
+    const auto& node_def = params.node_def;
+
+    float param[3];
+    all_weights_ = all_integers_ = true;
+    for (int i = 0; i < 3; i++) {
+      const auto& input = inputs.at(i);
+      all_integers_ &= input.TrtDType() == nvinfer1::DataType::kINT32;
+      if (input.is_weights()) {
+        switch (input.TrtDType()) {
+          case nvinfer1::DataType::kFLOAT:
+            param[i] = get_input_param<float>(input);
+            break;
+          case nvinfer1::DataType::kHALF:
+            param[i] = get_input_param<Eigen::half>(input);
+            break;
+          case nvinfer1::DataType::kINT32:
+            param[i] = get_input_param<int>(input);
+            break;
+          default:
+            return errors::InvalidArgument(
+                "Unsupported data type ", DebugString(input.TrtDType()),
+                " used for '", InputSpec()[i].name, "'");
+        }
+      } else {
+        all_weights_ = false;
+      }
+    }
+
+    if (!(all_weights_ || all_integers_)) {
+      // As of 06/03/2022, when at least one of the (start, limit, delta)
+      // is passed as a tensor, they must all be of type kINT32
+      return errors::Unimplemented(convert_range_expected_msg(node_def));
+    }
+
+    if (inputs.at(2).is_weights()) {
+      if ((delta_ = param[2]) == 0) {
+        return errors::InvalidArgument("The delta parameter of ", node_def.op(),
+                                       " operation cannot be equal to 0");
+      }
+
+      if (!all_weights_ && delta_ < 0) {
+        return errors::InvalidArgument(
+            "The delta parameter of Range operation "
+            "cannot be negative, when one of (start, limit) is passed as "
+            "a tensor, but got ",
+            delta_);
+      }
+    }
+
+    for (int i = 0; i < 3; i++) {
+      const auto& input = inputs.at(i);
+      const auto& dims = input.GetTrtDims();
+      if (dims.nbDims != 1 || dims.d[0] != 1) {
+        return errors::InvalidArgument("Dimension for '", InputSpec()[i].name,
+                                       "' of ", node_def.op(), " operator ",
+                                       "should be equal to 1");
+      }
+    }
+
+    if (all_weights_) {
+      const auto num_intervals_float =
+          (param[1] - (start_ = param[0])) / delta_;
+      if (num_intervals_float < 0) {
+        const auto error = convert_range_error_msg(start_, param[1], delta_);
+        return errors::InvalidArgument(error);
+      }
+
+      num_values_ = static_cast<int>(num_intervals_float);
+      if (start_ + delta_ * num_values_ != param[1]) {
+        num_values_++;
+      }
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& params = *this->params_;
+    const auto& inputs = params.inputs;
+    const TRT_TensorOrWeights& input = inputs.at(0);
+    TRT_TensorOrWeights value_input;
+    nvinfer1::Dims trt_dims{1};
+    auto builder = TRTNetworkBuilder::Create(params.converter->network(),
+                                             params.weight_store);
+    TRT_ENSURE_OK(builder);
+    ITensorProxyPtr dims_input_tensor = nullptr;
+    ITensorProxyPtr beta_tensor = nullptr;
+    ITensorProxyPtr scalar_tensor = nullptr;
+    if (!all_weights_) {
+      ITensorProxyPtr tensors[3];
+      for (int i = 0; i < 3; i++) {
+        TF_RETURN_IF_ERROR(
+            builder->get_tensor4TensorOrWeights(inputs.at(i), tensors + i));
+      }
+
+      StatusOr<nvinfer1::IElementWiseLayer*> num =
+          builder->Sub(/*limit*/ tensors[1]->trt_tensor(),
+                       /*start*/ tensors[0]->trt_tensor());
+
+      TRT_ENSURE_PTR_OK(num);
+      StatusOr<nvinfer1::IElementWiseLayer*> ceil_div = builder->FloorDiv(
+          (*num)->getOutput(0), (beta_tensor = tensors[2])->trt_tensor());
+      TRT_ENSURE_PTR_OK(ceil_div);
+      dims_input_tensor = (*ceil_div)->getOutput(0);
+      dims_input_tensor->setType(nvinfer1::DataType::kINT32);
+
+      nvinfer1::Dims scalar_dims{0};
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params.converter, params.inputs.at(0), scalar_dims, false,
+          &scalar_tensor, params.node_def));
+    } else {
+      DimsAdapter value_input_dims(std::vector<int>{1});
+      StatusOr<TRT_ShapedWeights> value_weights =
+          params.weight_store->GetTempWeights(input.TrtDType(),
+                                              value_input_dims);
+
+      TF_RETURN_IF_ERROR(value_weights.status());
+      TF_RETURN_IF_ERROR(value_weights->SetValues(start_));
+      value_input = TRT_TensorOrWeights(value_weights.value());
+
+      trt_dims.d[0] = num_values_;
+      StatusOr<nvinfer1::IConstantLayer*> const_layer =
+          builder->ConstantShape(value_input_dims);
+      TRT_ENSURE_PTR_OK(const_layer);
+      dims_input_tensor = (*const_layer)->getOutput(0);
+    }
+
+    TRT_TensorOrWeights dims_input(dims_input_tensor);
+
+    StatusOr<nvinfer1::ILayer*> layer =
+        builder->AddFill(value_input, dims_input, all_weights_, all_weights_, 1,
+                         trt_dims, scalar_tensor, beta_tensor, delta_);
+
+    ITensorProxyPtr output_tensor = (*layer)->getOutput(0);
+    if (all_integers_) {
+      output_tensor->setType(nvinfer1::DataType::kINT32);
+    }
+
+    this->AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+
+ private:
+  template <typename T>
+  float get_input_param(const TRT_TensorOrWeights& input) {
+    return static_cast<float>(*input.weights().GetPointer<T>());
+  }
+
+  float start_;
+  float delta_;
+  int num_values_;
+  bool all_weights_;
+  bool all_integers_;
+};
+
+std::string convert_range_error_msg(float start, float limit, float delta) {
+  constexpr const char* format_string =
+      "For parameters (start, limit) = (%.2f, %.2f) "
+      "of the Range operation delta cannot be %s, got %.2f";
+  return absl::StrFormat(format_string, start, limit,
+                         start < limit ? "negative" : "positive", delta);
+}
+
+std::string convert_range_expected_msg(const NodeDef& node_def) {
+  return "When at least one of parameters (start, limit, delta) of " +
+         node_def.op() + " operation in " + node_def.name() +
+         " is passed as a tensor, they must all be of type kINT32";
+}
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertFill>(), "Fill");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertRange>(),
+                                  "Range");
+
+#endif  // IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h

@@ -0,0 +1,715 @@
+/* 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_CONVERT_OPS_LAYER_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_LAYER_UTILS_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <type_traits>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "third_party/tensorrt/NvInfer.h"
+#include "third_party/tensorrt/NvInferRuntimeCommon.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+namespace convert {
+
+// Facilitates the creation of TensorRT layers inside a network. The user
+// provides a INetworkDefinition pointer during construction. They can then add
+// operations to the network through the provided functions. Each function
+// returns a struct which contains the symbolic result of the operation (ITensor
+// pointer) as well as a pointer to the last TensorRT ILayer created. Some
+// operations may create multiple layers in order to accomplish the desired
+// result (e.g. Sign).
+class TRTNetworkBuilder {
+ public:
+  static StatusOr<TRTNetworkBuilder> Create(
+      nvinfer1::INetworkDefinition* network, TrtWeightStore* weight_store) {
+    TRT_ENSURE(network);
+    TRT_ENSURE(weight_store);
+    return TRTNetworkBuilder(network, weight_store);
+  }
+
+ private:
+  TRTNetworkBuilder(nvinfer1::INetworkDefinition* network,
+                    TrtWeightStore* weight_store)
+      : network_(network), weight_store_(weight_store) {}
+
+ public:
+  // Adds an Add operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Add(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kSUM);
+    TRT_ENSURE(layer);
+    return layer;
+  };
+
+  // Adds an elementwise min(lhs, rhs) operation to the network. The output has
+  // the same data type as the input.
+  StatusOr<nvinfer1::IElementWiseLayer*> Min(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kMIN);
+    TRT_ENSURE(layer);
+    return layer;
+  };
+
+  // Adds an elementwise max(lhs, rhs) operation to the network. The output has
+  // the same datatype as the input.
+  StatusOr<nvinfer1::IElementWiseLayer*> Max(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kMAX);
+    TRT_ENSURE(layer);
+    return layer;
+  };
+
+  // Adds an absolute value operation to the network. Note that this unary
+  // operation will do an implicit float conversion. For int32 tensors, use
+  // "AbsInt".
+  StatusOr<nvinfer1::IUnaryLayer*> AbsFloat(nvinfer1::ITensor* input) noexcept {
+    TRT_ENSURE(input);
+    TRT_ENSURE(input->getType() != nvinfer1::DataType::kFLOAT &&
+               input->getType() != nvinfer1::DataType::kHALF);
+    nvinfer1::IUnaryLayer* layer =
+        network_->addUnary(*input, nvinfer1::UnaryOperation::kABS);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Performs Abs without implicit float conversion. The input should be of type
+  // kInt32. For float datatypes, use "Abs".
+  StatusOr<nvinfer1::IElementWiseLayer*> AbsInt(
+      nvinfer1::ITensor* input) noexcept {
+    TRT_ENSURE(input);
+    TRT_ENSURE(input->getType() == nvinfer1::DataType::kINT32);
+    StatusOr<nvinfer1::IElementWiseLayer*> sign = this->SignInt(input);
+    return this->Mul(input, (*sign)->getOutput(0));
+  }
+
+  // Returns elementwise sign(x) for int32 input tensors where sign(x) is
+  // defined as 1 where x > 0, -1 where x < 0 and 0 where x == 0.
+  StatusOr<nvinfer1::IElementWiseLayer*> SignInt(
+      nvinfer1::ITensor* input) noexcept {
+    TRT_ENSURE(input);
+
+    // Create constants +1 and -1.
+    StatusOr<nvinfer1::IConstantLayer*> one =
+        this->Constant<int32>(1, input->getDimensions().nbDims);
+    TRT_ENSURE_PTR_OK(one);
+
+    StatusOr<nvinfer1::IConstantLayer*> neg_one =
+        this->Constant<int32>(-1, input->getDimensions().nbDims);
+    TRT_ENSURE_PTR_OK(neg_one);
+
+    // Turn all negaitve elements into -1, positive and zero elements
+    // unaffected.
+    StatusOr<nvinfer1::IElementWiseLayer*> max =
+        this->Max(input, (*neg_one)->getOutput(0));
+    TRT_ENSURE_PTR_OK(max);
+
+    // Turn all positive elements into +1, negative and zero elements
+    // unaffected.
+    StatusOr<nvinfer1::IElementWiseLayer*> min =
+        this->Min((*max)->getOutput(0), (*one)->getOutput(0));
+    TRT_ENSURE_PTR_OK(min);
+    return min;
+  }
+
+  // Adds a Sub operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Sub(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kSUB);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds an Greater operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Greater(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kGREATER);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds an Equal operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Equal(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kEQUAL);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds a FloorDiv operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> FloorDiv(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kFLOOR_DIV);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Returns the equivalent of ceil_divide(abs(x)/abs(y))) operation. The inputs
+  // "lhs" and "rhs" should be int32 tensors.
+  StatusOr<nvinfer1::IElementWiseLayer*> AbsCeilDivInt(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    TRT_ENSURE(lhs->getType() == nvinfer1::DataType::kINT32);
+    TRT_ENSURE(rhs->getType() == nvinfer1::DataType::kINT32);
+
+    StatusOr<nvinfer1::IElementWiseLayer*> rhs_abs = this->AbsInt(rhs);
+    TRT_ENSURE_PTR_OK(rhs_abs);
+    StatusOr<nvinfer1::IElementWiseLayer*> lhs_abs = this->AbsInt(lhs);
+    TRT_ENSURE_PTR_OK(lhs_abs);
+    StatusOr<nvinfer1::IElementWiseLayer*> add1 =
+        this->Add((*lhs_abs)->getOutput(0), (*rhs_abs)->getOutput(0));
+    TRT_ENSURE_PTR_OK(add1);
+    StatusOr<nvinfer1::IConstantLayer*> one_const =
+        this->Constant<int32>(1, rhs->getDimensions().nbDims);
+    TRT_ENSURE_PTR_OK(one_const);
+    StatusOr<nvinfer1::IElementWiseLayer*> numerator =
+        this->Sub((*add1)->getOutput(0), (*one_const)->getOutput(0));
+    TRT_ENSURE_PTR_OK(numerator);
+    return FloorDiv((*numerator)->getOutput(0), (*rhs_abs)->getOutput(0));
+  }
+
+  // Adds an elementwise multiplication operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Mul(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kPROD);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds a sequence of elementwise multiplication operations to the network.
+  // The returned layer's output contains the cumulative elementwise product of
+  // all tensors in the input.
+  StatusOr<nvinfer1::ILayer*> CumulativeProd(
+      absl::Span<nvinfer1::ITensor*> inputs) noexcept {
+    TRT_ENSURE(!absl::c_any_of(
+        inputs, [](nvinfer1::ITensor* x) { return x == nullptr; }));
+    nvinfer1::ILayer* out = nullptr;
+    if (inputs.size() == 1) {
+      out = network_->addIdentity(*inputs[0]);
+      TRT_ENSURE(out != nullptr);
+      return out;
+    }
+    nvinfer1::ITensor* last = inputs[0];
+    for (int i = 1; i < inputs.size(); i++) {
+      StatusOr<nvinfer1::IElementWiseLayer*> mul = this->Mul(last, inputs[i]);
+      TRT_ENSURE_PTR_OK(mul);
+      out = *mul;
+      last = (*mul)->getOutput(0);
+    }
+    return out;
+  }
+
+  // Adds a Constant layer whose output is a TensorRT shape tensor. The shape
+  // tensor's size and values correspond to dim's nbDims and d[], respectively.
+  StatusOr<nvinfer1::IConstantLayer*> ConstantShape(
+      const DimsAdapter& shape_data) noexcept {
+    TRT_ENSURE(shape_data.NumDims() > 0);
+    nvinfer1::Dims shape_dims;
+    shape_dims.nbDims = 1;
+    shape_dims.d[0] = shape_data.NumDims();
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(nvinfer1::DataType::kINT32, shape_dims);
+    TRT_ENSURE_OK(const_weights);
+    absl::c_copy(shape_data, const_weights->GetPointer<int32>());
+    StatusOr<nvinfer1::Dims> trt_dims = const_weights->Shape().AsTrtDims();
+    TRT_ENSURE_OK(trt_dims);
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(*trt_dims, const_weights->GetTrtWeights());
+    TRT_ENSURE(const_layer);
+    nvinfer1::ITensor* output = const_layer->getOutput(0);
+    TRT_ENSURE(output);
+    TRT_ENSURE(output->getType() == nvinfer1::DataType::kINT32);
+    return const_layer;
+  }
+
+  // Adds a Constant layer whose output is a TensorRT shape tensor. The shape
+  // tensor's size and values correspond to dim's nbDims and d[], respectively.
+  StatusOr<nvinfer1::IConstantLayer*> Constant(
+      const std::vector<int>& data) noexcept {
+    nvinfer1::Dims shape_dims;
+    shape_dims.nbDims = 1;
+    shape_dims.d[0] = data.size();
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(nvinfer1::DataType::kINT32, shape_dims);
+    TRT_ENSURE_OK(const_weights);
+    int32* values = const_weights->GetPointer<int32>();
+    for (int i = 0; i < data.size(); i++) {
+      values[i] = static_cast<int32>(data[i]);
+    }
+    StatusOr<nvinfer1::Dims> trt_dims = const_weights->Shape().AsTrtDims();
+    TRT_ENSURE_OK(trt_dims);
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(*trt_dims, const_weights->GetTrtWeights());
+    TRT_ENSURE(const_layer);
+    nvinfer1::ITensor* output = const_layer->getOutput(0);
+    TRT_ENSURE(output);
+    TRT_ENSURE(output->getType() == nvinfer1::DataType::kINT32);
+    TRT_ENSURE(const_layer);
+    return const_layer;
+  }
+
+  // Adds a Constant layer that produces a tensor of shape "shape",
+  // type "data_type" and filled with value "scalar".
+  template <typename T>
+  StatusOr<nvinfer1::IConstantLayer*> Constant(
+      const T value, nvinfer1::Dims shape,
+      nvinfer1::DataType data_type) noexcept {
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(data_type, shape);
+    TRT_ENSURE_OK(const_weights);
+    TRT_ENSURE(const_weights->SetValues(value).ok());
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(shape, const_weights->GetTrtWeights());
+    TRT_ENSURE(const_layer);
+    return const_layer;
+  }
+
+  // Adds a Constant layer that produces a tensor with a single value "scalar".
+  // The tensor has "nb_dims" dimensions and each dimension has only one
+  // element. The data type of the tensor is determined by the data type of
+  // "scalar".
+  template <typename T, typename std::enable_if<std::is_trivially_copyable<
+                            T>::value>::type* = nullptr>
+  StatusOr<nvinfer1::IConstantLayer*> Constant(const T scalar,
+                                               const int nb_dims) noexcept {
+    TRT_ENSURE(nb_dims <= nvinfer1::Dims::MAX_DIMS);
+    auto data_type = nvinfer1::DataType::kINT32;
+    if (std::is_floating_point<T>::value) {
+      data_type = nvinfer1::DataType::kFLOAT;
+    }
+    nvinfer1::Dims zero_shape;
+    zero_shape.nbDims = nb_dims;
+    std::fill_n(zero_shape.d, nb_dims, 1);
+    return Constant<T>(scalar, zero_shape, data_type);
+  }
+
+  // Adds a Constant layer from a TRT_ShapedWeights object.
+  StatusOr<nvinfer1::IConstantLayer*> WeightsToConstant(
+      const nvinfer1::Weights& weights, const DimsAdapter& dims) noexcept {
+    StatusOr<int64_t> vol = dims.Volume();
+    TRT_ENSURE_OK(vol);
+    TRT_ENSURE(*vol == weights.count);
+    StatusOr<nvinfer1::Dims> trt_dims = dims.AsTrtDims();
+    TRT_ENSURE_OK(trt_dims);
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(*trt_dims, weights);
+    TRT_ENSURE(const_layer);
+    return const_layer;
+  }
+
+  Status get_tensor4TensorOrWeights(const TRT_TensorOrWeights& input,
+                                    ITensorProxyPtr* pTensor) {
+    if (input.is_weights()) {
+      StatusOr<nvinfer1::IConstantLayer*> const_layer = WeightsToConstant(
+          input.weights().GetTrtWeights(), input.GetTrtDims());
+      if (!const_layer.status().ok()) return const_layer.status();
+      *pTensor = (*const_layer)->getOutput(0);
+    } else {
+      *pTensor = input.tensor();
+    }
+    return OkStatus();
+  }
+
+  // Creates a nvinfer1::Weights object containing a single scalar.
+  template <typename T, typename std::enable_if<std::is_trivially_copyable<
+                            T>::value>::type* = nullptr>
+  StatusOr<nvinfer1::Weights> ScalarWeights(const T scalar,
+                                            const int nb_dims) noexcept {
+    TRT_ENSURE(nb_dims <= nvinfer1::Dims::MAX_DIMS);
+    auto data_type = nvinfer1::DataType::kINT32;
+    if (std::is_floating_point<T>::value) {
+      data_type = nvinfer1::DataType::kFLOAT;
+    }
+    nvinfer1::Dims weights_shape;
+    weights_shape.nbDims = nb_dims;
+    std::fill_n(weights_shape.d, nb_dims, 1);
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(data_type, weights_shape);
+    TRT_ENSURE_OK(const_weights);
+    const_weights->GetPointer<T>()[0] = scalar;
+    return const_weights->GetTrtWeights();
+  }
+
+  // Adds a TensorRT Slice operation to the network.
+  StatusOr<nvinfer1::ISliceLayer*> Slice(
+      nvinfer1::ITensor* input, const nvinfer1::Dims& begin,
+      const nvinfer1::Dims& size, const nvinfer1::Dims& stride) noexcept {
+    nvinfer1::ISliceLayer* layer =
+        network_->addSlice(*input, begin, size, stride);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds a TensorRT Concatenate operation to the network.
+  StatusOr<nvinfer1::IConcatenationLayer*> Concat(
+      absl::Span<nvinfer1::ITensor* const> inputs, const int axis) {
+    for (nvinfer1::ITensor* input : inputs) {
+      TRT_ENSURE(input);
+    }
+    nvinfer1::IConcatenationLayer* layer = network_->addConcatenation(
+        inputs.data(), static_cast<int32_t>(inputs.size()));
+    TRT_ENSURE(layer);
+    layer->setAxis(axis);
+    return layer;
+  }
+
+  // Adds a TensorRT Concatenate operation to the network.
+  StatusOr<nvinfer1::IConcatenationLayer*> Concat(
+      const std::vector<nvinfer1::ITensor*>& inputs, const int axis) {
+    return this->Concat(absl::MakeSpan(inputs), axis);
+  }
+
+  // Adds a TensorRT Shape operation, which determines the runtime shape of the
+  // input tensor, to the network.
+  StatusOr<nvinfer1::IShapeLayer*> Shape(nvinfer1::ITensor* input) {
+    TRT_ENSURE(input);
+    nvinfer1::IShapeLayer* layer = network_->addShape(*input);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Creates a Gather operation on the shape of the input tensor. The output of
+  // the gather operation is a 1D shape tensor where output[i] = (!sub_one ?
+  // input_shape[i] : input_shape[i] -1) if i is in "indices", otherwise zero.
+  StatusOr<nvinfer1::IGatherLayer*> GetPartialShapeOf(
+      nvinfer1::ITensor* input, absl::InlinedVector<int64, 4> indices,
+      bool sub_one = false) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(indices.size() <= nvinfer1::Dims::MAX_DIMS);
+
+    // Get the runtime shape of input;
+    StatusOr<nvinfer1::IShapeLayer*> shape_layer = this->Shape(input);
+    TRT_ENSURE_PTR_OK(shape_layer);
+    nvinfer1::ITensor* runtime_shape = (*shape_layer)->getOutput(0);
+
+    if (sub_one) {
+      StatusOr<nvinfer1::IConstantLayer*> ones = this->Constant<int32>(1, 1);
+      TRT_ENSURE_PTR_OK(ones);
+      StatusOr<nvinfer1::IElementWiseLayer*> sub =
+          this->Sub(runtime_shape, (*ones)->getOutput(0));
+      TRT_ENSURE_PTR_OK(sub);
+      runtime_shape = (*sub)->getOutput(0);
+    }
+
+    // Create a constant tensor containing the gather indices.
+    // For any dim not in "indices", we mark it size to gather a zero.
+    const int input_nb_dims = input->getDimensions().nbDims;
+    std::vector<int> indices_all(input_nb_dims, input_nb_dims);
+    for (auto idx : indices) {
+      TRT_ENSURE(idx < input_nb_dims);
+      indices_all[idx] = idx;
+    }
+
+    StatusOr<nvinfer1::IConstantLayer*> indices_result =
+        this->Constant(indices_all);
+    TRT_ENSURE_PTR_OK(indices_result);
+    nvinfer1::ITensor* gather_indices = (*indices_result)->getOutput(0);
+    TRT_ENSURE(gather_indices->getDimensions().nbDims == 1);
+    TRT_ENSURE(gather_indices->getType() == nvinfer1::DataType::kINT32);
+
+    // Append a zero to the shape tensor.
+    StatusOr<nvinfer1::IConstantLayer*> zero_result =
+        this->Constant(std::vector<int>{0});
+    TRT_ENSURE_PTR_OK(zero_result);
+    std::array<nvinfer1::ITensor*, 2> cat_inputs = {
+        runtime_shape, (*zero_result)->getOutput(0)};
+    nvinfer1::IConcatenationLayer* cat_layer =
+        network_->addConcatenation(cat_inputs.data(), cat_inputs.size());
+    TRT_ENSURE(cat_layer);
+    nvinfer1::ITensor* gather_input = cat_layer->getOutput(0);
+    TRT_ENSURE(gather_input);
+
+    // Finally, gather the indices from the input.
+    nvinfer1::IGatherLayer* gather =
+        network_->addGather(*gather_input, *gather_indices, 0);
+    TRT_ENSURE(gather);
+    return gather;
+  }
+
+  // Adds a scale layer that uniformly scales the input tensor by the specified
+  // amount.
+  StatusOr<nvinfer1::IScaleLayer*> AddUniformScale(nvinfer1::ITensor* input,
+                                                   float scale,
+                                                   const std::string& name) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(!name.empty());
+    StatusOr<nvinfer1::Weights> weight = this->ScalarWeights<float>(scale, 1);
+    TRT_ENSURE_OK(weight);
+    const nvinfer1::Weights empty_weights =
+        nvinfer1::Weights{nvinfer1::DataType::kFLOAT, nullptr, 0};
+    nvinfer1::IScaleLayer* scale_layer =
+        network_->addScale(*input, nvinfer1::ScaleMode::kUNIFORM, empty_weights,
+                           (*weight), empty_weights);
+    TRT_ENSURE(scale_layer != nullptr);
+    scale_layer->setName(name.c_str());
+    TRT_ENSURE((*scale_layer).getPower().count == 0);
+    TRT_ENSURE((*scale_layer).getShift().count == 0);
+    TRT_ENSURE((*scale_layer).getScale().count == 1);
+    return scale_layer;
+  }
+
+  StatusOr<nvinfer1::ILayer*> AddFill(const TRT_TensorOrWeights& value_input,
+                                      const TRT_TensorOrWeights& dims_input,
+                                      bool is_value_static, bool is_dims_static,
+                                      int nbDims,
+                                      const nvinfer1::Dims& trt_dims,
+                                      ITensorProxyPtr scalar_tensor = nullptr,
+                                      ITensorProxyPtr beta_tensor = nullptr,
+                                      const float delta = 0) {
+    // TensorRT IFillLayer requires a rank 0 scalar.
+    nvinfer1::Dims scalar_dims;
+    scalar_dims.nbDims = 0;
+    if (is_value_static) {
+      StatusOr<nvinfer1::IConstantLayer*> const_layer =
+          WeightsToConstant(value_input.weights().GetTrtWeights(), scalar_dims);
+      if (!const_layer.status().ok()) return const_layer.status();
+      scalar_tensor = (*const_layer)->getOutput(0);
+    } else {
+      if (scalar_tensor == nullptr) {
+        StatusOr<nvinfer1::IShuffleLayer*> shuffler_layer =
+            Reshape(value_input.tensor()->trt_tensor(), scalar_dims);
+        if (!shuffler_layer.status().ok()) return shuffler_layer.status();
+        scalar_tensor = (*shuffler_layer)->getOutput(0);
+      }
+    }
+
+    if (beta_tensor == nullptr) {
+      nvinfer1::Dims beta_shape{1, {nbDims}};
+      StatusOr<nvinfer1::IConstantLayer*> const_layer =
+          Constant(delta, beta_shape, value_input.TrtDType());
+      TF_RETURN_IF_ERROR(const_layer.status());
+      beta_tensor = (*const_layer)->getOutput(0);
+    }
+
+    nvinfer1::IFillLayer* layer =
+        network_->addFill(trt_dims, nvinfer1::FillOperation::kLINSPACE);
+    TRT_ENSURE(layer);
+    if (!is_dims_static) {
+      layer->setInput(0, *dims_input.tensor()->trt_tensor());
+    }
+    layer->setInput(1, *scalar_tensor->trt_tensor());
+    layer->setInput(2, *beta_tensor->trt_tensor());
+    return layer;
+  }
+
+  // Adds a quantization layer that uniformly scales the input tensor
+  // by the given multiplicative "scaling_factor", then rounds
+  // (round-to-nearest-ties-to-even) to the nearest integer and clamps in the
+  // range of [-128, 127].
+  StatusOr<nvinfer1::ILayer*> Quantize(nvinfer1::ITensor* input,
+                                       const float scaling_factor,
+                                       const std::string& name) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(!name.empty());
+    // Preprocessor usage here is unavoidable because TRT8 API is new.
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+    // The TensorRT IQuantizeLayer divides by the scale factor rather than
+    // multiplies. To be consistent, in this function we expect a multiplicative
+    // scale factor, so we take the reciprical.
+    StatusOr<nvinfer1::IConstantLayer*> scaling_const =
+        this->Constant<float>(1.0f / scaling_factor, 1);
+    TRT_ENSURE_PTR_OK(scaling_const);
+    (*scaling_const)->setDimensions(nvinfer1::Dims{0, {}});
+    nvinfer1::IQuantizeLayer* quant_layer =
+        network_->addQuantize(*input, *(*scaling_const)->getOutput(0));
+    TRT_ENSURE(quant_layer);
+    quant_layer->setAxis(1);
+    return quant_layer;
+#else
+    StatusOr<nvinfer1::IScaleLayer*> result =
+        this->AddUniformScale(input, scaling_factor, name);
+    TRT_ENSURE_PTR_OK(result);
+    (*result)->setOutputType(0, nvinfer1::DataType::kINT8);
+    (*result)->setPrecision(nvinfer1::DataType::kFLOAT);
+    return result;
+#endif
+  }
+
+  // Adds a dequantize layer that casts the input tensor to TensorRT float type
+  // and scales it uniformly by the given multiplicative "scaling_factor".
+  StatusOr<nvinfer1::ILayer*> Dequantize(nvinfer1::ITensor* input,
+                                         const float scaling_factor,
+                                         const std::string& name) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(!name.empty());
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+    StatusOr<nvinfer1::IConstantLayer*> scaling_const =
+        this->Constant<float>(scaling_factor, 1);
+    TRT_ENSURE_PTR_OK(scaling_const);
+    (*scaling_const)->setDimensions(nvinfer1::Dims{0, {}});
+    nvinfer1::IDequantizeLayer* dequant_layer =
+        network_->addDequantize(*input, *(*scaling_const)->getOutput(0));
+    dequant_layer->setAxis(1);
+    TRT_ENSURE(dequant_layer);
+    return dequant_layer;
+#else
+    StatusOr<nvinfer1::IScaleLayer*> result =
+        this->AddUniformScale(input, scaling_factor, name);
+    TRT_ENSURE_PTR_OK(result);
+    (*result)->setOutputType(0, nvinfer1::DataType::kFLOAT);
+    (*result)->setPrecision(nvinfer1::DataType::kINT8);
+    return result;
+#endif
+  }
+
+  // Adds TensorRT Q/DQ operations. This is for explicit precision mode.
+  StatusOr<nvinfer1::ILayer*> UniformQuantizeDequantizeExplicit(
+      nvinfer1::ITensor* input, float quantize_scale, float dequantize_scale,
+      const std::string& name) {
+    TRT_ENSURE(input);
+    if (!IS_TRT_VERSION_GE(8, 0, 0, 0)) {
+      TRT_ENSURE(network_->hasExplicitPrecision());
+    }
+    TRT_ENSURE(IS_TRT_VERSION_GE(7, 1, 0, 0));
+
+    static int count = 0;
+    TRT_ENSURE(input->getType() == nvinfer1::DataType::kFLOAT);
+    std::string quant_name = absl::StrCat(input->getName(), "_quant_", count);
+
+    StatusOr<nvinfer1::ILayer*> quant =
+        this->Quantize(input, quantize_scale, quant_name);
+    TRT_ENSURE_PTR_OK(quant);
+
+    std::string dequant_name =
+        absl::StrCat(input->getName(), "_dequant_", count);
+    StatusOr<nvinfer1::ILayer*> dequant = this->Dequantize(
+        (*quant)->getOutput(0), dequantize_scale, dequant_name);
+    TRT_ENSURE_PTR_OK(dequant);
+
+    count++;
+    return dequant;
+  }
+
+  StatusOr<nvinfer1::IShuffleLayer*> Reshape(nvinfer1::ITensor* input,
+                                             const nvinfer1::Dims& new_shape) {
+    TRT_ENSURE(input);
+    nvinfer1::IShuffleLayer* layer = network_->addShuffle(*input);
+    TRT_ENSURE(layer);
+    layer->setReshapeDimensions(new_shape);
+    return layer;
+  }
+
+  StatusOr<nvinfer1::ILayer*> FindProducerOf(const nvinfer1::ITensor* tensor) {
+    const char* name = tensor->getName();
+    const int num_layers = network_->getNbLayers();
+    for (int i = 0; i < num_layers; i++) {
+      nvinfer1::ILayer* layer = network_->getLayer(i);
+      const int num_outputs = layer->getNbOutputs();
+      for (int j = 0; j < num_outputs; j++) {
+        nvinfer1::ITensor* t = layer->getOutput(j);
+        if (std::string(t->getName()) == name) {
+          return layer;
+        }
+      }
+    }
+    return errors::NotFound("could not find producing layer of ", name);
+  }
+
+  StatusOr<nvinfer1::ILayer*> UniqueParentOf(const nvinfer1::ILayer* layer,
+                                             int input_idx = 0) {
+    return FindProducerOf(layer->getInput(input_idx));
+  }
+
+  nvinfer1::INetworkDefinition* Network() { return network_; }
+
+ private:
+  nvinfer1::INetworkDefinition* network_;
+  TrtWeightStore* weight_store_;
+};
+
+class ShuffleBuilder {
+ private:
+  explicit ShuffleBuilder(TRTNetworkBuilder* builder, nvinfer1::ITensor* input)
+      : builder_(builder) {
+    layer_ = builder->Network()->addShuffle(*input);
+  }
+
+ public:
+  static StatusOr<ShuffleBuilder> Create(TRTNetworkBuilder* builder,
+                                         nvinfer1::ITensor* input) {
+    TRT_ENSURE(builder != nullptr);
+    TRT_ENSURE(input != nullptr);
+    return ShuffleBuilder(builder, input);
+  }
+
+  ShuffleBuilder& SetReshape(const nvinfer1::Dims& dims) {
+    layer_->setReshapeDimensions(dims);
+    return *this;
+  }
+
+  ShuffleBuilder& SetReshape(nvinfer1::ITensor* shape) {
+    layer_->setInput(1, *shape);
+    return *this;
+  }
+
+  ShuffleBuilder& SetFirstTranspose(const nvinfer1::Permutation& perm) {
+    layer_->setFirstTranspose(perm);
+    return *this;
+  }
+
+  ShuffleBuilder& SetSecondTranspose(const nvinfer1::Permutation& perm) {
+    layer_->setSecondTranspose(perm);
+    return *this;
+  }
+
+  StatusOr<nvinfer1::ITensor*> Output() {
+    TRT_ENSURE(layer_ != nullptr);
+    TRT_ENSURE(layer_->getOutput(0) != nullptr);
+    return layer_->getOutput(0);
+  }
+
+ private:
+  TRTNetworkBuilder* builder_;
+  nvinfer1::IShuffleLayer* layer_;
+};
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_LAYER_UTILS_H_

tensorflow/compiler/tf2tensorrt/convert/ops/like_ops.cc

@@ -0,0 +1,94 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+template <int V>
+class ConvertLikeOps : public OpConverterBase<ConvertLikeOps<V>> {
+ public:
+  explicit ConvertLikeOps(const OpConverterParams *params)
+      : OpConverterBase<ConvertLikeOps<V>>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+    };
+  }
+  Status Validate() { return ConvertLikeOps<V>::NotSupportedInImplicitBatch(); }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    auto *network = params.converter->network();
+
+    const TRT_TensorOrWeights &input = inputs.at(0);
+    nvinfer1::Dims dims(input.GetTrtDims());
+
+    const std::vector<int> value_input_dims_data = {1};
+    const DimsAdapter value_input_dims(value_input_dims_data);
+    StatusOr<TRT_ShapedWeights> value_weights =
+        params.weight_store->GetTempWeights(input.TrtDType(), value_input_dims);
+    TF_RETURN_IF_ERROR(value_weights.status());
+    TF_RETURN_IF_ERROR(value_weights->SetValues(V));
+    TRT_TensorOrWeights value_input(value_weights.value());
+
+    const auto is_dims_static = HasStaticShape(dims);
+    auto builder = TRTNetworkBuilder::Create(network, params.weight_store);
+    ITensorProxyPtr dims_input_tensor;
+    if (!is_dims_static) {
+      StatusOr<nvinfer1::IShapeLayer *> shape_layer =
+          builder->Shape(input.tensor()->trt_tensor());
+      TF_RETURN_IF_ERROR(shape_layer.status());
+      dims_input_tensor = (*shape_layer)->getOutput(0);
+      dims.nbDims = 0;
+    }
+
+    TRT_TensorOrWeights dims_input(dims_input_tensor);
+    StatusOr<nvinfer1::ILayer *> layer =
+        builder->AddFill(value_input, dims_input, true, is_dims_static,
+                         input.GetTrtDims().nbDims, dims);
+    ITensorProxyPtr output_tensor = (*layer)->getOutput(0);
+    this->AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<0>>(),
+                                  "zeros_like");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<1>>(),
+                                  "ones_like");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<0>>(),
+                                  "ZerosLike");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<1>>(),
+                                  "OnesLike");
+
+#endif  // IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/log_softmax.cc

@@ -0,0 +1,104 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertLogSoftmax : public OpConverterBase<ConvertLogSoftmax> {
+ public:
+  explicit ConvertLogSoftmax(const OpConverterParams *params)
+      : OpConverterBase<ConvertLogSoftmax>(params) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("logits", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+    if (!num_trt_dims && params.use_implicit_batch) {
+      return errors::InvalidArgument(
+          "TensorRT LogSoftmax cannot apply on the batch dimension");
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    const auto &node_def = params.node_def;
+
+    // Perform LogSoftmax operation:
+    // `logsoftmax = logits - log(reduce_sum(exp(logits), axis))`
+
+    // Get the logits tensor.
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+
+    // Exponent of logits.
+    nvinfer1::IUnaryLayer *exp = params.converter->network()->addUnary(
+        *logits_tensor->trt_tensor(), nvinfer1::UnaryOperation::kEXP);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(exp, node_def.name());
+    params.converter->SetLayerName(exp, node_def, "exp");
+
+    // Reduce-sum operation across the final dimension.
+    nvinfer1::IReduceLayer *reduced_sum =
+        params.converter->network()->addReduce(
+            *exp->getOutput(0), nvinfer1::ReduceOperation::kSUM,
+            (1 << (num_trt_dims - 1)), /*Reduce across final dimension*/
+            true /*Keep reduced dims*/);
+    params.converter->SetLayerName(reduced_sum, node_def, "reduced_sum");
+
+    // Logarithm of reduced_sum.
+    nvinfer1::IUnaryLayer *log_reduced_sum =
+        params.converter->network()->addUnary(*reduced_sum->getOutput(0),
+                                              nvinfer1::UnaryOperation::kLOG);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(log_reduced_sum, node_def.name());
+    params.converter->SetLayerName(log_reduced_sum, node_def,
+                                   "log_reduced_sum");
+
+    // Finally, get the output by subtracting log_reduced_sum from logits.
+    nvinfer1::IElementWiseLayer *sub =
+        params.converter->network()->addElementWise(
+            *logits_tensor->trt_tensor(), *log_reduced_sum->getOutput(0),
+            nvinfer1::ElementWiseOperation::kSUB);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(sub, node_def.name());
+    params.converter->SetLayerName(sub, node_def, "sub");
+
+    params.outputs->push_back(TRT_TensorOrWeights(sub->getOutput(0)));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLogSoftmax>(),
+                                  "LogSoftmax");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.cc

@@ -0,0 +1,421 @@
+/* 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/convert/ops/quantization_ops.h"
+
+#include "absl/strings/str_format.h"
+#include "tensorflow/cc/ops//array_ops.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/weights.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+bool IsQuantizeAndDequantizeOp(const Node* node) {
+  return absl::c_find(kQuantizationOpNames, node->def().op()) !=
+         kQuantizationOpNames.end();
+}
+
+namespace {
+
+// Provides quantizing and dequantizing tensor scales for a given dynamic range.
+// Borrowed from TF quantization kernel logic.
+template <typename T>
+QuantizationScales<T, 1> ComputeQuantizationRange(bool signed_input,
+                                                  int num_bits,
+                                                  bool narrow_range,
+                                                  T* min_range, T* max_range) {
+  // Calculate the range for the simulated integer quantization:
+  // e.g. [-127,127] for signed = true, narrow_range = true, num_bits = 8,
+  // or [-128,127] for signed = true, narrow_range = false, num_bits = 8,
+  // or [0, 255] for signed = false, num_bits = 8.
+  const int64_t min_quantized =
+      signed_input ? narrow_range ? -(1ULL << (num_bits - 1)) + 1
+                                  : -(1ULL << (num_bits - 1))
+                   : 0;
+  const int64_t max_quantized =
+      signed_input ? (1ULL << (num_bits - 1)) - 1 : (1ULL << num_bits) - 1;
+  // Determine the maximum scaling factor that would scale
+  // [min_range, max_range] to not exceed [min_quantized, max_quantized],
+  // while keeping 0 unchanged.
+  const T scale_from_min_side = (min_quantized * *min_range > 0)
+                                    ? min_quantized / *min_range
+                                    : std::numeric_limits<T>::max();
+  const T scale_from_max_side = (max_quantized * *max_range > 0)
+                                    ? max_quantized / *max_range
+                                    : std::numeric_limits<T>::max();
+
+  QuantizationScales<T, 1> scales;
+  // Note: Avoids changing the side of the range that determines scale.
+  if (scale_from_min_side < scale_from_max_side) {
+    scales.quantize_scale[0] = scale_from_min_side;
+    scales.dequantize_scale[0] = *min_range / min_quantized;
+    *max_range = max_quantized * scales.dequantize_scale[0];
+  } else {
+    scales.quantize_scale[0] = scale_from_max_side;
+    scales.dequantize_scale[0] = *max_range / max_quantized;
+    *min_range = min_quantized * scales.dequantize_scale[0];
+  }
+  return scales;
+}
+
+// Prepares the input for a QDQ node in explicit precision mode, returning a
+// ITensor pointer. If the input is weights, we convert it to a ITensor by
+// adding a constant layer.
+StatusOr<nvinfer1::ITensor*> ExlicitQDQInputToTensor(
+    TRTNetworkBuilder* builder, const OpConverterParams* params,
+    const TRT_TensorOrWeights& input) {
+  if (input.is_tensor()) {
+    return input.tensor()->trt_tensor();
+  }
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && input.weights().count() > 1) {
+    LOG(WARNING) << absl::StrCat(
+        "QDQ per-channel for weights not "
+        "implemented, assuming uniform scaling");
+  }
+  TRT_ShapedWeights trt_weights = input.weights();
+  StatusOr<nvinfer1::IConstantLayer*> weights_const =
+      builder->WeightsToConstant(trt_weights.GetTrtWeights(),
+                                 trt_weights.Shape());
+  TRT_ENSURE_PTR_OK(weights_const);
+  params->converter->SetLayerName(*weights_const, params->node_def, "const");
+  nvinfer1::ITensor* qdq_input = (*weights_const)->getOutput(0);
+  std::string name = absl::StrCat((*weights_const)->getName(), "_output");
+  qdq_input->setName(name.c_str());
+  return qdq_input;
+}
+
+}  // namespace
+
+// Carries traits for each specific quantization op type for conversion.
+// Specialization for template parameter T should be given for each TF C++
+// quantization op.
+template <typename T>
+struct QDQOpSpec {};
+
+template <>
+struct QDQOpSpec<ops::QuantizeAndDequantizeV2> {
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+        InputArgSpec::Create("input_min", TrtInputArg::kWeight),
+        InputArgSpec::Create("input_max", TrtInputArg::kWeight),
+    };
+  }
+
+  struct Attrs {
+    float min_range;
+    float max_range;
+    bool narrow_range;
+    std::string round_mode;
+    UniformQuantizationScales scales;
+  };
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    AttrSlice attrs(node_def);
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "round_mode", &args->round_mode));
+    if (args->round_mode != "HALF_TO_EVEN") {
+      LOG(WARNING) << node_def.op() << ": " << node_def.name()
+                   << " has round_mode=" << args->round_mode
+                   << ", but for TensorRT conversion, "
+                      "round_mode=HALF_TO_EVEN is recommended.";
+    }
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "narrow_range", &args->narrow_range));
+    if (args->narrow_range) {
+      LOG(WARNING) << node_def.op() << ": " << node_def.name()
+                   << " has narrow_range=true, but for TensorRT conversion, "
+                      "narrow_range=false is recommended.";
+    }
+    args->min_range = inputs.at(1).weights().template GetPointer<float>()[0];
+    args->max_range = inputs.at(2).weights().template GetPointer<float>()[0];
+    const int num_bits = 8;
+    args->scales = ComputeQuantizationRange<float>(
+        /*signed_input=*/true, num_bits, args->narrow_range, &args->min_range,
+        &args->max_range);
+    TRT_ENSURE(args->scales.dequantize_scale[0] != 0);
+    TRT_ENSURE(args->scales.quantize_scale[0] != 0);
+    return OkStatus();
+  }
+
+  // Converts in explicit precision mode. In this mode, QDQ operations are
+  // directly converted into TensorRT quantizing and dequantizing scale
+  // operations.
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    const auto& node_def = params->node_def;
+
+    StatusOr<TRTNetworkBuilder> builder = TRTNetworkBuilder::Create(
+        params->converter->network(), params->weight_store);
+
+    StatusOr<nvinfer1::ITensor*> qdq_input =
+        ExlicitQDQInputToTensor(&*builder, params, params->inputs.at(0));
+    TRT_ENSURE_PTR_OK(qdq_input);
+
+    // TODO(cbate): check this condition exists for TRT8? Outline this block to
+    // a "reshape policy".
+    const int required_dims = params->use_implicit_batch ? 3 : 4;
+    const nvinfer1::Dims idims = (*qdq_input)->getDimensions();
+    nvinfer1::Dims intermediate_dims = idims;
+    TRT_ENSURE(idims.nbDims > 0);
+    if (idims.nbDims < required_dims) {
+      const int nb_extra_dims = required_dims - idims.nbDims;
+      intermediate_dims.nbDims = required_dims;
+      std::vector<int> ones(nb_extra_dims, 1);
+      TRT_ENSURE(ones.size() == nb_extra_dims && nb_extra_dims > 0);
+
+      if (!params->use_implicit_batch) {
+        intermediate_dims.d[0] = idims.d[0];
+        std::copy(ones.begin(), ones.end(), intermediate_dims.d + 1);
+        std::copy_n(idims.d + 1, idims.nbDims - 1,
+                    intermediate_dims.d + ones.size() + 1);
+      } else {
+        std::copy(ones.begin(), ones.end(), intermediate_dims.d);
+        std::copy_n(idims.d, idims.nbDims, intermediate_dims.d + ones.size());
+      }
+
+      LOG(WARNING) << absl::StrCat(
+          node_def.name(), ":", node_def.op(), ": tensor ",
+          (*qdq_input)->getName(), " has shape ", DebugString(idims),
+          " but TRT scale layer requires at least 3 dims excluding batch dim, "
+          "trying to recover by inserting 1's to create shape ",
+          DebugString(intermediate_dims));
+      StatusOr<nvinfer1::IShuffleLayer*> reshape =
+          builder->Reshape(*qdq_input, intermediate_dims);
+      TRT_ENSURE_PTR_OK(reshape);
+      *qdq_input = (*reshape)->getOutput(0);
+    }
+
+    VLOG(1) << "[ExplicitPrecision]" << node_def.op() << ": " << node_def.name()
+            << " computed scales: " << args.scales << " from min/max ranges "
+            << args.min_range << "/" << args.max_range;
+
+    StatusOr<nvinfer1::ILayer*> qdq =
+        builder->UniformQuantizeDequantizeExplicit(
+            *qdq_input, args.scales.quantize_scale[0],
+            args.scales.dequantize_scale[0], node_def.name());
+    TRT_ENSURE_PTR_OK(qdq);
+    ITensorProxyPtr final_output = (*qdq)->getOutput(0);
+    if (idims.nbDims != intermediate_dims.nbDims) {
+      StatusOr<nvinfer1::IShuffleLayer*> undo_reshape =
+          builder->Reshape(*qdq_input, idims);
+      TRT_ENSURE_PTR_OK(undo_reshape);
+      final_output = (*undo_reshape)->getOutput(0);
+    }
+    params->outputs->push_back(final_output);
+    return OkStatus();
+  }
+};
+
+template <>
+
+struct QDQOpSpec<ops::QuantizeAndDequantizeV3> {
+  static constexpr std::array<InputArgSpec, 4> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+        InputArgSpec::Create("min", TrtInputArg::kWeight),
+        InputArgSpec::Create("max", TrtInputArg::kWeight),
+        InputArgSpec::Create("num_bits", TrtInputArg::kWeight),
+    };
+  }
+  // Use same attributes and conversion functions as QDQV2.
+  using Attrs = QDQOpSpec<ops::QuantizeAndDequantizeV2>::Attrs;
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    return QDQOpSpec<
+        ops::QuantizeAndDequantizeV2>::ValidateQDQForExplicitPrecision(inputs,
+                                                                       node_def,
+                                                                       args);
+  }
+
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    return QDQOpSpec<ops::QuantizeAndDequantizeV2>::ConvertExplicit(params,
+                                                                    args);
+  }
+};
+
+template <>
+
+struct QDQOpSpec<ops::FakeQuantWithMinMaxVars> {
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+        InputArgSpec::Create("min", TrtInputArg::kWeight),
+        InputArgSpec::Create("max", TrtInputArg::kWeight),
+    };
+  }
+  struct Attrs {
+    int num_bits;
+    bool narrow_range;
+  };
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    return errors::Unimplemented("");
+  }
+
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    return errors::Unimplemented("");
+  }
+};
+
+template <>
+
+struct QDQOpSpec<ops::FakeQuantWithMinMaxArgs> {
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+    };
+  }
+
+  struct Attrs {
+    float min;
+    float max;
+    int num_bits;
+    bool narrow_range;
+  };
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    return errors::Unimplemented("");
+  }
+
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    return errors::Unimplemented("");
+  }
+};
+
+// Converts QDQ operations in non-explicit precision mode. This is the original
+// "ConvertQuantize" function. In this mode, Q/DQ operations are no-ops and are
+// instead used to set the dynamic range of the input tensor.
+Status ConvertDynamicRangeMode(const OpConverterParams* params) {
+  const auto& inputs = params->inputs;
+  const auto& node_def = params->node_def;
+  float min_range = 0.0f;
+  float max_range = 0.0f;
+  const auto& op_name = node_def.op();
+  if (op_name == "FakeQuantWithMinMaxArgs") {
+    AttrSlice attrs(node_def);
+    // Get ranges via node attributes.
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "min", &min_range));
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "max", &max_range));
+  } else if (op_name == "FakeQuantWithMinMaxVars" ||
+             op_name == "QuantizeAndDequantizeV2" ||
+             op_name == "QuantizeAndDequantizeV3") {
+    // Get ranges via inputs.
+    auto get_weights_value = [&inputs](int index) {
+      const auto* raw_weights = inputs.at(index).weights().GetPointer<float>();
+      return raw_weights[0];
+    };
+    min_range = get_weights_value(1);
+    max_range = get_weights_value(2);
+  } else {
+    return errors::InvalidArgument("Unknown quantization op ", op_name, ", at ",
+                                   node_def.name());
+  }
+  if (params->validation_only) {
+    return OkStatus();
+  }
+
+  // Store ranges for tensor
+  ITensorProxyPtr input0 = inputs.at(0).tensor();
+  params->converter->ProvideQuantizationRange(&input0, min_range, max_range);
+  // Sometimes, TRT may not quantize a tensor, either because it chooses to
+  // execute a higher precision kernel or because of op fusion. In these
+  // cases, accuracy will suffer if the model was trained to expect
+  // quantization at that tensor. We should consider adding a clip(tensor,
+  // min_range, max_range) operation here to ensure that any arbitrarily
+  // placed quantize node will execute as expected. However, this will
+  // negatively affect performance. If users train their models in a way which
+  // models inference as close as possible (i.e. not quantizing in place where
+  // fusion will occur), then there is no problem with the current
+  // implementation.
+  params->outputs->push_back(inputs.at(0));
+  return OkStatus();
+}
+
+template <typename TFOpType>
+class ConvertQDQ : public OpConverterBase<ConvertQDQ<TFOpType>> {
+ public:
+  explicit ConvertQDQ(const OpConverterParams* params)
+      : OpConverterBase<ConvertQDQ<TFOpType>>(params) {}
+
+  static constexpr auto InputSpec() { return QDQOpSpec<TFOpType>::InputSpec(); }
+
+  // Disable the non-applicable data type check by providing empty string.
+  static constexpr const char* NodeDefDataTypeAttributeName() { return ""; }
+
+  Status ValidateDynamicRangeINT8Mode() {
+    // The condition ensures we only call the conversion once. We should break
+    // this function up into validation and conversion.
+    if (this->params_->validation_only) {
+      return ConvertDynamicRangeMode(this->params_);
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    if (!this->params_->use_explicit_precision) {
+      return ValidateDynamicRangeINT8Mode();
+    }
+    return OpSpec::ValidateQDQForExplicitPrecision(
+        this->params_->inputs, this->params_->node_def, &attrs_);
+  }
+
+  Status Convert() {
+    if (!this->params_->use_explicit_precision) {
+      return ConvertDynamicRangeMode(this->params_);
+    }
+    return OpSpec::ConvertExplicit(this->params_, attrs_);
+  }
+
+  using OpSpec = QDQOpSpec<TFOpType>;
+  using OpSpecAttrs = typename QDQOpSpec<TFOpType>::Attrs;
+  OpSpecAttrs attrs_;
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::QuantizeAndDequantizeV2>>(),
+    "QuantizeAndDequantizeV2");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::QuantizeAndDequantizeV3>>(),
+    "QuantizeAndDequantizeV3");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::FakeQuantWithMinMaxVars>>(),
+    "FakeQuantWithMinMaxVars");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::FakeQuantWithMinMaxArgs>>(),
+    "FakeQuantWithMinMaxArgs");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_QUANTIZATION_OPS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_QUANTIZATION_OPS_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+constexpr std::array<const char*, 4> kQuantizationOpNames = {
+    "QuantizeAndDequantizeV2",
+    "QuantizeAndDequantizeV3",
+    "FakeQuantWithMinMaxVars",
+    "FakeQuantWithMinMaxArgs",
+};
+
+// Operations with supported conversion to Q/DQ ops in TensorRT explicit
+// precision mode.
+constexpr std::array<const char*, 1> kExplicitQuantizationOpNames = {
+    "QuantizeAndDequantizeV2",
+};
+
+// Contains two scaling factors for quantization and dequantization
+// respectively. A shift factor is omitted as TensorRT only supports symmetric
+// quantization.
+template <typename T, size_t N>
+struct QuantizationScales {
+  std::array<T, N> quantize_scale;
+  std::array<T, N> dequantize_scale;
+};
+
+// In TensorRT 7 and 8, only uniform tensor scaling is supported for
+// activations.
+using UniformQuantizationScales = QuantizationScales<float, 1>;
+
+// Per-channel scaling is supported for weights in TensorRT version >= 8.0.
+template <size_t ChannelDimSize>
+using PerChannelQuantizationScales = QuantizationScales<float, ChannelDimSize>;
+
+template <typename T, size_t N>
+std::ostream& operator<<(std::ostream& os,
+                         const QuantizationScales<T, N>& scales) {
+  os << absl::StrFormat("QuantizationScales[quantize={%s},dequantize={%s}]",
+                        absl::StrJoin(scales.quantize_scale, ","),
+                        absl::StrJoin(scales.dequantize_scale, ","));
+  return os;
+}
+
+// Returns true if the Tensorflow node is a quantize and dequantize operation.
+bool IsQuantizeAndDequantizeOp(const Node*);
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_QUANTIZATION_OPS_H_

tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops_test.cc

@@ -0,0 +1,618 @@
+/* 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/convert/ops/quantization_ops.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/linalg_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/compiler/jit/shape_inference.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/trt_convert_api.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+namespace ops = ::tensorflow::ops;
+using ::tensorflow::testing::StatusIs;
+
+// This anonymous namespace contains helper functions for instantiating small TF
+// building blocks. These are used below to construct specific graph patterns
+// which test end-to-end conversion of the TF graph to an explciit-precision
+// enabled TensorRT network.
+namespace {
+
+enum class ConvEpilogueType {
+  kNone,
+  kReLU,
+  kBatchNorm,
+  kReLUBatchnorm,
+  kBatchnormReLU
+};
+
+std::ostream& operator<<(std::ostream& os, ConvEpilogueType epilogue) {
+  switch (epilogue) {
+    case ConvEpilogueType::kNone:
+      return os << "None";
+    case ConvEpilogueType::kReLU:
+      return os << "ReLU only";
+    case ConvEpilogueType::kBatchNorm:
+      return os << "BatchNorm Only";
+    case ConvEpilogueType::kReLUBatchnorm:
+      return os << "ReLU+Batchnorm";
+    case ConvEpilogueType::kBatchnormReLU:
+      return os << "BatchNorm+ReLU";
+  }
+}
+
+std::string DebugString(ConvEpilogueType epilogue) {
+  std::stringstream ss;
+  ss << epilogue;
+  return ss.str();
+}
+
+// Adds a 2D 3x3, single channel input with specified data_format. data_format
+// must be NHWC,NCHW or NHW.
+ops::Placeholder AddInput(Scope scope, int input_idx,
+                          const std::string data_format,
+                          std::array<int, 3> size_chw = {1, 3, 3}) {
+  PartialTensorShape input_shape;
+  if (data_format == "NCHW") {
+    input_shape =
+        PartialTensorShape({1, size_chw[0], size_chw[1], size_chw[2]});
+  } else if (data_format == "NHWC") {
+    input_shape =
+        PartialTensorShape({1, size_chw[1], size_chw[2], size_chw[0]});
+  } else if (data_format == "NHW") {
+    input_shape = PartialTensorShape({1, size_chw[1], size_chw[2]});
+  } else {
+    LOG(FATAL) << "Unknown input shape type " << data_format;
+  }
+  auto input_attrs = ops::Placeholder::Attrs().Shape(input_shape);
+  return ops::Placeholder(scope.WithOpName(absl::StrCat("input_", input_idx)),
+                          DT_FLOAT, input_attrs);
+}
+
+// Adds QDQ op with min = -1.0f, max = 1.0f.
+Output AddQDQV2(Scope scope, Input input) {
+  // Create scaling factors.
+  auto input_min =
+      ops::Const<float>(scope.WithOpName("in_min"), -1.0f, TensorShape{});
+  auto input_max =
+      ops::Const<float>(scope.WithOpName("in_max"), 1.0f, TensorShape{});
+  return ops::QuantizeAndDequantizeV2(scope.WithOpName("qdq"), input, input_min,
+                                      input_max);
+}
+
+Output AddOutput(Scope scope, Output input, int idx, bool add_qdq) {
+  Output out = input;
+  if (add_qdq) {
+    out = AddQDQV2(scope, input);
+  }
+  return ops::Identity(scope.WithOpName(StrCat("output_", idx)), out);
+}
+
+// Adds a 3x3x1x1 Conv2D op and optional bias weights, followed by ReLU
+// activation. Puts QDQ between (weights, op). Puts QDQ between (input, op)
+// when qdq_on_output=false. Otherwise, puts QDQ between (op, output).
+Output AddConv2D(Scope scope, Input input, int in_channels, int out_channels,
+                 std::array<int, 2> filter_size = {1, 1},
+                 std::array<int, 2> stride = {1, 1},
+                 const std::string& data_format = "NCHW", bool with_bias = true,
+                 ConvEpilogueType epilogue = ConvEpilogueType::kBatchnormReLU,
+                 bool qdq_on_output = false) {
+  // Create 3x3 non-quantized weights weights.
+  auto weights_const = ops::Const(
+      scope.WithOpName("weights"), 1.0f,
+      TensorShape({filter_size[0], filter_size[1], in_channels, out_channels}));
+
+  // Add QDQ to input if we don't add QDQ to output.
+  auto conv_input =
+      !qdq_on_output ? AddQDQV2(scope.WithOpName("qdq_input"), input) : input;
+
+  Output result = ops::Conv2D(
+      scope.WithOpName("conv2d"), conv_input, AddQDQV2(scope, weights_const),
+      /*strides=*/{1, 1, 1, 1},
+      /*padding=*/"SAME", ops::Conv2D::Attrs().DataFormat(data_format));
+
+  if (with_bias) {
+    auto bias_const = ops::Const(scope.WithOpName("bias_weights"), 1.0f,
+                                 TensorShape({
+                                     out_channels,
+                                 }));
+    result = ops::BiasAdd(scope.WithOpName("bias"), result, bias_const,
+                          ops::BiasAdd::Attrs().DataFormat(data_format));
+  }
+
+  auto add_bn = [scope, data_format](Input input,
+                                     const int channels) -> Output {
+    TensorShape constant_shape = TensorShape({channels});
+    auto bn_scale =
+        ops::Const(scope.WithOpName("bn_scale"), 1.0f, constant_shape);
+    auto bn_offset =
+        ops::Const(scope.WithOpName("bn_offset"), 1.0f, constant_shape);
+    auto bn_mean =
+        ops::Const(scope.WithOpName("bn_mean"), 0.1f, TensorShape({channels}));
+    auto bn_var =
+        ops::Const(scope.WithOpName("bn_var"), 1.0f, TensorShape({channels}));
+    Input conv_bn_input = IS_TRT_VERSION_GE(8, 0, 1, 0)
+                              ? input
+                              : AddQDQV2(scope.WithOpName("qdq_input"), input);
+    return ops::FusedBatchNormV3(
+               scope.WithOpName("bn"), conv_bn_input, bn_scale, bn_offset,
+               bn_mean, bn_var,
+               ops::FusedBatchNormV3::Attrs().IsTraining(false).DataFormat(
+                   data_format))
+        .y;
+  };
+
+  switch (epilogue) {
+    case ConvEpilogueType::kBatchNorm: {
+      result = add_bn(result, out_channels);
+      break;
+    }
+    case ConvEpilogueType::kReLU: {
+      result = ops::Relu(scope.WithOpName("relu"), result);
+      break;
+    }
+    case ConvEpilogueType::kReLUBatchnorm: {
+      result = ops::Relu(scope.WithOpName("relu"), result);
+      result = add_bn(result, out_channels);
+      break;
+    }
+    case ConvEpilogueType::kBatchnormReLU: {
+      result = add_bn(result, out_channels);
+      result = ops::Relu(scope.WithOpName("relu"), result);
+      break;
+    }
+    case ConvEpilogueType::kNone:
+      break;
+  }
+
+  if (qdq_on_output) {
+    result = AddQDQV2(scope.WithOpName("qdq_out"), result);
+  }
+  return result;
+}
+
+// Adds a batch matrix multiplication V2 operation, which commonly appears in
+// fully connected layers. Puts QDQ between (input, op) as well as between
+// (weights, op).
+ops::BatchMatMulV2 AddMatMul(Scope scope, const std::string& name,
+                             Input input) {
+  // Add QDQ to input.
+  auto input_qdq = AddQDQV2(scope, input);
+
+  // Add 3x3 weights with QDQ.
+  auto weights_const =
+      ops::Const(scope.WithOpName(name + "_weights"),
+                 {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f},
+                 TensorShape({3, 3}));
+  auto weights_qdq = AddQDQV2(scope.WithOpName("weights_qdq"), weights_const);
+  return ops::BatchMatMulV2(scope.WithOpName(name), input_qdq, weights_qdq);
+}
+}  // namespace
+
+struct QDQTestOptions {
+  bool conv_has_bias{true};
+
+  // TRT7 may have issues with optimizing redundant transpose operations between
+  // QDQ and Op introduced by TF-TRT when format is not "NCHW". This allows to
+  // test both cases as well as WAR feasibility.
+  std::string data_format{"NCHW"};
+
+  // Tests whether placing QDQ on outputs rather than inputs is handled
+  // correctly.
+  bool qdq_on_output{false};
+
+  // Option for testing whether TRT build succeeds without a final QDQ before
+  // the output.
+  bool final_qdq{true};
+
+  // Whether to add activations (relu) to conv operations
+  ConvEpilogueType conv_epilogue;
+
+  // TF-TRT API Options
+  TfTrtConversionParams conversion_params{};
+};
+
+std::ostream& operator<<(std::ostream& os, const QDQTestOptions opts) {
+  return os << absl::StrCat(
+             "QDQTestOptions(conv_has_bias=",
+             static_cast<int>(opts.conv_has_bias),
+             ", qdq_on_output=", static_cast<int>(opts.qdq_on_output),
+             ", data_format=", opts.data_format,
+             ", conv_epilogue=", DebugString(opts.conv_epilogue),
+             ", final_qdq=", opts.final_qdq, ")");
+}
+
+std::vector<QDQTestOptions> EnumerateQDQTestOptions() {
+  std::vector<QDQTestOptions> result;
+  for (const absl::string_view data_format : {"NCHW", "NHWC"}) {
+    for (auto use_bias : {true, false}) {
+      for (auto qdq_on_output : {false, true}) {
+        // For now, always append a QDQ before output. For small single-op tests
+        // (besides QDQ), TensorRT7 sometimes has trouble.
+        for (auto final_qdq : {true, false}) {
+          for (auto conv_epilogue :
+               {ConvEpilogueType::kReLU, ConvEpilogueType::kNone,
+                ConvEpilogueType::kBatchnormReLU}) {
+            // Currently batch norm converter only supports NHWC.
+            if (data_format == "NHWC" &&
+                (conv_epilogue == ConvEpilogueType::kBatchnormReLU ||
+                 conv_epilogue == ConvEpilogueType::kBatchNorm ||
+                 conv_epilogue == ConvEpilogueType::kBatchnormReLU)) {
+              continue;
+            }
+            QDQTestOptions opts{};
+            opts.conv_has_bias = use_bias;
+            opts.data_format = data_format;
+            opts.qdq_on_output = qdq_on_output;
+            opts.final_qdq = final_qdq;
+            opts.conv_epilogue = conv_epilogue;
+            result.push_back(opts);
+          }
+        }
+      }
+    }
+  }
+  return result;
+}
+
+// This class is a test fixture for running graph conversion and evaluating
+// numerical results.
+class QDQExplicitTest : public ::testing::Test,
+                        public ::testing::WithParamInterface<QDQTestOptions> {
+ public:
+  static StatusOr<PartialTensorShape> GetShape(const std::string& name,
+                                               const GraphShapeInfo& shapes) {
+    TRT_ENSURE(shapes.find(name) != shapes.end());
+    TRT_ENSURE(shapes.at(name).size() == 1);
+    return shapes.at(name)[0].shape;
+  }
+
+  StatusOr<MetaGraphDef> GetModel(const GraphDef& graph_def,
+                                  const std::vector<const NodeDef*>& inputs,
+                                  const std::vector<const NodeDef*>& outputs,
+                                  const GraphShapeInfo& shapes) {
+    TRT_ENSURE(!inputs.empty());
+    TRT_ENSURE(!outputs.empty());
+
+    MetaGraphDef out;
+    out.mutable_graph_def()->CopyFrom(graph_def);
+
+    SignatureDef signature_def;
+    auto& mutable_inputs = *signature_def.mutable_inputs();
+    for (int i = 0; i < inputs.size(); i++) {
+      std::string input_name = inputs[i]->name();
+      auto& input = mutable_inputs[input_name];
+      input.set_name(input_name);
+      input.set_dtype(DT_FLOAT);
+      TRT_ENSURE(shapes.find(input_name) != shapes.end());
+      TRT_ENSURE(shapes.at(input_name).size() == 1);
+      PartialTensorShape input_shape = shapes.at(input_name)[0].shape;
+      input_shape.AsProto(input.mutable_tensor_shape());
+    }
+
+    auto& mutable_outputs = *signature_def.mutable_outputs();
+    for (int i = 0; i < outputs.size(); i++) {
+      std::string output_name = outputs[i]->name();
+      auto& output = mutable_outputs[output_name];
+      output.set_name(output_name);
+      output.set_dtype(DT_FLOAT);
+      TRT_ENSURE(shapes.find(output_name) != shapes.end());
+      TRT_ENSURE(shapes.at(output_name).size() == 1);
+      PartialTensorShape output_shape = shapes.at(output_name)[0].shape;
+      output_shape.AsProto(output.mutable_tensor_shape());
+    }
+
+    (*out.mutable_signature_def())["serving_default"] = signature_def;
+    return out;
+  }
+
+  // Confirms that we have a TRT node with the correct attributes.
+  static Status CheckTrtNode(const GraphDef& converted_graph_def) {
+    int n_trt_ops = 0;
+    string op_name{"TRTEngineOp"};
+    for (const auto& node : converted_graph_def.node()) {
+      if (op_name == node.op()) {
+        n_trt_ops++;
+        const auto& attr = node.attr();
+        TRT_ENSURE(attr.at("static_engine").b());
+        VLOG(2) << "Found serialized segment with size "
+                << attr.at("serialized_segment").s().size();
+        TRT_ENSURE(!attr.at("serialized_segment").s().empty());
+      }
+    }
+    TRT_ENSURE(n_trt_ops == 1);
+    return OkStatus();
+  }
+
+  Status ConvertAndRun(Scope* scope) {
+    std::vector<const NodeDef*> inputs;
+    std::vector<const NodeDef*> outputs;
+
+    GraphDef gdef;
+    TF_RETURN_IF_ERROR(scope->ToGraphDef(&gdef));
+
+    std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+    TF_RETURN_IF_ERROR(scope->ToGraph(graph.get()));
+
+    GraphShapeInfo shape_info;
+    TF_RETURN_IF_ERROR(InferShapes(graph.get(), /*arg_shapes=*/{},
+                                   /*fnlib_def=*/nullptr, &shape_info));
+
+    for (const NodeDef& node : gdef.node()) {
+      if (absl::StartsWith(node.name(), "input_")) {
+        inputs.push_back(&node);
+      } else if (absl::StartsWith(node.name(), "output_")) {
+        outputs.push_back(&node);
+      }
+    }
+
+    StatusOr<MetaGraphDef> meta_graph_def =
+        GetModel(gdef, inputs, outputs, shape_info);
+    TRT_ENSURE_OK(meta_graph_def);
+
+    // Create a list of input tensors, they will be used to build the engines.
+    std::vector<Tensor> input_tensors;
+    std::vector<std::string> input_names;
+    for (const auto& input : inputs) {
+      input_names.push_back(input->name());
+
+      StatusOr<PartialTensorShape> input_shape =
+          GetShape(input->name(), shape_info);
+      TRT_ENSURE_OK(input_shape);
+
+      TensorShape shape;
+      input_shape->AsTensorShape(&shape);
+      Tensor tensor(DT_FLOAT, shape);
+      test::FillIota(&tensor, 1.0f);
+      input_tensors.push_back(tensor);
+    }
+
+    std::vector<std::string> output_names;
+    for (const auto& output : outputs) {
+      output_names.push_back(output->name());
+    }
+
+    TfTrtConversionParams conversion_params;
+    conversion_params.allow_build_at_runtime = true;
+    conversion_params.precision_mode = TrtPrecisionMode::INT8;
+    conversion_params.use_calibration = false;
+    conversion_params.convert_to_static_engine = true;
+    TRT_ENSURE(input_names.size() == input_tensors.size());
+    StatusOr<GraphDef> converted_gdef = tensorrt::ConvertAndBuild(
+        meta_graph_def->graph_def(), input_names, output_names, {input_tensors},
+        conversion_params);
+    TRT_ENSURE_OK(converted_gdef);
+    return CheckTrtNode(*converted_gdef);
+  }
+
+ protected:
+  TfTrtConversionParams params_;
+  TrtUniquePtrType<nvinfer1::ICudaEngine> engine_;
+};
+
+class TestQDQSuite : public QDQExplicitTest {};
+
+#define EXPECT_QDQ_ON_OUTPUT_FAILURE(params, scope)                  \
+  if ((params).qdq_on_output) {                                      \
+    EXPECT_THAT(ConvertAndRun(&(scope)), StatusIs(error::INTERNAL)); \
+    return;                                                          \
+  }
+#define EXPECT_NO_FINAL_QDQ_FAILURE(params, scope)                   \
+  if (!(params).final_qdq) {                                         \
+    EXPECT_THAT(ConvertAndRun(&(scope)), StatusIs(error::INTERNAL)); \
+    return;                                                          \
+  }
+
+#define EXPECT_BUILD_OK(scope) TF_EXPECT_OK(ConvertAndRun(&(scope)))
+
+#define POLICY_TRT7(params, scope)               \
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0)) {          \
+    EXPECT_QDQ_ON_OUTPUT_FAILURE(params, scope); \
+    EXPECT_NO_FINAL_QDQ_FAILURE(params, scope);  \
+    EXPECT_BUILD_OK(scope);                      \
+  }
+
+#define POLICY_TRT8(params, scope)                                          \
+  if (IS_TRT_VERSION_GE(8, 0, 0, 0)) {                                      \
+    if (((params).conv_epilogue == ConvEpilogueType::kBatchNorm ||          \
+         (params).conv_epilogue == ConvEpilogueType::kBatchnormReLU ||      \
+         (params).conv_epilogue == ConvEpilogueType::kReLUBatchnorm) &&     \
+        (params).data_format == "NHWC") {                                   \
+      EXPECT_THAT(ConvertAndRun(&(scope)), StatusIs(error::UNIMPLEMENTED)); \
+      return;                                                               \
+    }                                                                       \
+    EXPECT_BUILD_OK(scope);                                                 \
+  }
+
+#define SKIP_TRT7(x)                           \
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && (x)) { \
+    GTEST_SKIP();                              \
+  }
+
+// Tests single convolution operation conversion.
+TEST_P(TestQDQSuite, TestConv2DBasic) {
+  SKIP_TRT7(GetParam().qdq_on_output);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+  SKIP_TRT7(!GetParam().final_qdq);
+
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, GetParam().data_format, {3, 28, 28});
+
+  Output out = input;
+  const int num_conv = 1;
+  std::array<int, 2> in_channels = {3, 16};
+  std::array<int, 2> out_channels = {16, 32};
+  for (int i = 0; i < num_conv; i++) {
+    out = AddConv2D(scope.WithOpName(absl::StrCat("conv_", i)), out,
+                    in_channels[i], out_channels[i], /*filter_size=*/{3, 3},
+                    /*stride=*/{1, 1}, GetParam().data_format,
+                    GetParam().conv_has_bias, GetParam().conv_epilogue,
+                    GetParam().qdq_on_output);
+  }
+  out = AddOutput(scope, out, 0, GetParam().final_qdq);
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+// Tests single convolution operation conversion.
+TEST_P(TestQDQSuite, TestMatMulBasic) {
+  // Some param's don't apply, so pick one combination and skip otherwise.
+  if (GetParam().data_format != "NCHW" || !GetParam().conv_has_bias ||
+      GetParam().qdq_on_output ||
+      GetParam().conv_epilogue != ConvEpilogueType::kReLU) {
+    GTEST_SKIP();
+  }
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, "NHW");
+  auto matmul_op = AddMatMul(scope, "matmul", input);
+  auto out = AddOutput(scope, matmul_op, 0, GetParam().final_qdq);
+
+  TF_EXPECT_OK(ConvertAndRun(&scope));
+}
+
+// A single input goes through two different Conv2D. Outputs of Conv2D are
+// added together, with QQQ on both branches of ADD.
+TEST_P(TestQDQSuite, AddBothBranchesQDQConvSingleInput) {
+  SKIP_TRT7(!GetParam().final_qdq);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+
+  Scope scope = Scope::NewRootScope();
+  auto input1 = AddInput(scope, 0, GetParam().data_format,
+                         /*size_chw=*/{3, 28, 28});
+
+  auto conv1 =
+      AddConv2D(scope, input1, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/{1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+
+  auto conv2 =
+      AddConv2D(scope, input1, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/
+                {1, 1}, GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+
+  // In the case of "qdq on output", we don't need to add QDQ.
+  auto add =
+      ops::Add(scope.WithOpName("add"),
+               !GetParam().qdq_on_output ? AddQDQV2(scope, conv1) : conv1,
+               !GetParam().qdq_on_output ? AddQDQV2(scope, conv2) : conv2);
+
+  auto conv3 =
+      AddConv2D(scope.WithOpName("conv3"), conv2, 16, 16, {1, 1}, {1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+
+  auto out =
+      AddOutput(scope.WithOpName("output"), conv3, 0, GetParam().final_qdq);
+
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+// Tests adding a single tensor to itself, with QQQ on both branches of ADD.
+TEST_P(TestQDQSuite, AddBothBranchesQDQMultipleInput) {
+  // TRT7 QDQ optimizer makes single-input restriction.
+  SKIP_TRT7(true);
+
+  Scope scope = Scope::NewRootScope();
+  auto input1 = AddInput(scope, 0, GetParam().data_format);
+  auto input2 = AddInput(scope, 1, GetParam().data_format);
+  auto add =
+      ops::Add(scope.WithOpName("add"),
+               !GetParam().qdq_on_output ? AddQDQV2(scope, input1) : input1,
+               !GetParam().qdq_on_output ? AddQDQV2(scope, input2) : input2);
+  auto output = AddOutput(scope, add, 0, true);
+  TF_EXPECT_OK(ConvertAndRun(&scope));
+}
+
+// Tests Conv-MaxPool combination
+TEST_P(TestQDQSuite, TestConvMaxpool) {
+  SKIP_TRT7(!GetParam().final_qdq);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, GetParam().data_format,
+                        /*size_chw=*/{3, 28, 28});
+  auto conv1 =
+      AddConv2D(scope, input, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/{1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+  ops::MaxPool maxpool =
+      ops::MaxPool(scope.WithOpName("maxpool"),
+                   AddQDQV2(scope.WithOpName("mp_qdq_in"), conv1), {1, 1, 1, 1},
+                   {1, 1, 1, 1}, "SAME",
+                   ops::MaxPool::Attrs().DataFormat(GetParam().data_format));
+  auto output =
+      AddOutput(scope.WithOpName("output"), maxpool, 0, GetParam().final_qdq);
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+// Tests QDQ(Conv(QDQ(MaxPool(Conv(QDQ(x))))))
+TEST_P(TestQDQSuite, TestConvMaxpoolConv) {
+  SKIP_TRT7(!GetParam().final_qdq);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, GetParam().data_format,
+                        /*size_chw=*/{3, 28, 28});
+  auto conv1 =
+      AddConv2D(scope, input, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/{1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+  ops::MaxPool maxpool =
+      ops::MaxPool(scope.WithOpName("maxpool"),
+                   AddQDQV2(scope.WithOpName("mp_qdq_in"), conv1), {1, 1, 1, 1},
+                   {1, 1, 1, 1}, "SAME",
+                   ops::MaxPool::Attrs().DataFormat(GetParam().data_format));
+  auto conv2 = AddConv2D(scope, maxpool, 16, 16, {3, 3}, {1, 1},
+                         GetParam().data_format, GetParam().conv_has_bias,
+                         GetParam().conv_epilogue, GetParam().qdq_on_output);
+  auto output =
+      AddOutput(scope.WithOpName("out"), conv2, 0, GetParam().final_qdq);
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+INSTANTIATE_TEST_SUITE_P(TestQDQSuiteInst, TestQDQSuite,
+                         ::testing::ValuesIn(EnumerateQDQTestOptions()));
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // IS_TRT_VERSION_GE(8, 0, 0, 0)
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/selectv2.cc

@@ -0,0 +1,220 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+/*  The ConvertSelectV2 is working only for cond_input passed as a boolean
+ *  tensor, which could be created only for TRT >= 8.2.
+ */
+class ConvertSelectBase : public OpConverterBase<ConvertSelectBase> {
+ public:
+  explicit ConvertSelectBase(const OpConverterParams* params,
+                             const std::string& layer_name)
+      : OpConverterBase<ConvertSelectBase>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}),
+        layer_name_(layer_name) {}
+
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return std::array<InputArgSpec, 3>{
+        InputArgSpec::Create("cond", TrtInputArg::kBoth),
+        InputArgSpec::Create("then", TrtInputArg::kBoth),
+        InputArgSpec::Create("else", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+
+    const auto& params = *this->params_;
+    const auto& inputs = params.inputs;
+    const auto& i_cond = inputs.at(0);
+    const auto& node = params.node_def;
+    TF_RETURN_IF_ERROR(
+        check_type(i_cond.TrtDType(), nvinfer1::DataType::kBOOL, node));
+
+    if (i_cond.is_weights()) {
+      // According to
+      // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#constant-layer
+      // Boolean weights are not supported in TRT version 8.4.
+      return errors::InvalidArgument(bool_weight_error_msg(node));
+    }
+
+    const auto& i_then = inputs.at(1);
+    const auto& i_else = inputs.at(2);
+    const auto type_then = i_then.TrtDType();
+    const auto type_else = i_else.TrtDType();
+    if (type_then != type_else && (type_then == nvinfer1::DataType::kINT32 ||
+                                   type_else == nvinfer1::DataType::kINT32)) {
+      // Both or none of (type_then, type_else) should be equal to kINT32.
+      return errors::InvalidArgument(
+          then_else_dtypes_error_msg(type_then, type_else, node));
+    }
+
+    bool cond_is_vector = false;
+    const auto& shape_cond = i_cond.GetTrtDims();
+    if (layer_name_ == "select") {
+      const auto& shape_then = i_then.GetTrtDims();
+      const auto& shape_else = i_else.GetTrtDims();
+      TF_RETURN_IF_ERROR(compare_shapes(shape_then, shape_else));
+      TF_RETURN_IF_ERROR(
+          compare_shapes(shape_cond, shape_then, &cond_is_vector));
+    }
+
+    nvinfer1::Dims cond_dims(shape_cond);
+    if (cond_is_vector) {
+      cond_dims.nbDims = i_then.GetTrtDims().nbDims;
+      const std::vector<int> ones(cond_dims.d[0], 1);
+      std::copy(ones.begin(), ones.end(), cond_dims.d + 1);
+    }
+
+    const TRT_TensorOrWeights new_cond(nvinfer1::DataType::kBOOL, cond_dims,
+                                       i_cond.batch_size());
+    nvinfer1::Dims broadcasted_dims[3];
+    for (int i = 1; i < 3; i++) {
+      TF_RETURN_IF_ERROR(GetTrtBroadcastShape(new_cond, inputs.at(i), true,
+                                              false, broadcasted_dims,
+                                              broadcasted_dims + i));
+    }
+
+    for (int i = 0; i < tensor_.size(); i++) {
+      // This will also convert constants to tensors.
+      tensor_[i] = std::make_unique<TRT_TensorOrWeights>(inputs.at(i));
+      TF_RETURN_IF_ERROR(
+          ApplyBroadcast(tensor_[i], broadcasted_dims[i], this->params_, 0));
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& params = *this->params_;
+    auto* converter = params.converter;
+
+    nvinfer1::ISelectLayer* select_layer = converter->network()->addSelect(
+        *tensor_[0].get()->as_tensor(params_)->trt_tensor(),  // cond_tensor
+        *tensor_[1].get()->as_tensor(params_)->trt_tensor(),  // then_tensor
+        *tensor_[2].get()->as_tensor(params_)->trt_tensor()   // else_tensor
+    );
+
+    converter->SetLayerName(select_layer, params.node_def.name(), layer_name_);
+    AddOutput(TRT_TensorOrWeights(select_layer->getOutput(0)));
+    return OkStatus();
+  }
+
+ private:
+  Status compare_shapes(const nvinfer1::Dims& shape1,
+                        const nvinfer1::Dims& shape2,
+                        bool* cond_is_vector = nullptr) const {
+    const bool then_vs_else = cond_is_vector == nullptr;
+    bool same_shapes = shape1 == shape2;
+    if (!same_shapes && shape1.nbDims == shape2.nbDims) {
+      // We can't check size equivalent when dynamic shapes are involved.
+      // In this case, the two shapes should be equal at runtime. Therefore,
+      // the shapes still should be considered as equal if at least one of
+      // them is a tensor with dynamic shape,
+      same_shapes = DynamicShapeInput(this->params_->inputs, then_vs_else);
+    }
+    if (!same_shapes) {
+      if (then_vs_else || !(*cond_is_vector = (shape1.nbDims == 1 &&
+                                               shape1.d[0] == shape2.d[0]))) {
+        const auto err = input_shapes_error_msg(
+            shape1, shape2, this->params_->node_def, then_vs_else);
+        return errors::InvalidArgument(err);
+      }
+    }
+    return OkStatus();
+  }
+
+  bool DynamicShapeInput(const std::vector<TRT_TensorOrWeights>& inputs,
+                         bool then_vs_else) const {
+    const int idx = then_vs_else ? 1 : 0;
+    for (int i = 0; i < 2; ++i) {
+      const auto& input = inputs.at(i + idx);
+      if (input.is_tensor() && !HasStaticShape(input.GetTrtDims())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  std::array<std::unique_ptr<TRT_TensorOrWeights>, 3> tensor_;
+  const std::string layer_name_;
+};
+
+class ConvertSelect : public ConvertSelectBase {
+ public:
+  explicit ConvertSelect(const OpConverterParams* params)
+      : ConvertSelectBase(params, "select") {}
+};
+
+class ConvertSelectV2 : public ConvertSelectBase {
+ public:
+  explicit ConvertSelectV2(const OpConverterParams* params)
+      : ConvertSelectBase(params, "selectv2") {}
+};
+
+std::string op_node_info(const NodeDef& node) {
+  return " of the '" + node.op() + "' operation at the node '" + node.name() +
+         "' ";
+}
+
+std::string bool_weight_error_msg(const NodeDef& node) {
+  return "The boolean parameter '" + node.input(0) + "'" + op_node_info(node) +
+         "cannot be passed as a weight in TRT version 8.4.";
+}
+
+std::string then_else_dtypes_error_msg(nvinfer1::DataType type_then,
+                                       nvinfer1::DataType type_else,
+                                       const NodeDef& node) {
+  return "DataTypes (" + DebugString(type_then) + ", " +
+         DebugString(type_else) + ") of parameters (" + node.input(1) + ", " +
+         node.input(2) + ")" + op_node_info(node) + "are incompatible.";
+}
+
+std::string input_shapes_error_msg(const nvinfer1::Dims& shape1,
+                                   const nvinfer1::Dims& shape2,
+                                   const NodeDef& node, bool then_vs_else) {
+  const std::string& param_names =
+      then_vs_else ? "'then' and 'else'" : "'cond' and 'then'";
+  std::string error_msg = "The shapes of the " + param_names + " parameters" +
+                          op_node_info(node) + "must be the same";
+  if (!then_vs_else) {
+    error_msg +=
+        " OR 'cond' must be a vector with N elements, "
+        "where N is a batch size (the first shape dimension for 'then')";
+  }
+  return error_msg + ", got " + DebugString(shape1) + " vs. " +
+         DebugString(shape2) + ".";
+}
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertSelect>(),
+                                  "Select");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertSelectV2>(),
+                                  "SelectV2");
+#endif
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/slice_ops.cc

@@ -0,0 +1,300 @@
+/* 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/convert/ops/slice_ops.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <bitset>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_tensor_proxy.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/platform/errors.h"
+#include "tensorflow/core/util/strided_slice_op.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+// Adds a set of operations to the network which set the parameters for the
+// given "slice_layer" in order to handle dynamic input shape.
+Status HandleDynamicStridedSliceInput(
+    TRTNetworkBuilder* builder, nvinfer1::ISliceLayer* slice_layer,
+    const StridedSliceShapeSpec& strided_slice_spec,
+    const absl::InlinedVector<int64, 4>& dynamic_input_size_indices,
+    nvinfer1::Dims begin_dims, nvinfer1::Dims stride_dims,
+    nvinfer1::Dims end_dims);
+
+Status ConvertStridedSliceHelper(
+    const OpConverterParams* params, const TRT_TensorOrWeights& input,
+    const PartialTensorShape& input_dims, const SliceDims& begin,
+    const SliceDims& stride, const SliceDims& end,
+    std::optional<nvinfer1::Dims> final_shape, std::optional<int> op_instance,
+    std::optional<StridedSliceShapeSpec> strided_slice_spec) {
+  const auto& node_def = params->node_def;
+
+  auto begin_dims = DimsAdapter::Create(begin, params->use_implicit_batch);
+  auto stride_dims = DimsAdapter::Create(stride, params->use_implicit_batch);
+  auto end_dims = DimsAdapter::Create(end, params->use_implicit_batch);
+  TRT_ENSURE_OK(begin_dims);
+  TRT_ENSURE_OK(stride_dims);
+  TRT_ENSURE_OK(end_dims);
+
+  // For each dimension, gather information about static vs dynamic dimension
+  // and slice size.
+  nvinfer1::Dims size_dims = begin_dims->AsTrtDims();
+  absl::InlinedVector<int64, 4> static_input_size_indices;
+  absl::InlinedVector<int64, 4> dynamic_input_size_indices;
+  for (int i = 0; i < begin_dims->NumDims(); i++) {
+    size_dims.d[i] = (std::abs(end_dims->dim(i) - begin_dims->dim(i)) +
+                      std::abs(stride_dims->dim(i)) - 1) /
+                     std::abs(stride_dims->dim(i));
+    // When begin tensor has negative values, currently range can't be computed.
+    if (begin_dims->dim(i) < 0) {
+      return errors::Unimplemented(
+          "Negative values in begin weight tensor are unsupported");
+    }
+    if (input_dims.dim_size(i) < 0) {
+      // end_dims and begin_dims do not have valid information yet.
+      dynamic_input_size_indices.push_back(i);
+    } else {
+      static_input_size_indices.push_back(i);
+      if (end_dims->dim(i) < begin_dims->dim(i) && stride_dims->dim(i) > 0) {
+        return errors::InvalidArgument(
+            "\"size\" cannot be negative for StridedSlice");
+      }
+    }
+  }
+
+  if (!dynamic_input_size_indices.empty()) {
+    if (strided_slice_spec == std::nullopt) {
+      return errors::InvalidArgument(
+          "The argument `strided_slice_spec` is "
+          "`std::nullopt` with `dynamic_input_size_indices` non empty.");
+    }
+    if (params->use_implicit_batch) {
+      return errors::InvalidArgument(
+          "In implicit batch mode, dynamic input size is not supported.");
+    }
+  }
+
+  if (params->validation_only) return OkStatus();
+
+  StatusOr<TRTNetworkBuilder> builder = TRTNetworkBuilder::Create(
+      params->converter->network(), params->weight_store);
+  TRT_ENSURE_OK(builder);
+
+  // Create the slice operation. For dynamic dims, the inputs of the operations
+  // may be reassigned later.
+  StatusOr<nvinfer1::ISliceLayer*> slice =
+      builder->Slice(input.tensor()->trt_tensor(), begin_dims->AsTrtDims(),
+                     size_dims, stride_dims->AsTrtDims());
+  TRT_ENSURE_PTR_OK(slice);
+
+  // Handle dynamic input shapes.
+  if (!dynamic_input_size_indices.empty()) {
+    TF_RETURN_IF_ERROR(HandleDynamicStridedSliceInput(
+        &*builder, *slice, *strided_slice_spec, dynamic_input_size_indices,
+        begin_dims->AsTrtDims(), stride_dims->AsTrtDims(),
+        end_dims->AsTrtDims()));
+  }
+
+  params->converter->SetLayerName(*slice, params->node_def, "slice",
+                                  op_instance);
+  ITensorProxyPtr tensor = (*slice)->getOutput(0);
+
+  // Reshape for shrink axis, ellipsis masks based on the shape computed by
+  // ValidateStridedSliceOp or HandleDynamicStridedSliceInput.
+  nvinfer1::Dims dims = tensor->trt_tensor()->getDimensions();
+  std::vector<int> slice_input_dims(dims.d, dims.d + dims.nbDims);
+  StridedSliceShapeSpec empty_spec;
+  empty_spec.shrink_axis_dense_mask = 0;
+  auto shrink_axis_mask =
+      strided_slice_spec.value_or(empty_spec).shrink_axis_dense_mask;
+  if (final_shape) {
+    if (shrink_axis_mask) {
+      int shrink_idx = params->use_implicit_batch ? 1 : 0;
+      const auto bShrink_axis_mask = std::bitset<32>(shrink_axis_mask);
+      for (int idx = 0; idx < slice_input_dims.size(); ++idx, ++shrink_idx) {
+        const bool shrink_axis = bShrink_axis_mask[shrink_idx];
+        if (shrink_axis) {
+          slice_input_dims[idx] = 0;
+        }
+      }
+      TF_RETURN_IF_ERROR(params->converter->SqueezeTensor(
+          tensor, &slice_input_dims, params, &tensor, op_instance));
+    } else {
+      /* To do: pmajety:
+            Remove the else condition when shrink_axis_mask is always defined */
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params->converter, TRT_TensorOrWeights(tensor), *final_shape,
+          /*validation_only=*/false, &tensor, node_def, op_instance));
+    }
+  }
+  params->outputs->push_back(TRT_TensorOrWeights(tensor));
+  return OkStatus();
+}
+
+Status HandleDynamicStridedSliceInput(
+    TRTNetworkBuilder* builder, nvinfer1::ISliceLayer* slice_layer,
+    const StridedSliceShapeSpec& strided_slice_spec,
+    const absl::InlinedVector<int64, 4>& dynamic_input_size_indices,
+    nvinfer1::Dims begin_dims, nvinfer1::Dims stride_dims,
+    nvinfer1::Dims end_dims) {
+  TRT_ENSURE(builder);
+  TRT_ENSURE(slice_layer);
+
+  nvinfer1::ITensor* input_tensor = slice_layer->getInput(0);
+  TRT_ENSURE(input_tensor);
+
+  // When begin_mask or end_mask are set, we have to disregard the begin_tensor
+  // and end_tensor values. In static indices cases, ValidateStridedSliceOp
+  // returns the correct begin_tensor and end_tensor values, however with
+  // dynamic indices the correct shape has to be computed.
+
+  VLOG(3) << "begin_dims before: " << DebugString(begin_dims);
+  VLOG(3) << "end_dims before: " << DebugString(end_dims);
+  const auto begin_mask = std::bitset<32>(strided_slice_spec.begin_dense_mask);
+  const auto end_mask = std::bitset<32>(strided_slice_spec.end_dense_mask);
+  const auto shrink_axis_mask =
+      std::bitset<32>(strided_slice_spec.shrink_axis_dense_mask);
+  nvinfer1::Dims dims = input_tensor->getDimensions();
+
+  for (int idx = 0; idx < dims.nbDims; ++idx) {
+    VLOG(3) << "begin_mask[" << idx << "]: " << begin_mask[idx];
+    VLOG(3) << "end_mask[" << idx << "]: " << end_mask[idx];
+    VLOG(3) << "shrink_mask[" << idx << "]: " << shrink_axis_mask[idx];
+    if (begin_mask[idx]) {
+      begin_dims.d[idx] = 0;
+    }
+    if (end_mask[idx]) {
+      end_dims.d[idx] = dims.d[idx];
+    }
+    if (shrink_axis_mask[idx]) {
+      end_dims.d[idx] = begin_dims.d[idx] + 1;
+    }
+  }
+
+  VLOG(2) << "begin_dims after shrink_axis_mask correction: "
+          << DebugString(begin_dims);
+  VLOG(2) << "end_dims after shrink_axis_mask correction: "
+          << DebugString(end_dims);
+
+  // For each dynamic input dimension of the input, do some preprocessing based
+  // on whether this dimension is set in "begin_mask" or "end_mask" and the sign
+  // of the dimension's stride value.
+  // When stride is negative:
+  //   - If "begin_mask[dynamic_idx]" is set, then we need to adjust the slice
+  //     start of dimension[i] to the dynamic size.
+  //   - If "end_mask[dynamic_idx]" is set, it suffices to set
+  //     end_dims[dynamic_idx] to -1.
+  // When stride is positive:
+  //   - If "begin_mask[dynamic_idx]" is set, it suffices to set
+  //     begin_dims[dynamic_idx] to zero.
+  //   - If "end_mask[dynamic_idx]" is set, we need to adjust slice end to the
+  //     dynamic size of dimension "dynamic_idx".
+  absl::InlinedVector<int64, 4> dynamic_begin_indices;
+  absl::InlinedVector<int64, 4> dynamic_end_indices;
+
+  for (int i = 0; i < dynamic_input_size_indices.size(); i++) {
+    auto dynamic_idx = dynamic_input_size_indices[i];
+    if (begin_mask[dynamic_idx]) {
+      begin_dims.d[dynamic_idx] = 0;
+      if (stride_dims.d[dynamic_idx] < 0) {
+        dynamic_begin_indices.push_back(dynamic_idx);
+      }
+    }
+    if (end_mask[dynamic_idx] && !shrink_axis_mask[dynamic_idx]) {
+      end_dims.d[dynamic_idx] = stride_dims.d[dynamic_idx] > 0 ? 0 : -1;
+      if (stride_dims.d[dynamic_idx] > 0) {
+        dynamic_end_indices.push_back(dynamic_idx);
+      }
+    }
+  }
+
+  VLOG(2) << " Dynamic begin indices: " << DebugString(dynamic_begin_indices)
+          << " Dynamic end indices: " << DebugString(dynamic_end_indices);
+
+  // Create ITensors for each of the begin/stride/end constants.
+  StatusOr<nvinfer1::IConstantLayer*> begin_const = builder->Constant(
+      std::vector<int>(begin_dims.d, begin_dims.d + begin_dims.nbDims));
+  TRT_ENSURE_PTR_OK(begin_const);
+  nvinfer1::ITensor* begin_tensor = (*begin_const)->getOutput(0);
+  StatusOr<nvinfer1::IConstantLayer*> stride_const = builder->Constant(
+      std::vector<int>(stride_dims.d, stride_dims.d + stride_dims.nbDims));
+  TRT_ENSURE_PTR_OK(stride_const);
+  StatusOr<nvinfer1::IConstantLayer*> end_const = builder->Constant(
+      std::vector<int>(end_dims.d, end_dims.d + end_dims.nbDims));
+  TRT_ENSURE_PTR_OK(end_const);
+  nvinfer1::ITensor* end_tensor = (*end_const)->getOutput(0);
+
+  // Make corrections based on the begin_mask/end_mask values.
+  if (dynamic_end_indices.size() > 0) {
+    StatusOr<nvinfer1::IGatherLayer*> dynamic_end_masked_tensor =
+        builder->GetPartialShapeOf(input_tensor, dynamic_end_indices,
+                                   /*sub_one=*/false);
+    TRT_ENSURE_PTR_OK(dynamic_end_masked_tensor);
+    StatusOr<nvinfer1::IElementWiseLayer*> end_corrected =
+        builder->Add((*dynamic_end_masked_tensor)->getOutput(0), end_tensor);
+    TRT_ENSURE_PTR_OK(end_corrected);
+    end_tensor = (*end_corrected)->getOutput(0);
+  }
+  if (dynamic_begin_indices.size() > 0) {
+    StatusOr<nvinfer1::IGatherLayer*> dynamic_begin_masked_tensor =
+        builder->GetPartialShapeOf(input_tensor, dynamic_begin_indices,
+                                   /*sub_one=*/true);
+    TRT_ENSURE_PTR_OK(dynamic_begin_masked_tensor);
+
+    // Add back the original "begin" values for static dimensions.
+    StatusOr<nvinfer1::IElementWiseLayer*> begin_corrected = builder->Add(
+        (*dynamic_begin_masked_tensor)->getOutput(0), begin_tensor);
+    TRT_ENSURE_PTR_OK(begin_corrected);
+    begin_tensor = (*begin_corrected)->getOutput(0);
+  }
+
+  // Calculate the final size of the slice dynamically.
+  nvinfer1::ITensor* size_tensor;
+  {
+    StatusOr<nvinfer1::IElementWiseLayer*> num =
+        builder->Sub(end_tensor, begin_tensor);
+    TRT_ENSURE_PTR_OK(num);
+    StatusOr<nvinfer1::IElementWiseLayer*> ceil_div = builder->AbsCeilDivInt(
+        (*num)->getOutput(0), (*stride_const)->getOutput(0));
+    TRT_ENSURE_PTR_OK(ceil_div);
+    size_tensor = (*ceil_div)->getOutput(0);
+  }
+
+  slice_layer->setInput(1, *begin_tensor);
+  slice_layer->setInput(2, *size_tensor);
+  slice_layer->setInput(3, *(*stride_const)->getOutput(0));
+
+  return OkStatus();
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/slice_ops.h

@@ -0,0 +1,70 @@
+/* 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_CONVERT_OPS_SLICE_OPS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/util/strided_slice_op.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+using SliceDims = absl::InlinedVector<int64, 4>;
+
+// Creates a strided slice operation using the given information. This function
+// expects that the begin, stride, and end vectors have already been validated.
+// This function converts the [begin:stride:end] specification to the TensorRT
+// [begin:stride:size] ISliceLayer specification. The following algorithm is
+// used to perform this conversion: 1) The given (input_dims,
+// [begin:stride:end]) specification is dividied into
+//  "static dimensions" and "dynamic dimensions". "Dynamic dimensions"
+//  includes all dimensions of the slice where input_dims[i] == -1.
+// 2a) If there are no dynamic dimensions, then the "begin", "stride", and
+//  "size" variables are passed to the ISLiceLayer creation as build-time
+//  constants in the form of nvinfer1::Dims objects.
+// 2b) If there are any dynamic dimensions, then the "begin", "stride", and
+//  "size" variables are treated as runtime dynamic shape Tensors in the
+//  TensorRT graph. In this case, we must calculate "size" at runtime for all
+//  dynamic dimensions, while static dimensions use the constant values.
+//
+// Note that when any dynamic indices are present (2b), the "strided_slice_spec"
+// must be specified. This structure can be obtained through the
+// "tensorflow::ValidateStridedSliceOp" function, or it can be constructed
+// directly. When the ValidateStridedSliceOp helper function is used, it will
+// also return the "begin", "stride", and "end" vectors. When all dimensions are
+// static (2a), the "strided_slice_spec" variable is not required.
+//
+// If the "final_shape" variable is specified, then a reshape operation will be
+// added to the graph to achieve this shape. The shape must be fully specified.
+//
+// "op_instance" is only required if the caller needs to pass this variable
+// through to the Converter functions optionally accept it (SetLayerName,
+// PrepareTensorForShape).
+Status ConvertStridedSliceHelper(
+    const OpConverterParams* params, const TRT_TensorOrWeights& input,
+    const PartialTensorShape& input_dims, const SliceDims& begin,
+    const SliceDims& stride, const SliceDims& end,
+    std::optional<nvinfer1::Dims> final_shape = std::nullopt,
+    std::optional<int> op_instance = std::nullopt,
+    std::optional<StridedSliceShapeSpec> strided_slice_spec = std::nullopt);
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_

tensorflow/compiler/tf2tensorrt/convert/ops/softmax.cc

@@ -0,0 +1,81 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertSoftmax : public OpConverterBase<ConvertSoftmax> {
+ public:
+  explicit ConvertSoftmax(const OpConverterParams *params)
+      : OpConverterBase<ConvertSoftmax>(params) {}
+
+  static constexpr std::array<DataType, 3> AllowedDataTypes() {
+    return {DataType::DT_FLOAT, DataType::DT_HALF};
+  }
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("logits", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+    if (!num_trt_dims && params.use_implicit_batch) {
+      return errors::InvalidArgument(
+          "TensorRT Softmax cannot apply on the batch dimension");
+    }
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    const auto &node_def = params.node_def;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+
+    // Perform Softmax operation:
+    nvinfer1::ISoftMaxLayer *layer =
+        params.converter->network()->addSoftMax(*logits_tensor->trt_tensor());
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+    params.converter->SetLayerName(layer, node_def);
+    // Tensorflow SoftMax applies softmax operation over the last dimension.
+    layer->setAxes(1 << (num_trt_dims - 1));
+
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+    params.outputs->push_back(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertSoftmax>(),
+                                  "Softmax");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/tile.cc

@@ -0,0 +1,207 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertTile : public OpConverterBase<ConvertTile> {
+ public:
+  explicit ConvertTile(const OpConverterParams *params)
+      : OpConverterBase<ConvertTile>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("input_tensor", TrtInputArg::kBoth),
+        InputArgSpec::Create("weight", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+
+    const auto &repl = inputs.at(1);
+    if (params.use_implicit_batch && repl.is_tensor()) {
+      return errors::InvalidArgument(
+          "Conversion for Tile is not implemented for multipliers "
+          "passed as a tensor in implicit batch mode.");
+    }
+
+    nvinfer1::DataType dtype;
+    const int *multiplies;
+    if (repl.is_weights()) {
+      TFTRT_CHECK_SHAPE_TENSOR(repl.weights().GetTensor());
+      dtype = repl.weights().TrtDType();
+      multiplies = repl.weights().GetPointer<int>();
+    } else {
+      dtype = repl.tensor()->getType();
+      multiplies = nullptr;
+    }
+
+    const auto &node = params.node_def;
+    TF_RETURN_IF_ERROR(check_type(dtype, nvinfer1::DataType::kINT32, node, 1));
+
+    const auto dims = inputs.at(0).GetTrtDims();
+    const auto nb_dims =
+        dims.nbDims +
+        (params.use_implicit_batch && inputs.at(0).is_tensor() ? 1 : 0);
+    if (multiplies) {
+      const int mult_numb = repl.weights().count();
+      if (mult_numb != nb_dims) {
+        return errors::InvalidArgument(
+            "The length of the replication vector (", mult_numb,
+            ") of the Tile operation in '", node.name(),
+            "' is expected to be equal to the rank of the input vector (",
+            nb_dims, ").");
+      }
+
+      if (std::any_of(multiplies, multiplies + nb_dims,
+                      [](int i) { return i <= 0; })) {
+        const auto &mul = absl::StrJoin(multiplies, multiplies + nb_dims, ", ");
+        return errors::InvalidArgument(
+            "All replications of the Tile operation in '", node.name(),
+            "' should be positive, got (", mul, ").");
+      }
+
+      if (params.use_implicit_batch && multiplies[0] > 1) {
+        return errors::Unimplemented(
+            "The Tile operation along the batch dimension in '", node.name(),
+            "' is not implemented.");
+      }
+    } else {
+      const auto &repl_dims = repl.GetTrtDims();
+      if (repl_dims.nbDims != 1) {
+        return errors::InvalidArgument(
+            "When replications are defined as a tensor, that tensor must be "
+            "1-dimensional. Got ",
+            repl_dims.nbDims, "-dimensional tensor.");
+      }
+
+      // Check the number of elements in multiplyer for tensors with non-dynamic
+      // shape
+      if (repl_dims.d[0] >= 0 && repl_dims.d[0] != nb_dims) {
+        return errors::InvalidArgument(
+            "When replications are defined as a tensor, "
+            "the number of its elements (",
+            repl_dims.d[0], ") must be equal to the rank of the input tensor (",
+            nb_dims, ").");
+      }
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    auto *converter = params.converter;
+    auto *network = converter->network();
+    const auto &tensor = inputs.at(0);
+    const auto &replics = inputs.at(1);
+    const auto dims = tensor.GetTrtDims();
+    const auto nb_dims = dims.nbDims;
+
+    nvinfer1::Dims output_size{nb_dims, {1}};
+    bool dynamic_flag = replics.is_tensor() || !HasStaticShape(dims);
+
+    if (!dynamic_flag) {
+      // If input0 is a tensor, and we're in implicit batch mode, then we need
+      // dim_offset.
+      const auto dim_offset =
+          params.use_implicit_batch && tensor.is_tensor() ? 1 : 0;
+      const auto *input_size = dims.d;
+      const int *pReplics = replics.weights().GetPointer<int>() + dim_offset;
+      for (int i = 0; i < nb_dims; i++)
+        output_size.d[i] = pReplics[i] * input_size[i];
+    }
+
+    StatusOr<TRTNetworkBuilder> builder;
+    if (tensor.is_weights() || (dynamic_flag && replics.is_weights())) {
+      builder =
+          TRTNetworkBuilder::Create(converter->network(), params.weight_store);
+      TRT_ENSURE_OK(builder);
+    }
+
+    ITensorProxyPtr input_tensor;
+    if (tensor.is_weights()) {
+      StatusOr<nvinfer1::IConstantLayer *> weights_const =
+          builder->WeightsToConstant(tensor.weights().GetTrtWeights(), dims);
+      TRT_ENSURE_PTR_OK(weights_const);
+      input_tensor = (*weights_const)->getOutput(0);
+    } else {
+      input_tensor = tensor.tensor();
+    }
+
+    auto &input_trt_tensor = *input_tensor->trt_tensor();
+    nvinfer1::ITensor *target_shape = nullptr;
+    if (dynamic_flag) {
+      nvinfer1::ITensor *mult;
+      if (replics.is_weights()) {
+        StatusOr<nvinfer1::IConstantLayer *> weights_const =
+            builder->WeightsToConstant(replics.weights().GetTrtWeights(),
+                                       replics.GetTrtDims());
+        TRT_ENSURE_PTR_OK(weights_const);
+        mult = (*weights_const)->getOutput(0);
+      } else {
+        const ITensorProxyPtr multiplies = replics.tensor()->trt_tensor();
+        mult = multiplies->trt_tensor();
+      }
+
+      nvinfer1::ITensor *shape =
+          network->addShape(input_trt_tensor)->getOutput(0);
+      target_shape = network
+                         ->addElementWise(*shape, *mult,
+                                          nvinfer1::ElementWiseOperation::kPROD)
+                         ->getOutput(0);
+    }
+
+    nvinfer1::Dims start{nb_dims, {}};
+    DimsAdapter stride(std::vector<int>(nb_dims, 1));
+    auto layer = network->addSlice(input_trt_tensor, start, output_size,
+                                   stride.AsTrtDims());
+    layer->setMode(nvinfer1::SliceMode::kWRAP);
+    if (target_shape) layer->setInput(2, *target_shape);
+
+    converter->SetLayerName(layer, params.node_def.name(), "to_tile");
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+    if (tensor.is_weights() && params.use_implicit_batch) {
+      // Reshape output tensor by removing first dimension.
+      DimsAdapter adap(output_tensor->getDimensions());
+      TF_RETURN_IF_ERROR(adap.RemoveBatchDimension());
+
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params.converter, TRT_TensorOrWeights(output_tensor),
+          adap.AsTrtDims(), false, &output_tensor, params.node_def));
+    }
+
+    AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertTile>(), "Tile");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/unary_ops.cc

@@ -0,0 +1,251 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+const UnaryOperationMapType* UnaryOperationMap() {
+  static auto* const m = new UnaryOperationMapType({
+    {"Exp", nvinfer1::UnaryOperation::kEXP},
+        {"Log", nvinfer1::UnaryOperation::kLOG},
+        {"Sqrt", nvinfer1::UnaryOperation::kSQRT},
+        {"Rsqrt", nvinfer1::UnaryOperation::kSQRT},
+        {"Reciprocal", nvinfer1::UnaryOperation::kRECIP},
+        {"Abs", nvinfer1::UnaryOperation::kABS},
+        {"Neg", nvinfer1::UnaryOperation::kNEG},
+        {"Sin", nvinfer1::UnaryOperation::kSIN},
+        {"Cos", nvinfer1::UnaryOperation::kCOS},
+        {"Tan", nvinfer1::UnaryOperation::kTAN},
+        {"Sinh", nvinfer1::UnaryOperation::kSINH},
+        {"Cosh", nvinfer1::UnaryOperation::kCOSH},
+        {"Asin", nvinfer1::UnaryOperation::kASIN},
+        {"Acos", nvinfer1::UnaryOperation::kACOS},
+        {"Atan", nvinfer1::UnaryOperation::kATAN},
+        {"Asinh", nvinfer1::UnaryOperation::kASINH},
+        {"Acosh", nvinfer1::UnaryOperation::kACOSH},
+        {"Atanh", nvinfer1::UnaryOperation::kATANH},
+        {"Ceil", nvinfer1::UnaryOperation::kCEIL},
+        {"Floor", nvinfer1::UnaryOperation::kFLOOR},
+        {"Erf", nvinfer1::UnaryOperation::kERF},
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+        {"Round", nvinfer1::UnaryOperation::kROUND},
+        {"Sign", nvinfer1::UnaryOperation::kSIGN},
+#endif
+  });
+  return m;
+}
+
+const UnaryOperationMapType* UnaryBooleanOperationMap() {
+  static auto* const m = new UnaryOperationMapType({
+      {"LogicalNot", nvinfer1::UnaryOperation::kNOT},
+  });
+  return m;
+}
+
+const ActivationTypeMapType* ActivationTypeMap() {
+  static auto* const m = new ActivationTypeMapType({
+      {"LeakyRelu", nvinfer1::ActivationType::kLEAKY_RELU},
+      {"Relu", nvinfer1::ActivationType::kRELU},
+      {"Relu6", nvinfer1::ActivationType::kCLIP},
+      {"Sigmoid", nvinfer1::ActivationType::kSIGMOID},
+      {"Tanh", nvinfer1::ActivationType::kTANH},
+      {"Elu", nvinfer1::ActivationType::kELU},
+      {"Selu", nvinfer1::ActivationType::kSELU},
+      {"Softsign", nvinfer1::ActivationType::kSOFTSIGN},
+      {"Softplus", nvinfer1::ActivationType::kSOFTPLUS},
+  });
+  return m;
+}
+
+template <typename T>
+class ConvertUnaryImpl {
+ protected:
+  ConvertUnaryImpl(const OperationMap<T>* pOperMap) : pOperMap_(pOperMap) {}
+
+  Status ValidateImpl(const OpConverterParams& params,
+                      const std::vector<string>& not_supported_ops = {}) {
+    const auto& node = params.node_def;
+    const auto& op = node.op();
+    if (pOperMap_->find(op) == pOperMap_->end()) {
+      return errors::Unimplemented("Unary op: ", op, " not supported");
+    }
+    DimsAdapter input_dims(params.inputs.at(0).GetTrtDims());
+    if (!input_dims.NumDims()) {
+      return errors::InvalidArgument(
+          "At least 1 dimension is required for UNARY operation '", op, "'");
+    }
+
+    if (!not_supported_ops.empty() && params.use_implicit_batch) {
+      const auto& end = not_supported_ops.end();
+      if (std::find(not_supported_ops.begin(), end, op) != end) {
+        const auto& err =
+            convert_not_supported_implicit(op, node.name(), "Unary");
+        return errors::Unimplemented(err);
+      }
+    }
+
+    return OkStatus();
+  }
+
+  Status ConvertImpl(const OpConverterParams& params) {
+    const auto& node_def = params.node_def;
+    auto* converter = params.converter;
+    const auto op_pair = pOperMap_->find(node_def.op());
+    ITensorProxyPtr tensor = params.inputs.at(0).tensor();
+    nvinfer1::IUnaryLayer* layer =
+        converter->network()->addUnary(*tensor->trt_tensor(), op_pair->second);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+    converter->SetLayerName(layer, node_def);
+    if (node_def.op() == "Rsqrt") {
+      layer = converter->network()->addUnary(*layer->getOutput(0),
+                                             nvinfer1::UnaryOperation::kRECIP);
+      TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+      converter->SetLayerName(layer, node_def, "recip");
+    }
+    params.outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0)));
+    return OkStatus();
+  }
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("x", TrtInputArg::kTensor)};
+  }
+
+ protected:
+  const OperationMap<T>* pOperMap_;
+};
+
+class ConvertUnary : public OpConverterBase<ConvertUnary>,
+                     protected ConvertUnaryImpl<nvinfer1::UnaryOperation> {
+ public:
+  explicit ConvertUnary(const OpConverterParams* params)
+      : OpConverterBase<ConvertUnary>(
+            params,
+            params->node_def.op() == "Sign"
+                ? std::vector<DataType>{DataType::DT_FLOAT, DataType::DT_HALF,
+                                        DataType::DT_INT8, DT_INT32}
+                : std::vector<DataType>{DataType::DT_FLOAT, DataType::DT_HALF,
+                                        DataType::DT_INT8}),
+        ConvertUnaryImpl(UnaryOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return ConvertUnaryImpl::InputSpec();
+  }
+
+  Status Validate() { return ValidateImpl(*params_, {"Sign", "Round"}); }
+  Status Convert() { return ConvertImpl(*params_); }
+};
+
+class ConvertBooleanUnary : public OpConverterBase<ConvertBooleanUnary>,
+                            public ConvertUnaryImpl<nvinfer1::UnaryOperation> {
+ public:
+  explicit ConvertBooleanUnary(const OpConverterParams* params)
+      : OpConverterBase<ConvertBooleanUnary>(params, {DataType::DT_BOOL}),
+        ConvertUnaryImpl(UnaryBooleanOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return ConvertUnaryImpl::InputSpec();
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "LogicalNot"
+      input: "..."
+    }
+    */
+    return "";
+  }
+  Status Validate() {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    return ValidateImpl(*params_, {"LogicalNot"});
+#else
+    return errors::Unimplemented("Boolean op: ", params_->node_def.op(),
+                                 " is not supported in TRT version < 8.2");
+#endif
+  }
+  Status Convert() { return ConvertImpl(*params_); }
+};
+
+class ConvertActivation : public OpConverterBase<ConvertActivation>,
+                          protected ConvertUnaryImpl<nvinfer1::ActivationType> {
+ public:
+  explicit ConvertActivation(const OpConverterParams* params)
+      : OpConverterBase<ConvertActivation>(params),
+        ConvertUnaryImpl(ActivationTypeMap()) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("input", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(ValidateImpl(*params_));
+    const auto& node_def = params_->node_def;
+    if (node_def.op() == "LeakyRelu") {
+      return GetNodeAttr(AttrSlice(node_def), "alpha", &alpha_);
+    }
+    alpha_ = 1.0f;
+    return OkStatus();
+  }
+  Status Convert() {
+    auto* converter = params_->converter;
+    const auto& inputs = params_->inputs;
+    const auto& node_def = params_->node_def;
+    const auto& op = node_def.op();
+    const auto op_pair = pOperMap_->find(op);
+    nvinfer1::IActivationLayer* layer = converter->network()->addActivation(
+        *inputs.at(0).tensor()->trt_tensor(), op_pair->second);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+    converter->SetLayerName(layer, node_def, "activation");
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+    // Set parameters.
+    if (op == "Selu") {
+      // From tensorflow/core/kernels/relu_op_functor.h
+      alpha_ = 1.7580993408473768599402175208123f;
+      layer->setBeta(1.0507009873554804934193349852946f);
+    } else if (op == "Softplus") {
+      layer->setBeta(1.0f);
+    } else if (op == "Relu6") {
+      layer->setBeta(6.0f);
+      converter->ProvideQuantizationRange(&output_tensor, alpha_ = 0.0f, 6.0f);
+    }
+    layer->setAlpha(alpha_);
+    params_->outputs->push_back(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+
+ private:
+  float alpha_ = 0.f;
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertUnary>(),
+                                  GetOperationNames(*UnaryOperationMap()));
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertBooleanUnary>(),
+    GetOperationNames(*UnaryBooleanOperationMap()));
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertActivation>(),
+                                  GetOperationNames(*ActivationTypeMap()));
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/variable_ops.cc

@@ -0,0 +1,356 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "third_party/tensorrt/NvInfer.h"
+#include "third_party/tensorrt/NvInferRuntimeCommon.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+struct VarAttributes {
+  TensorShapeProto shape_proto;
+  TensorShape shape;
+  string name;
+  DataType dtype;
+  string shared_name;
+  string container;
+};
+
+template <typename T, bool is_resource>
+Status ReadVariableHelper(const OpConverterParams* params,
+                          const VarAttributes& attrs,
+                          TRT_ShapedWeights* weights) {
+  Tensor tensor(attrs.dtype, attrs.shape);
+  auto ctx = params->converter->context();
+  TRT_ENSURE(ctx != nullptr);
+  auto tensor_flat = tensor.flat<T>();
+
+  // Clone function library runtime in order to get a mutable library
+  // definition to add and run a function with the variable operation.
+  auto lib = ctx->function_library();
+  std::unique_ptr<FunctionLibraryDefinition> lib_def;
+  std::unique_ptr<ProcessFunctionLibraryRuntime> lib_pflr;
+  FunctionLibraryRuntime* lib_clone;  // Not owned.
+  TF_RETURN_IF_ERROR(lib->Clone(&lib_def, &lib_pflr, &lib_clone));
+
+  // Create function definition.
+  FunctionDef fdef;
+  std::vector<Tensor> args;
+  string func_name = attrs.name + "/func";
+  if (is_resource) {
+    // Create input tensor with the resource handle.
+    const auto& inputs = params->inputs;
+    const TRT_TensorOrWeights& handle = inputs.at(0);
+    args.emplace_back(handle.resource());
+
+    fdef = FunctionDefHelper::Define(
+        func_name,                                             // Name
+        {"in: resource"},                                      // Args
+        {absl::StrCat("out: ", DataTypeString(attrs.dtype))},  // Returns
+        {},                                                    // Attr def
+        // Nodes
+        {{{attrs.name},
+          "ReadVariableOp",
+          {"in"},  // Name of the Placeholder or VarHandleOp
+          {{"dtype", attrs.dtype}}},
+         {{"out"}, "Identity", {attrs.name}, {{"T", attrs.dtype}}}});
+  } else {
+    fdef = FunctionDefHelper::Define(
+        func_name,                                             // Name
+        {},                                                    // Args
+        {absl::StrCat("out: ", DataTypeString(attrs.dtype))},  // Returns
+        {},                                                    // Attr def
+        // Nodes
+        {{{attrs.name},
+          "VariableV2",
+          {},
+          {{"dtype", attrs.dtype},
+           {"shape", attrs.shape_proto},
+           {"container", attrs.container},
+           {"shared_name", attrs.shared_name}}},
+         {{"out"}, "Identity", {attrs.name}, {{"T", attrs.dtype}}}});
+  }
+
+  // Add function definition to the library.
+  TF_RETURN_IF_ERROR(lib_def->AddFunctionDef(fdef));
+
+  // Instantiate function.
+  FunctionLibraryRuntime::Handle func_handle;
+  FunctionLibraryRuntime::InstantiateOptions inst_ops;
+  inst_ops.state_handle = "";
+  inst_ops.target = ctx->device()->name();
+  AttrValueMap attr_list;
+  TF_RETURN_IF_ERROR(lib_clone->Instantiate(func_name, AttrSlice(&attr_list),
+                                            inst_ops, &func_handle));
+
+  FunctionLibraryRuntime::Options opts;
+  opts.rendezvous = ctx->rendezvous();
+  opts.cancellation_manager = ctx->cancellation_manager();
+  opts.runner = ctx->runner();
+
+  std::vector<Tensor>* rets = new std::vector<Tensor>();
+  std::unique_ptr<std::vector<Tensor>> outputs_wrapper(rets);
+
+  // Run the new function synchronously.
+  TF_RETURN_IF_ERROR(lib_clone->RunSync(opts, func_handle, args, rets));
+
+  TRT_ENSURE(ctx->op_device_context() != nullptr);
+  TRT_ENSURE(ctx->op_device_context()->stream() != nullptr);
+
+  // Copy tensor.
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(CHECK_NOTNULL(
+      ctx->op_device_context()->stream()->platform_specific_handle().stream));
+
+  auto ret = cudaMemcpyAsync(tensor_flat.data(), rets->at(0).flat<T>().data(),
+                             rets->at(0).NumElements() * sizeof(T),
+                             cudaMemcpyDeviceToHost, stream);
+  if (ret != 0) {
+    return errors::Internal("Could not copy the variable ", attrs.name);
+  }
+  cudaStreamSynchronize(stream);
+
+  TF_RETURN_IF_ERROR(
+      TfTensorToTrtWeights(tensor, params->weight_store, weights));
+
+  return OkStatus();
+}
+
+class ConvertVariableV2 : public OpConverterBase<ConvertVariableV2> {
+ public:
+  ConvertVariableV2(const OpConverterParams* params)
+      : OpConverterBase<ConvertVariableV2>(params) {}
+
+  static constexpr std::array<InputArgSpec, 0> InputSpec() { return {}; }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "VariableV2"
+      ...
+      attr {
+        key: "dtype"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      ...
+    }
+    */
+    return "dtype";
+  }
+
+  template <typename T>
+  Status ValidateImpl() {
+    const auto& node_def = params_->node_def;
+
+    // Verify and consume node attributes.
+    StatusOr<TensorShapeProto> shape_proto =
+        GetAttrValue<TensorShapeProto>("shape");
+    StatusOr<string> shared_name = GetAttrValue<string>("shared_name");
+    StatusOr<string> container = GetAttrValue<string>("container");
+    TRT_ENSURE_OK(shape_proto);
+    TRT_ENSURE_OK(shared_name);
+    TRT_ENSURE_OK(container);
+
+    attrs_.shape_proto = *shape_proto;
+    attrs_.shape = TensorShape(*shape_proto);
+    attrs_.name = node_def.name();
+    attrs_.shared_name = *shared_name;
+    attrs_.container = *container;
+
+    Tensor tensor(attrs_.dtype, attrs_.shape);
+    auto tensor_flat = tensor.flat<T>();
+    for (int64_t i = 0; i < tensor_flat.size(); i++) {
+      tensor_flat(i) = T(0.0f);
+    }
+
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(
+        TfTensorToTrtWeights(tensor, params_->weight_store, &weights));
+
+    // Only push outputs during validation and when outputs are expected.
+    if (params_->validation_only && params_->outputs != nullptr) {
+      AddOutput(TRT_TensorOrWeights(weights));
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    const auto& node_def = params_->node_def;
+    StatusOr<DataType> dtype = GetAttrValue<DataType>("dtype");
+    TRT_ENSURE_OK(dtype);
+    attrs_.dtype = *dtype;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ValidateImpl<float>();
+      case DT_HALF:
+        return ValidateImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+  template <typename T>
+  Status ConvertImpl() {
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(ReadVariableHelper<T, false>(params_, attrs_, &weights));
+    AddOutput(TRT_TensorOrWeights(weights));
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& node_def = params_->node_def;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ConvertImpl<float>();
+      case DT_HALF:
+        return ConvertImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+ private:
+  VarAttributes attrs_{};
+};
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertVariableV2>(),
+                                  {"VariableV2"});
+
+class ConvertReadVariableOp : public OpConverterBase<ConvertReadVariableOp> {
+ public:
+  ConvertReadVariableOp(const OpConverterParams* params)
+      : OpConverterBase<ConvertReadVariableOp>(params) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return {InputArgSpec::Create("resource", TrtInputArg::kResource)};
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    return "dtype";
+  }
+
+  template <typename T>
+  Status ValidateImpl() {
+    const auto& node_def = params_->node_def;
+
+    // Verify and consume node attributes.
+    StatusOr<TensorShapeProto> shape_proto =
+        GetAttrValue<TensorShapeProto>("_shape");
+    TRT_ENSURE_OK(shape_proto);
+
+    attrs_.shape_proto = *shape_proto;
+    attrs_.shape = TensorShape(*shape_proto);
+    attrs_.name = node_def.name();
+
+    Tensor tensor(attrs_.dtype, attrs_.shape);
+    auto tensor_flat = tensor.flat<T>();
+    for (int64_t i = 0; i < tensor_flat.size(); i++) {
+      tensor_flat(i) = T(0.0f);
+    }
+
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(
+        TfTensorToTrtWeights(tensor, params_->weight_store, &weights));
+
+    // Only push outputs during validation and when outputs are expected.
+    if (params_->validation_only && params_->outputs != nullptr) {
+      AddOutput(TRT_TensorOrWeights(weights));
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    const auto& node_def = params_->node_def;
+    if (params_->use_implicit_batch) {
+      return errors::Unimplemented("Implicit batch mode not supported, at ",
+                                   node_def.name());
+    }
+
+    StatusOr<DataType> dtype = GetAttrValue<DataType>("dtype");
+    TRT_ENSURE_OK(dtype);
+    attrs_.dtype = *dtype;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ValidateImpl<float>();
+      case DT_HALF:
+        return ValidateImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+  template <typename T>
+  Status ConvertImpl() {
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(ReadVariableHelper<T, true>(params_, attrs_, &weights));
+    AddOutput(TRT_TensorOrWeights(weights));
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& node_def = params_->node_def;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ConvertImpl<float>();
+      case DT_HALF:
+        return ConvertImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+ private:
+  VarAttributes attrs_{};
+};
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertReadVariableOp>(), {"ReadVariableOp"});
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/timing_cache.cc

@@ -0,0 +1,86 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/timing_cache.h"
+
+#include <unordered_map>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/core/platform/errors.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+StatusOr<TimingCacheRegistry::TimingCachePtr> TimingCacheRegistry::LookUp(
+    const string& name, nvinfer1::IBuilderConfig* builder_config) {
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+  TRT_ENSURE(builder_config != nullptr);
+  mutex_lock scoped_lock(mu_);
+  if (map_.find(name) != map_.end()) {
+    const std::vector<uint8_t>& data = map_[name];
+    return std::unique_ptr<nvinfer1::ITimingCache>(
+        builder_config->createTimingCache(data.data(), data.size()));
+  }
+
+  // If no such timing cache exists, create a new timing cache.
+  return std::unique_ptr<nvinfer1::ITimingCache>(
+      builder_config->createTimingCache(nullptr, 0));
+#endif  // IS_TRT_VERSION_GE(8, 0, 0, 0)
+  return errors::Unimplemented(
+      "serializable timing cache does not exist in TensorRT versions < 8.0");
+}
+
+void TimingCacheRegistry::Upsert(const string& name, TimingCache* cache) {
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+  nvinfer1::IHostMemory* memory = cache->serialize();
+  if (memory == nullptr) {
+    return;
+  }
+
+  if (map_.find(name) == map_.end()) {
+    // If the timing cache with the given name does not exist, emplace the
+    // serialized buffer.
+    std::vector<uint8_t> mem(memory->size());
+    std::copy_n(static_cast<uint8_t*>(memory->data()), memory->size(),
+                mem.begin());
+    {
+      mutex_lock scoped_lock(mu_);
+      map_.emplace(name, std::move(mem));
+    }
+  } else {
+    // If the timing cache does exist, use the existing buffer.
+    mutex_lock scoped_lock(mu_);
+    std::vector<uint8_t>& mem = map_[name];
+    mem.resize(memory->size());
+    std::copy_n(static_cast<uint8_t*>(memory->data()), memory->size(),
+                mem.begin());
+  }
+  memory->destroy();
+#endif  // IS_TRT_VERSION_GE(8, 0, 0, 0)
+}
+
+TimingCacheRegistry* GetTimingCacheRegistry() {
+  static TimingCacheRegistry* registry = new TimingCacheRegistry();
+  return registry;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/timing_cache.h

@@ -0,0 +1,70 @@
+/* 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_CONVERT_TIMING_CACHE_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TIMING_CACHE_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include <unordered_map>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+// A registry for holding serialized TensorRT autotuner timing caches.
+// For TensorRT versions < 8.0, the timing cache is not serializable, so these
+// operations become no-ops.
+class TimingCacheRegistry {
+ public:
+  TimingCacheRegistry() = default;
+  ~TimingCacheRegistry() = default;
+
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+  using TimingCache = nvinfer1::ITimingCache;
+  using TimingCachePtr = std::unique_ptr<TimingCache>;
+#else
+  struct TimingCache {};
+  using TimingCachePtr = std::unique_ptr<TimingCache>;
+#endif
+
+  // Insert or update a registry into the map using the given name. The cache
+  // will be serialized before being placed into the map.
+  void Upsert(const string& name, TimingCache* cache);
+
+  // Find a timing cache using the given name. The provided BuilderConfig is
+  // used to deserialize the cache. If no timing cache is found, a new timing
+  // cache is returned.
+  StatusOr<TimingCachePtr> LookUp(const string& name,
+                                  nvinfer1::IBuilderConfig* builder_config);
+
+ private:
+  using SerializedTimingCache = std::vector<uint8_t>;
+
+  mutex mu_;
+  std::unordered_map<std::string, SerializedTimingCache> map_;
+};
+
+TimingCacheRegistry* GetTimingCacheRegistry();
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TIMING_CACHE_H_

tensorflow/compiler/tf2tensorrt/convert/trt_layout_optimization_pass.cc

@@ -0,0 +1,97 @@
+/* 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/convert/trt_layout_optimization_pass.h"
+
+#include "absl/strings/ascii.h"
+#include "absl/strings/escaping.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
+#include "tensorflow/core/grappler/utils/functions.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/stacktrace.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+using absl::AsciiStrToUpper;
+using absl::StrAppend;
+using absl::StrCat;
+
+TRTLayoutOptimizationPass::TRTLayoutOptimizationPass(const string& name)
+    : name_(name),
+      trt_logger_name_("DefaultLogger"),
+      minimum_segment_size_(3),
+      is_dynamic_op_(false),
+      max_cached_batches_(1),
+      max_workspace_size_bytes_(256LL << 20) {
+  VLOG(1) << "Constructing " << name_;
+}
+
+Status TRTLayoutOptimizationPass::Optimize(grappler::Cluster* cluster,
+                                           const grappler::GrapplerItem& item,
+                                           GraphDef* optimized_graph) {
+  GraphDef modified_graph_def = item.graph;
+
+  // Construct a GrapplerItem using the modified graph_def and the input
+  // grappler_item.
+  grappler::GrapplerItem grappler_item =
+      grappler_item.WithGraph(std::move(modified_graph_def));
+  const GraphDef& graph_def = grappler_item.graph;
+
+  // Convert graphdef to graph.
+  FunctionLibraryDefinition flib(OpRegistry::Global(), graph_def.library());
+  Graph graph(flib);
+  TF_RETURN_IF_ERROR(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), graph_def, &graph));
+
+  // Algorithm steps:
+  // 1. We iterate over the graph to find any Conv (or other layout sensitive
+  // op)
+  // 2. If found, we continue, else we return
+  // 3. We iterate over the nodes and replace the layout-sensitive params
+  // 3. We add Transpose before the inputs and after the outputs
+
+  grappler::GraphProperties static_graph_properties(grappler_item);
+
+  std::cout << "TRTLayoutOptimizationPass: reading nodes..." << std::endl;
+  for (Node* node : graph.nodes()) {
+    std::cout << node->name() << std::endl;
+  }
+
+  // TODO: assign output *optimized_graph =;
+}
+
+Status TRTLayoutOptimizationPass::Init(
+    const RewriterConfig_CustomGraphOptimizer* config) {
+  std::cout << "Do nothing for now" << std::endl;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/trt_layout_optimization_pass.h

@@ -0,0 +1,69 @@
+/* Copyright 20121 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_CONVERT_TRT_LAYOUT_OPTIMIZATION_PASS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_LAYOUT_OPTIMIZATION_PASS_H_
+
+#include <string>
+
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#if !IS_TRT_VERSION_GE(7, 0, 0, 0)
+#error From version 2.6, we only support NVIDIA TensorRT version 7 or newer.
+#error Please update your environment and relaunch the compilation.
+#endif
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+class TRTLayoutOptimizationPass : public grappler::CustomGraphOptimizer {
+ public:
+  TRTLayoutOptimizationPass(const string& name = "TRTLayoutOptimizationPass");
+
+  string name() const override { return name_; };
+
+  bool UsesFunctionLibrary() const override { return true; }
+
+  Status Init(
+      const RewriterConfig_CustomGraphOptimizer* config = nullptr) override;
+
+  Status Optimize(grappler::Cluster* cluster,
+                  const grappler::GrapplerItem& item,
+                  GraphDef* optimized_graph) override;
+
+  /*  void PrintDebugInfo(grappler::Cluster* cluster,
+                        const grappler::GrapplerItem& item);
+  */
+
+ private:
+  const string name_;
+  string trt_logger_name_;
+  int minimum_segment_size_;
+  bool is_dynamic_op_;
+  int max_cached_batches_;
+  int64_t max_workspace_size_bytes_;
+};
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_LAYOUT_OPTIMIZATION_PASS_H_

tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc

@@ -0,0 +1,251 @@
+/* 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/convert/trt_optimization_pass.h"
+
+#include <memory>
+
+#include "absl/strings/ascii.h"
+#include "absl/strings/escaping.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/trt_parameters.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/op_types.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
+#include "tensorflow/core/grappler/utils/functions.h"
+#include "tensorflow/core/grappler/utils/topological_sort.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+using absl::AsciiStrToUpper;
+using absl::StrAppend;
+using absl::StrCat;
+
+namespace {
+
+bool ShouldUseExplicitPrecision(const GraphDef& gdef) {
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0)) {
+    return false;
+  }
+  return absl::c_any_of(gdef.node(), [](const auto& node) {
+    return (absl::c_find(kExplicitQuantizationOpNames, node.op()) !=
+            kExplicitQuantizationOpNames.end());
+  });
+}
+
+StatusOr<bool> ShouldConvertFunction(const grappler::GrapplerItem& item) {
+  if (item.id == "tf_graph") {
+    return false;
+  }
+  const auto& func_item =
+      tensorflow::down_cast<const grappler::GrapplerFunctionItem&>(item);
+  const AttrSlice& attr = func_item.func_attr();
+  const AttrValue* attr_value = attr.FindByString("_tftrt_convert_function");
+  if (attr_value != nullptr) {
+    bool result = false;
+    TF_RETURN_IF_ERROR(GetNodeAttr(attr, "_tftrt_convert_function", &result));
+    return result;
+  }
+  VLOG(1) << "Attribute _tftrt_convert_function was not found.";
+  return false;
+}
+
+// Converts function conversion attributes to conversion parameters.
+Status UpdateFunctionSpecificConversionParams(
+    TRTOptimizationPass::ConversionParams& cp,
+    const tensorflow::AttrSlice& attr) {
+  auto get_size_attr = [](const AttrSlice& attr, absl::string_view name,
+                          size_t* dst) -> Status {
+    int tmp = 0;
+    TF_RETURN_IF_ERROR(GetNodeAttr(attr, name, &tmp));
+    *dst = static_cast<size_t>(tmp);
+    return OkStatus();
+  };
+
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(attr, "_tftrt_trt_logger_name", &cp.trt_logger_name));
+  TF_RETURN_IF_ERROR(
+      get_size_attr(attr, "_tftrt_max_batch_size", &cp.max_batch_size));
+  TF_RETURN_IF_ERROR(get_size_attr(attr, "_tftrt_max_workspace_size_bytes",
+                                   &cp.max_workspace_size_bytes));
+  std::string precision_mode;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(attr, "_tftrt_precision_mode", &precision_mode));
+  TF_RETURN_IF_ERROR(
+      TrtPrecisionModeFromName(precision_mode, &cp.precision_mode));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attr, "_tftrt_minimum_segment_size",
+                                 &cp.minimum_segment_size));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attr, "_tftrt_is_dyn_op", &cp.is_dynamic_op));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(attr, "_tftrt_max_cached_engines", &cp.max_cached_engines));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(attr, "_tftrt_use_calibration", &cp.use_calibration));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(attr, "_tftrt_use_implicit_batch", &cp.use_implicit_batch));
+  std::string profile_strategy;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(attr, "_tftrt_profile_strategy", &profile_strategy));
+  TF_RETURN_IF_ERROR(
+      ProfileStrategyFromName(profile_strategy, &cp.profile_strategy));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attr, "_tftrt_allow_build_at_runtime",
+                                 &cp.allow_build_at_runtime));
+  return OkStatus();
+}
+}  // namespace
+
+Status TRTOptimizationPass::Init(
+    const RewriterConfig_CustomGraphOptimizer* config) {
+  if (config == nullptr) {
+    return OkStatus();
+  }
+  const auto params = config->parameter_map();
+  if (params.count("minimum_segment_size")) {
+    params_.minimum_segment_size = params.at("minimum_segment_size").i();
+  }
+  if (params.count("max_batch_size")) {
+    params_.max_batch_size = params.at("max_batch_size").i();
+  }
+  if (params.count("is_dynamic_op")) {
+    params_.is_dynamic_op = params.at("is_dynamic_op").b();
+  }
+  if (params.count("maximum_cached_engines")) {
+    params_.max_cached_engines = params.at("maximum_cached_engines").i();
+  }
+  if (params.count("max_workspace_size_bytes")) {
+    params_.max_workspace_size_bytes =
+        params.at("max_workspace_size_bytes").i();
+  }
+  if (params.count("precision_mode")) {
+    TF_RETURN_IF_ERROR(TrtPrecisionModeFromName(
+        AsciiStrToUpper(params.at("precision_mode").s()),
+        &params_.precision_mode));
+  }
+  if (params.count("use_calibration")) {
+    params_.use_calibration = params.at("use_calibration").b();
+  }
+  if (params.count("trt_logger")) {
+    params_.trt_logger_name = params.at("trt_logger").s();
+  }
+  if (params.count("allow_build_at_runtime")) {
+    params_.allow_build_at_runtime = params.at("allow_build_at_runtime").b();
+  }
+  if (params.count("use_implicit_batch")) {
+    params_.use_implicit_batch = params.at("use_implicit_batch").b();
+  }
+  if (params.count("profile_strategy")) {
+    TF_RETURN_IF_ERROR(ProfileStrategyFromName(
+        params.at("profile_strategy").s(), &params_.profile_strategy));
+  }
+  return OkStatus();
+}
+
+static bool ExplicitPrecisionModePolicy() {
+  return IS_TRT_VERSION_GE(8, 0, 0, 0);
+}
+
+Status TRTOptimizationPass::Optimize(grappler::Cluster* cluster,
+                                     const grappler::GrapplerItem& item,
+                                     GraphDef* optimized_graph) {
+  VLOG(1) << "Called TRTOptimization Pass " << name_
+          << " on a grappler item with id=" << item.id;
+  TF_ASSIGN_OR_RETURN(bool do_function_conversion, ShouldConvertFunction(item));
+  // Optimizing the main graph(identified with `item.id == "tf_graph"`) with
+  // `minimim_segment_size == -1` indicates skipping main graph conversion.
+  if ((params_.minimum_segment_size == -1 && item.id == "tf_graph") ||
+      (item.id != "tf_graph" && !do_function_conversion)) {
+    VLOG(1) << "Not optimizing this grappler item: " << item.id;
+    *optimized_graph = item.graph;
+    return OkStatus();
+  }
+
+  if (params_.use_calibration &&
+      params_.precision_mode != TrtPrecisionMode::INT8) {
+    LOG(WARNING) << "Calibration with FP32 or FP16 is not implemented. "
+                 << "Falling back to use_calibration = False."
+                 << "Note that the default value of use_calibration is True.";
+    params_.use_calibration = false;
+  }
+
+  params_.use_explicit_precision = ShouldUseExplicitPrecision(item.graph);
+  if (params_.use_explicit_precision) {
+    LOG(INFO) << "[TF-TRT] Using explicit QDQ mode";
+    if (params_.precision_mode != TrtPrecisionMode::INT8 ||
+        params_.use_calibration) {
+      LOG(WARNING)
+          << "Explicit precision mode with calibration or FP32/FP16 mode is "
+             "not supported."
+          << " Setting precision mode to INT8 and calibration to false.";
+      params_.precision_mode = TrtPrecisionMode::INT8;
+      params_.use_calibration = false;
+    }
+  }
+
+  // Create a copy of the graph to optimize.
+  grappler::GrapplerItem optimized_item(item);
+
+  std::vector<string> nodes_to_preserve;
+  const auto& old_nodes_to_preserve = item.NodesToPreserve();
+  nodes_to_preserve.reserve(old_nodes_to_preserve.size());
+  for (const auto& n : old_nodes_to_preserve) {
+    auto tokens = str_util::Split(n, ":");
+    string s = tokens.at(0);
+    for (int i = 1; i < tokens.size() - 1; ++i) {
+      StrAppend(&s, ":", tokens.at(i));
+    }
+    int dumm_port = -1;
+    // If the last token is not an integer, it must be part of the name.
+    // Otherwise it is port number.
+    if (tokens.size() > 1 &&
+        !strings::safe_strto32(tokens.back(), &dumm_port)) {  // non-absl ok
+      StrAppend(&s, ":", tokens.back());
+    }
+    nodes_to_preserve.push_back(s);
+  }
+
+  if (item.id != "tf_graph" && do_function_conversion) {
+    const grappler::GrapplerFunctionItem& func_item =
+        tensorflow::down_cast<const grappler::GrapplerFunctionItem&>(item);
+    TF_RETURN_IF_ERROR(
+        UpdateFunctionSpecificConversionParams(params_, func_item.func_attr()));
+  }
+
+  return ConvertGraph(params_, optimized_item, nodes_to_preserve, cluster,
+                      optimized_graph);
+}
+
+static grappler::CustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
+    []() {
+      VLOG(1)
+          << "Instantiating CustomOptimizationPass object TensorRTOptimizer";
+      return new TRTOptimizationPass("TensorRTOptimizer");
+    },
+    ("TensorRTOptimizer"));
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h

@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
+
+#include <memory>
+#include <string>
+
+#include "tensorflow/compiler/tf2tensorrt/convert/trt_parameters.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/grappler/costs/graph_properties.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
+#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#if !IS_TRT_VERSION_GE(7, 0, 0, 0)
+#error From version 2.6, we only support NVIDIA TensorRT version 7 or newer.
+#error Please update your environment and relaunch the compilation.
+#endif
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class TRTOptimizationPass : public grappler::CustomGraphOptimizer {
+ public:
+  struct ConversionParams {
+    string trt_logger_name = "DefaultLogger";
+    size_t max_batch_size = -1;
+    size_t max_workspace_size_bytes = 1 << 30;
+    TrtPrecisionMode precision_mode = TrtPrecisionMode::FP32;
+    int minimum_segment_size = 3;
+    // Whether to create engine on conversion or execution time
+    bool is_dynamic_op = false;
+    // maximum number of cached engines
+    int max_cached_engines = 1;
+    bool use_calibration = true;
+    bool use_implicit_batch = true;
+    ProfileStrategy profile_strategy = ProfileStrategy::kRange;
+    bool allow_build_at_runtime = true;
+    bool use_explicit_precision = false;
+  };
+
+  TRTOptimizationPass(const string& name = "TRTOptimizationPass")
+      : name_(name) {}
+
+  string name() const override { return name_; };
+
+  bool UsesFunctionLibrary() const override { return true; }
+
+  Status Init(
+      const RewriterConfig_CustomGraphOptimizer* config = nullptr) override;
+
+  Status Optimize(grappler::Cluster* cluster,
+                  const grappler::GrapplerItem& item,
+                  GraphDef* optimized_graph) override;
+
+ private:
+  const string name_;
+
+  ConversionParams params_;
+
+  std::vector<int> batches_;
+};
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_

tensorflow/compiler/tf2tensorrt/convert/trt_parameters.cc

@@ -0,0 +1,102 @@
+/* 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/convert/trt_parameters.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <algorithm>
+#include <cctype>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+Status TrtPrecisionModeToName(const TrtPrecisionMode mode, string* name) {
+  const char* kUnknown = "UNKNOWN";
+  *name = *kUnknown;
+  switch (mode) {
+    case TrtPrecisionMode::FP32:
+      *name = "FP32";
+      break;
+    case TrtPrecisionMode::FP16:
+      *name = "FP16";
+      break;
+    case TrtPrecisionMode::INT8:
+      *name = "INT8";
+      break;
+  }
+  if (name->compare(kUnknown) == 0)
+    return errors::OutOfRange("Unknown precision mode");
+  return OkStatus();
+}
+
+Status TrtPrecisionModeFromName(const string& name, TrtPrecisionMode* mode) {
+  if (name == "FP32") {
+    *mode = TrtPrecisionMode::FP32;
+  } else if (name == "FP16") {
+    *mode = TrtPrecisionMode::FP16;
+  } else if (name == "INT8") {
+    *mode = TrtPrecisionMode::INT8;
+  } else {
+    return errors::InvalidArgument("Invalid precision mode name: ", name);
+  }
+  return OkStatus();
+}
+
+string DebugString(const TrtPrecisionMode mode) {
+  string mode_str;
+  TF_CHECK_OK(TrtPrecisionModeToName(mode, &mode_str));
+  return absl::StrCat("TrtPrecisionMode::", mode_str);
+}
+
+string ProfileStrategyToName(const ProfileStrategy strategy) {
+  switch (strategy) {
+    case ProfileStrategy::kRange:
+      return "Range";
+    case ProfileStrategy::kOptimal:
+      return "Optimal";
+    case ProfileStrategy::kRangeOptimal:
+      return "Range+Optimal";
+    case ProfileStrategy::kImplicitBatchModeCompatible:
+      return "ImplicitBatchModeCompatible";
+  }
+  return "Unknown";
+}
+
+Status ProfileStrategyFromName(const string& name, ProfileStrategy* strategy) {
+  std::string name_lowercase = absl::AsciiStrToLower(name);
+  if (name_lowercase == "range") {
+    *strategy = ProfileStrategy::kRange;
+  } else if (name_lowercase == "optimal") {
+    *strategy = ProfileStrategy::kOptimal;
+  } else if (name_lowercase == "range+optimal") {
+    *strategy = ProfileStrategy::kRangeOptimal;
+  } else if (name_lowercase == "implicitbatchmodecompatible") {
+    *strategy = ProfileStrategy::kImplicitBatchModeCompatible;
+  } else {
+    return errors::InvalidArgument("Invalid profile strategy: ", name);
+  }
+  return OkStatus();
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/trt_parameters.h

@@ -0,0 +1,72 @@
+/* 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_CONVERT_TRT_PARAMETERS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_PARAMETERS_H_
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// The PrecisionMode controls the precision used in TRT converted parts of the
+// model. Setting PrecisionMode other than FP32 enables TensorRT to select
+// lower-precision implementations when searching for the fastest kernels.
+//
+// For regularized models whose input dynamic range is approximately one, this
+// typically produces significant speedups with negligible change in accuracy.
+// There is additional complexity when working with INT8, see Calibration.
+//
+// - FP32
+// - FP16 Enable FP16 layer selection, with FP32 fallback.
+// - INT8 Enable Int8 layer selection, with FP32 and FP16 fallback.
+//
+// Note that TensorRT will still choose a higher-precision kernel if it results
+// in overall lower runtime, or if no low-precision implementation exists.
+enum class TrtPrecisionMode { FP32, FP16, INT8 };
+
+Status TrtPrecisionModeToName(const TrtPrecisionMode mode, string* name);
+
+Status TrtPrecisionModeFromName(const string& name, TrtPrecisionMode* mode);
+
+string DebugString(const TrtPrecisionMode mode);
+
+// Optimization profile generation strategies.
+// - `kRange`: create one profile that works for inputs with dimension values
+//   in the range of [min_dims, max_dims] where min_dims and max_dims are
+//   derived from the provided inputs.
+// - `kOptimal`: create one profile for each input. The profile only works for
+//   inputs with the same dimensions as the input it is created for. The GPU
+//   engine will be run with optimal performance with such inputs.
+// - `kRangeOptimal`: create the profiles for both `Range` and `Optimal`.
+// - `kImplicitBatchModeCompatible`: create the profiles that will produce the
+//   same GPU engines as the implicit_batch_mode would produce.
+enum class ProfileStrategy {
+  kRange,
+  kOptimal,
+  kRangeOptimal,
+  kImplicitBatchModeCompatible,
+};
+
+string ProfileStrategyToName(const ProfileStrategy strategy);
+Status ProfileStrategyFromName(const string& name, ProfileStrategy* strategy);
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_PARAMETERS_H_

tensorflow/compiler/tf2tensorrt/convert/utils.cc

@@ -0,0 +1,280 @@
+/* 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/convert/utils.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "absl/strings/ascii.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+string DebugString(const nvinfer1::Dims& dims) {
+  string out = StrCat("nvinfer1::Dims(nbDims=", dims.nbDims, ", d=");
+  for (int i = 0; i < std::max(dims.nbDims, 0); ++i) {
+    StrAppend(&out, dims.d[i]);
+    StrAppend(&out, ",");
+  }
+  StrAppend(&out, ")");
+  return out;
+}
+
+string DebugString(const DataType tf_type) {
+  switch (tf_type) {
+    case DT_FLOAT:
+      return "DT_FLOAT";
+    case DT_HALF:
+      return "DT_HALF";
+    case DT_INT32:
+      return "DT_INT32";
+    case DT_INT8:
+      return "DT_INT8";
+    case DT_BOOL:
+      return "DT_BOOL";
+    case DT_UINT8:
+      return "DT_UINT8";
+    default:
+      return "Unknow TF DataType";
+  }
+}
+
+string DebugString(const nvinfer1::DataType trt_dtype) {
+  switch (trt_dtype) {
+    case nvinfer1::DataType::kFLOAT:
+      return "kFLOAT";
+    case nvinfer1::DataType::kHALF:
+      return "kHALF";
+    case nvinfer1::DataType::kINT8:
+      return "kINT8";
+    case nvinfer1::DataType::kINT32:
+      return "kINT32";
+    case nvinfer1::DataType::kBOOL:
+      return "kBOOL";
+#if IS_TRT_VERSION_GE(8, 5, 0, 0)
+    case nvinfer1::DataType::kUINT8:
+      return "kUINT8";
+#endif
+#if IS_TRT_VERSION_GE(8, 6, 0, 0)
+    case nvinfer1::DataType::kFP8:
+      return "kFP8";
+#endif
+    default:
+      return "Invalid TRT data type";
+  }
+}
+
+string DebugString(const nvinfer1::Permutation& permutation, int len) {
+  string out = "nvinfer1::Permutation(";
+  for (int i = 0; i < len; ++i) {
+    StrAppend(&out, permutation.order[i], ",");
+  }
+  StrAppend(&out, ")");
+  return out;
+}
+
+string DebugString(const ITensorProxyPtr& tensor) {
+  return StrCat(
+      tensor->is_trt_tensor() ? "nvinfer1::ITensor(@" : "SimpleItensor(@",
+      reinterpret_cast<uintptr_t>(&tensor), ", name=", tensor->getName(),
+      ", dtype=", DebugString(tensor->getType()),
+      ", dims=", DebugString(tensor->getDimensions()), ")");
+}
+
+string DebugString(const nvinfer1::ITensor& tensor) {
+  return StrCat("nvinfer1::ITensor(@", reinterpret_cast<uintptr_t>(&tensor),
+                ", name=", tensor.getName(),
+                ", dtype=", DebugString(tensor.getType()),
+                ", dims=", DebugString(tensor.getDimensions()), ")");
+}
+
+string DebugString(const std::vector<nvinfer1::Dims>& dimvec) {
+  return absl::StrCat("[",
+                      absl::StrJoin(dimvec, ",",
+                                    [](std::string* out, nvinfer1::Dims in) {
+                                      out->append(DebugString(in));
+                                    }),
+                      "]");
+}
+
+string DebugString(const std::vector<TensorShape>& shapes) {
+  return TensorShapeUtils::ShapeListString(shapes);
+}
+
+string DebugString(const std::vector<PartialTensorShape>& shapes) {
+  return PartialTensorShapeUtils::PartialShapeListString(shapes);
+}
+
+// Checks whether actual_shapes are compatible with cached_shapes. This should
+// only be used in implicit batch mode (in explicit batch mode one needs to
+// check the profile ranges). Therefore implicit batch mode is assumed.
+// It is also assumed that both actual_shapes and cached_shapes have been
+// verified by TRTEngineOp::VerifyInputShapes, which ensures that the batch size
+// for all tensors are the same.
+bool AreShapesCompatible(const std::vector<TensorShape>& actual_shapes,
+                         const std::vector<TensorShape>& cached_shapes) {
+  auto match_shape = [](const TensorShape& actual_shape,
+                        const TensorShape& cached_shape) {
+    // Match the rank.
+    if (actual_shape.dims() != cached_shape.dims()) return false;
+    // Match the batch size. In implicit batch mode cached_shape.dim_size(0) is
+    // the max batch size, which can be larger than the actual batch size.
+    if (actual_shape.dim_size(0) > cached_shape.dim_size(0)) return false;
+    // Match remaining dimensions.
+    for (int i = 1; i < actual_shape.dims(); ++i) {
+      if (actual_shape.dim_size(i) != cached_shape.dim_size(i)) return false;
+    }
+    return true;
+  };
+  for (int i = 0; i < actual_shapes.size(); ++i) {
+    if (!match_shape(actual_shapes[i], cached_shapes[i])) {
+      return false;
+    }
+  }
+  return true;
+}
+Status GetNetworkInputShapes(const nvinfer1::INetworkDefinition* network,
+                             std::vector<PartialTensorShape>* input_shapes) {
+  const int n_inputs = network->getNbInputs();
+  input_shapes->resize(n_inputs);
+  for (int i = 0; i < n_inputs; i++) {
+    const ITensorProxyPtr input = network->getInput(i);
+    TF_RETURN_IF_ERROR(DimsAdapter(input->getDimensions())
+                           .PartialTensorShape(&input_shapes->at(i)));
+  }
+  return OkStatus();
+}
+
+Status TfTypeToTrtType(DataType tf_type, nvinfer1::DataType* trt_type) {
+  switch (tf_type) {
+    case DT_FLOAT:
+      *trt_type = nvinfer1::DataType::kFLOAT;
+      break;
+    case DT_HALF:
+      *trt_type = nvinfer1::DataType::kHALF;
+      break;
+    case DT_INT32:
+      *trt_type = nvinfer1::DataType::kINT32;
+      break;
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    case DT_BOOL:
+      *trt_type = nvinfer1::DataType::kBOOL;
+      break;
+#endif
+#if IS_TRT_VERSION_GE(8, 5, 0, 0)
+    case DT_UINT8:
+      *trt_type = nvinfer1::DataType::kUINT8;
+      break;
+#endif
+    default:
+      return errors::InvalidArgument("Unsupported tensorflow data type ",
+                                     DataTypeString(tf_type));
+  }
+  return OkStatus();
+}
+
+Status TrtTypeToTfType(nvinfer1::DataType trt_type, DataType* tf_type) {
+  switch (trt_type) {
+    case nvinfer1::DataType::kFLOAT:
+      *tf_type = DT_FLOAT;
+      break;
+    case nvinfer1::DataType::kHALF:
+      *tf_type = DT_HALF;
+      break;
+    case nvinfer1::DataType::kINT32:
+      *tf_type = DT_INT32;
+      break;
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    case nvinfer1::DataType::kBOOL:
+      *tf_type = DT_BOOL;
+      break;
+#endif
+#if IS_TRT_VERSION_GE(8, 5, 0, 0)
+    case nvinfer1::DataType::kUINT8:
+      *tf_type = DT_UINT8;
+      break;
+#endif
+#if IS_TRT_VERSION_GE(8, 6, 0, 0)
+    case nvinfer1::DataType::kFP8:
+      *tf_type = DT_FLOAT8_E4M3FN;
+      break;
+#endif
+    default:
+      return errors::InvalidArgument("Invalid TRT data type");
+  }
+  return OkStatus();
+}
+
+int GetNumberOfEngineInputs(const nvinfer1::ICudaEngine* engine) {
+  int n_bindings = engine->getNbBindings();
+  int n_input = 0;
+  for (int i = 0; i < n_bindings; i++) {
+    if (engine->bindingIsInput(i)) n_input++;
+  }
+  // According to TensorRT 7 doc: "If the engine has been built for K profiles,
+  // the first getNbBindings() / K bindings are used by profile number 0, the
+  // following getNbBindings() / K bindings are used by profile number 1 etc."
+  // Therefore, to get the number of input tensors, we need to divide by the
+  // the number of profiles.
+  int n_profiles = engine->getNbOptimizationProfiles();
+  return n_input / n_profiles;
+}
+
+absl::string_view GetDeviceName(const Node* node) {
+  if (node->has_assigned_device_name()) {
+    return node->assigned_device_name();
+  }
+  return node->requested_device();
+}
+
+std::optional<DeviceNameUtils::ParsedName> GetDeviceParsedName(
+    const Node* node) {
+  absl::string_view device_name = GetDeviceName(node);
+  DeviceNameUtils::ParsedName parsed_name;
+  if (!DeviceNameUtils::ParseFullName(device_name, &parsed_name)) {
+    return std::nullopt;
+  }
+  return parsed_name;
+}
+
+std::optional<DeviceNameUtils::ParsedName> MergeIfCompatible(
+    const DeviceNameUtils::ParsedName& a,
+    const DeviceNameUtils::ParsedName& b) {
+  DeviceNameUtils::ParsedName merged_name = a;
+  if (!DeviceNameUtils::MergeDevNames(&merged_name, b,
+                                      /*allow_soft_placement=*/false)
+           .ok()) {
+    return std::nullopt;
+  }
+  return merged_name;
+}
+
+std::optional<DeviceNameUtils::ParsedName> MergeIfCompatible(
+    const DeviceNameUtils::ParsedName& a, absl::string_view b) {
+  DeviceNameUtils::ParsedName b_parsed_name;
+  if (!DeviceNameUtils::ParseFullName(b, &b_parsed_name)) {
+    return std::nullopt;
+  }
+
+  return MergeIfCompatible(a, b_parsed_name);
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/utils.h

@@ -0,0 +1,399 @@
+/* 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_CONVERT_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_UTILS_H_
+
+#include <algorithm>
+#include <iterator>
+#include <memory>
+#include <type_traits>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_tensor_proxy.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/util/env_var.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "third_party/tensorrt/NvInfer.h"
+
+#define TFTRT_ERROR(func, ...)                                              \
+  do {                                                                      \
+    return func("TFTRT::", __FUNCTION__, ":", __LINE__, ": ", __VA_ARGS__); \
+  } while (0)
+
+#define TFTRT_CHECK_SHAPE_TENSOR(tensor)                                 \
+  if (!IsTrtShapeTensorCompatible(tensor)) {                             \
+    TFTRT_ERROR(errors::InvalidArgument, "Tensor of type ",              \
+                DebugString(tensor.dtype()), " having shape ",           \
+                tensor.shape().DebugString(), " is not TRT compatible"); \
+  }
+
+namespace tensorflow {
+namespace tensorrt {
+
+static constexpr char kCastOutputTypeAttrName[] = "DstT";
+
+#if !IS_TRT_VERSION_GE(8, 2, 0, 0)
+template <typename T>
+struct TrtDestroyer {
+  void operator()(T* t) {
+    if (t) t->destroy();
+  }
+};
+template <typename T>
+using TrtUniquePtrType = std::unique_ptr<T, TrtDestroyer<T>>;
+#else
+template <typename T>
+using TrtUniquePtrType = std::unique_ptr<T>;
+#endif
+
+// Define a hash function for vector<TensorShape> because it is used as the key
+// for the engine cache.
+struct VectorTensorShapeHasher {
+  std::size_t operator()(const std::vector<TensorShape>& key) const {
+    return std::hash<std::string>()(TensorShapeUtils::ShapeListString(key));
+  }
+};
+
+using absl::StrAppend;
+using absl::StrCat;
+
+// This utility template converts an arithmetic type to a string. This function
+// is necessary to allow the following function to behave recursively:
+// `string DebugString(const std::vector<CType>&)`.
+template <typename CType, typename = typename std::enable_if<
+                              std::is_arithmetic<CType>::value, CType>::type>
+string DebugString(const CType& el) {
+  string el_str = std::to_string(el);
+  // Prettify std::to_string which can sometimes returns 1.50000 instead of 1.5.
+  // In short it removes trailing 0s in a string-formatted number.
+  el_str.erase(el_str.find_last_not_of('0') + 1, std::string::npos);
+  return el_str;
+}
+// This utility template converts nested vectors to a string for debug purposes.
+template <typename CType>
+string DebugString(const std::vector<CType>& vector) {
+  string tmp_s = "";
+  for (const auto el : vector) {
+    StrAppend(&tmp_s, StrCat(DebugString(el), ", "));
+  }
+  return StrCat("{", tmp_s.substr(0, tmp_s.length() - 2), "}");
+}
+string DebugString(const nvinfer1::Dims& dims);
+string DebugString(const nvinfer1::DataType trt_dtype);
+string DebugString(const DataType tf_type);
+string DebugString(const nvinfer1::Permutation& permutation, int len);
+string DebugString(const ITensorProxyPtr& tensor);
+string DebugString(const nvinfer1::ITensor& tensor);
+string DebugString(const std::vector<nvinfer1::Dims>& dimvec);
+string DebugString(const std::vector<TensorShape>& shapes);
+string DebugString(const std::vector<PartialTensorShape>& shapes);
+
+template <size_t N>
+string DebugString(const absl::InlinedVector<int64, N>& data) {
+  return absl::StrCat("[", absl::StrJoin(data, ","), "]");
+}
+
+inline bool HasStaticShape(const nvinfer1::Dims& dims) {
+  if (dims.nbDims < 0) return false;
+  for (int d = 0; d < dims.nbDims; ++d) {
+    if (dims.d[d] < 0) return false;
+  }
+  return true;
+}
+
+template <typename T>
+bool HasStaticShape(const T& dims) {
+  return !absl::c_any_of(dims, [](int i) { return i < 0; });
+}
+
+// Returns whether a shape is compatible with a TRT shape tensor.
+template <typename TensorShapeType>
+inline bool IsTrtShapeTensorCompatible(const TensorShapeType& shape) {
+  return (
+      shape.dims() == 0 ||
+      (shape.dims() == 1 && shape.num_elements() <= nvinfer1::Dims::MAX_DIMS));
+}
+
+// Returns whether a TF tensor could be interpreted as a TRT shape tensor.
+inline bool IsTrtShapeTensorCompatible(const Tensor& tensor) {
+  return tensor.dtype() == DT_INT32 &&
+         IsTrtShapeTensorCompatible(tensor.shape());
+}
+
+// Adapts various representations of shape (TF Shape, TRT Dims, plain
+// containers) and provides methods for properties (length, volume) and
+// conversion between types. Note that unlike TF's TensorShape, the underlying
+// storage will only contain active dimensions. In the case of scalar shapes,
+// `NumDims` is allowed to return 0 or 1, but the `storage_` vector will contain
+// 1 element in both cases. In the non-scalar case, `NumDims() ==
+// storage_.size()`.
+class DimsAdapter {
+ public:
+  using StorageType = absl::InlinedVector<int64_t, 4>;
+
+ private:
+  template <typename T>
+  using EnableIfNotTensorShapeType =
+      std::enable_if_t<!std::is_base_of<TensorShapeBase<T>, T>::value>;
+
+  template <typename T>
+  using EnableIfInt = std::enable_if_t<std::is_arithmetic<T>::value &&
+                                       std::is_integral<T>::value>;
+
+ public:
+  //----- Constructors ------
+
+  // Constructs from an absl::Span.
+  template <typename T>
+  explicit DimsAdapter(absl::Span<T> shape)
+      : num_dims_(static_cast<int32_t>(shape.size())) {
+    absl::c_copy(shape, std::back_inserter(storage_));
+  }
+
+  // Constructs from an absl::Span.
+  template <typename T>
+  explicit DimsAdapter(const std::vector<T>& shape)
+      : num_dims_(static_cast<int32_t>(shape.size())) {
+    absl::c_copy(shape, std::back_inserter(storage_));
+  }
+
+  // Constructs from a TRT dims object.
+  DimsAdapter(const nvinfer1::Dims& dims) : num_dims_(dims.nbDims) {
+    absl::c_copy(absl::MakeSpan(dims.d, dims.d + std::max(dims.nbDims, 0)),
+                 std::back_inserter(storage_));
+  }
+
+  // Constructs explicitly specifying num_dims and storage data.
+  DimsAdapter(int32_t num_dims, StorageType data)
+      : num_dims_(num_dims), storage_(std::forward<StorageType>(data)) {}
+
+  // Constructs from a TensorShape or PartialTensorShape.
+  template <typename T>
+  static StatusOr<DimsAdapter> Create(const TensorShapeBase<T>& shape,
+                                      bool ignore_first_dim = false) {
+    if (shape.dims() > nvinfer1::Dims::MAX_DIMS)
+      return errors::InvalidArgument("dims of TensorShape exceed MAX_DIMS");
+    if (ignore_first_dim && shape.dims() <= 0)
+      return errors::InvalidArgument(
+          "removing first dim requires explicit batch dimension");
+    if (shape.dims() == -1) {
+      return DimsAdapter(-1, StorageType{});
+    }
+    if (shape.dims() == 0) {
+      return DimsAdapter(0, StorageType{1});
+    }
+    auto offt = (ignore_first_dim ? 1 : 0);
+    return DimsAdapter(
+        absl::MakeSpan(shape.dim_sizes().begin() + offt, shape.dims() - offt));
+  }
+
+  // Constructs from a container.
+  template <typename InputSequence,
+            typename = EnableIfNotTensorShapeType<InputSequence>>
+  static StatusOr<DimsAdapter> Create(const InputSequence& shape,
+                                      bool ignore_first_dim = false) {
+    if (ignore_first_dim && shape.size() <= 0) {
+      return errors::InvalidArgument(
+          "removing first dim requires explicit batch dimension");
+    }
+    return DimsAdapter(
+        absl::MakeSpan(shape).subspan(ignore_first_dim ? 1 : 0, shape.size()));
+  }
+
+  //----- Conversion Utilities ------
+
+  //  Converts to an nvinfers::Dims and assign the result to the object passed
+  //  in via the result pointer.
+  void TrtDims(nvinfer1::Dims* result) const {
+    result->nbDims = num_dims_;
+    absl::c_copy(storage_, static_cast<int32_t*>(result->d));
+  }
+
+  // Converts to an nvinfer1::Dims and return by value.
+  nvinfer1::Dims AsTrtDims() const {
+    nvinfer1::Dims result;
+    TrtDims(&result);
+    return result;
+  }
+
+  // Converts to a TensorShape and assigns the result to the object passed in
+  // via the shape pointer.
+  Status TensorShape(TensorShape* shape,
+                     std::optional<int> batch_size = std::nullopt) const {
+    TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(
+        static_cast<const int64_t*>(storage_.data()), storage_.size(), shape));
+    if (batch_size) shape->InsertDim(0, *batch_size);
+    return OkStatus();
+  }
+
+  // Converts to a PartialTensorShape and assigns the result to the object
+  // passed in via the shape pointer.
+  Status PartialTensorShape(
+      PartialTensorShape* shape,
+      std::optional<int> batch_size = std::nullopt) const {
+    TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(
+        static_cast<const int64_t*>(storage_.data()), storage_.size(), shape));
+    if (batch_size) shape->InsertDim(0, *batch_size);
+    return OkStatus();
+  }
+
+  // Copies the dimension values to the vector passed in via the shape pointer.
+  template <typename T, typename = EnableIfInt<T>>
+  Status Vector(std::vector<T>* shape) const {
+    shape->clear();
+    absl::c_copy(storage_, std::back_inserter(*shape));
+    return OkStatus();
+  }
+
+  //----- Property Accessors ------
+
+  // Returns true if the shape has no dynamic dimensions.
+  bool IsStatic() const {
+    return !absl::c_any_of(storage_, [](auto i) { return i < 0; });
+  }
+
+  // Returns product of all dimensions.
+  int64_t Volume() const {
+    return absl::c_accumulate(storage_, static_cast<int64_t>(1),
+                              std::multiplies<>());
+  }
+
+  int32_t NumDims() const { return num_dims_; }
+
+  // Returns true if the shape should be interpreted as a scalar. This follows
+  // TensorRT conversions: a scalar shape can have NumDims()==1 or NumDims()==0,
+  // but the underlying storage_ container has a single dimension of size 1.
+  bool IsScalar() const {
+    return (num_dims_ == 0 || num_dims_ == 1) && storage_.size() == 1 &&
+           storage_[0] == 1;
+  }
+
+  // Returns true if the dimension storage is empty. This indicates an empty
+  // shape in both the scalar and non-scalar case.
+  bool IsEmpty() const { return storage_.empty(); }
+
+  string DebugString() const {
+    auto vol = absl::c_accumulate(storage_, static_cast<int64_t>(1),
+                                  std::multiplies<>());
+    return absl::StrCat("DimsAdapter(num_dims=", num_dims_, ",shape=[",
+                        absl::StrJoin(storage_, ","), "],", "vol=", vol, ")");
+  }
+
+  // Returns beginning iterator for the underlying storage.
+  StorageType::const_iterator begin() const { return storage_.begin(); }
+
+  // Returns ending iterator for the underlying storage.
+  StorageType::const_iterator end() const { return storage_.end(); }
+
+  // Returns the size of the dimension at `idx`.
+  StorageType::value_type dim(size_t idx) const { return storage_[idx]; }
+
+  // Returns a references to the dimension at `idx`.
+  StorageType::value_type& dim(size_t idx) { return storage_[idx]; }
+
+  //----- Non-Const Operators ------
+
+  DimsAdapter& Append(int32_t dim) {
+    StatusOr<bool> is_scalar = IsScalar();
+    if (!is_scalar.ok()) return *this;
+    num_dims_ = *is_scalar ? 2 : num_dims_ + 1;
+    storage_.push_back(dim);
+    return *this;
+  }
+
+  DimsAdapter& Prepend(std::optional<int32_t> dim) {
+    if (dim) {
+      num_dims_ = IsScalar() ? 2 : num_dims_ + 1;
+      storage_.insert(storage_.begin(), *dim);
+    }
+    return *this;
+  }
+
+  Status RemoveBatchDimension() {
+    if (storage_.empty())
+      return errors::InvalidArgument(
+          "attempted to remove batch dim from scalar");
+    num_dims_ -= 1;
+    storage_.erase(storage_.begin());
+    return OkStatus();
+  }
+
+  //----- Comparison Operators ------
+
+  bool operator==(const DimsAdapter& rhs) const {
+    if (rhs.num_dims_ != num_dims_) return false;
+    for (int i = 0; i < num_dims_; i++) {
+      if (rhs.storage_[i] != storage_[i]) return false;
+    }
+    return true;
+  }
+
+  bool operator!=(const DimsAdapter& rhs) const { return !(*this == rhs); }
+
+ private:
+  int32_t num_dims_{0};
+  StorageType storage_{};
+};
+
+Status GetNetworkInputShapes(const nvinfer1::INetworkDefinition* network,
+                             std::vector<PartialTensorShape>* input_shapes);
+
+Status TfTypeToTrtType(DataType tf_type, nvinfer1::DataType* trt_type);
+Status TrtTypeToTfType(nvinfer1::DataType trt_type, DataType* tf_type);
+
+// Returns true if an engine built for cached_shapes can also run actual_shapes.
+bool AreShapesCompatible(const std::vector<TensorShape>& actual_shapes,
+                         const std::vector<TensorShape>& cached_shapes);
+
+// Returns the number of inputs for the engine, which also correspends to the
+// number of input tensors for the network. This can differ from the number of
+// input bindings, because the number of total input bindings equals the number
+// of profiles times the number of engine inputs.
+int GetNumberOfEngineInputs(const nvinfer1::ICudaEngine* engine);
+
+// Returns the string representation for the assigned device or the requested
+// device of the given node.
+absl::string_view GetDeviceName(const Node* node);
+
+// Returns the ParsedName representation for the assigned device or the
+// requested device string of the given node. If the device string is invalid,
+// returns std::nullopt.
+std::optional<DeviceNameUtils::ParsedName> GetDeviceParsedName(
+    const Node* node);
+
+// If the given two device assignments as compatible, returns the merge of the
+// two assignments. Otherwise, returns std::nullopt.
+std::optional<DeviceNameUtils::ParsedName> MergeIfCompatible(
+    const DeviceNameUtils::ParsedName& a, const DeviceNameUtils::ParsedName& b);
+// Similar to the above, except that the second device assignment is represented
+// by a string_view.
+std::optional<DeviceNameUtils::ParsedName> MergeIfCompatible(
+    const DeviceNameUtils::ParsedName& a, absl::string_view b);
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_UTILS_H_

tensorflow/compiler/tf2tensorrt/convert/weights.cc

@@ -0,0 +1,214 @@
+/* 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/convert/weights.h"
+
+#include <functional>
+#include <numeric>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+
+namespace convert {
+
+TRT_ShapedWeights::TRT_ShapedWeights(nvinfer1::DataType type)
+    : shape_(0, DimsAdapter::StorageType{}), type_(type), volume_(0) {}
+
+StatusOr<TRT_ShapedWeights> TRT_ShapedWeights::CreateWithTensor(
+    nvinfer1::DataType type, DimsAdapter dims, Tensor tensor) {
+  TRT_ShapedWeights weights(type);
+  weights.shape_ = dims;
+  weights.tensor_ = std::forward<Tensor>(tensor);
+  weights.volume_ = weights.shape_.Volume();
+  if (weights.shape_.NumDims() == 0) {
+    DCHECK(weights.shape_.IsEmpty() || weights.shape_.IsScalar());
+  }
+  return weights;
+}
+
+nvinfer1::Weights TRT_ShapedWeights::GetTrtWeights() const {
+  return nvinfer1::Weights{type_, GetPointer<int8>(), volume_};
+}
+
+Status TRT_ShapedWeights::SetShape(DimsAdapter dims) {
+  if (volume_ != dims.Volume()) {
+    VLOG(2) << "Changing shape from " << shape_.DebugString() << ", to "
+            << dims.DebugString();
+    return errors::Internal("SetShape would change number of elements");
+  }
+  shape_ = std::move(dims);
+  return OkStatus();
+}
+
+size_t TRT_ShapedWeights::size_bytes() const {
+  size_t data_type_size = -1;
+  switch (type_) {
+    case nvinfer1::DataType::kFLOAT:
+    case nvinfer1::DataType::kINT32:
+      data_type_size = 4;
+      break;
+    case nvinfer1::DataType::kHALF:
+      data_type_size = 2;
+      break;
+#if IS_TRT_VERSION_GE(8, 5, 0, 0)
+    case nvinfer1::DataType::kUINT8:
+#endif
+#if IS_TRT_VERSION_GE(8, 6, 0, 0)
+    case nvinfer1::DataType::kFP8:
+#endif
+    case nvinfer1::DataType::kINT8:
+    case nvinfer1::DataType::kBOOL:
+      data_type_size = 1;
+      break;
+  }
+  return volume_ * data_type_size;
+}
+
+string TRT_ShapedWeights::DebugString() const {
+  return absl::StrCat(
+      "TRT_ShapedWeights(shape=", shape_.DebugString(),
+      ", type=", tensorflow::tensorrt::DebugString(type_),
+      ", values=", reinterpret_cast<uintptr_t>(GetPointer<int8>()), ")");
+}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(ITensorProxyPtr tensor)
+    : tensor_proxy_ptr_(tensor),
+      initialized_(true),
+      arg_type_(TRT_ArgumentType::TENSOR) {}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(ITensorProxyPtr tensor, int batch_size)
+    : tensor_proxy_ptr_(tensor),
+      batch_size_(batch_size),
+      initialized_(true),
+      arg_type_(TRT_ArgumentType::TENSOR) {}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(nvinfer1::ITensor* tensor,
+                                         int batch_size)
+    : tensor_proxy_ptr_(tensor),
+      batch_size_(batch_size),
+      initialized_(true),
+      arg_type_(TRT_ArgumentType::TENSOR) {}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(nvinfer1::DataType trt_dtype,
+                                         const nvinfer1::Dims& trt_dims,
+                                         int batch_size)
+    : tensor_proxy_ptr_(new SimpleITensor(trt_dtype, trt_dims)),
+      batch_size_(batch_size),
+      initialized_(true),
+      arg_type_(TRT_ArgumentType::TENSOR) {}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(const TRT_ShapedWeights& weights)
+    : weights_(weights),
+      initialized_(true),
+      arg_type_(TRT_ArgumentType::WEIGHTS) {}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(const ResourceHandle& resource)
+    : resource_(resource),
+      initialized_(true),
+      arg_type_(TRT_ArgumentType::RESOURCE) {}
+
+TRT_TensorOrWeights::TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs)
+    : tensor_proxy_ptr_(rhs.tensor_proxy_ptr_),
+      batch_size_(rhs.batch_size_),
+      resource_(rhs.resource_),
+      weights_(rhs.weights_),
+      initialized_(rhs.initialized_),
+      arg_type_(rhs.arg_type_) {}
+
+void TRT_TensorOrWeights::operator=(const TRT_TensorOrWeights& rhs) {
+  tensor_proxy_ptr_ = rhs.tensor_proxy_ptr_;
+  batch_size_ = rhs.batch_size_;
+  weights_ = rhs.weights_;
+  resource_ = rhs.resource_;
+  initialized_ = rhs.initialized_;
+  arg_type_ = rhs.arg_type_;
+}
+
+ITensorProxyPtr TRT_TensorOrWeights::tensor() const {
+  DCHECK(is_tensor());
+  return tensor_proxy_ptr_;
+}
+
+ResourceHandle TRT_TensorOrWeights::resource() const {
+  DCHECK(is_resource());
+  return resource_;
+}
+
+nvinfer1::Dims TRT_TensorOrWeights::GetTrtDims() const {
+  switch (arg_type_) {
+    case TRT_ArgumentType::TENSOR:
+      return tensor()->getDimensions();
+    case TRT_ArgumentType::WEIGHTS:
+      return weights().Shape().AsTrtDims();
+    case TRT_ArgumentType::RESOURCE:
+      return {0, {}};  // Scalar.
+  }
+}
+
+Status TRT_TensorOrWeights::GetTfType(DataType* tf_type) const {
+  if (!initialized_) {
+    return errors::Internal("The object is not initialized");
+  }
+  switch (arg_type_) {
+    case TRT_ArgumentType::TENSOR: {
+      nvinfer1::DataType trt_type = tensor()->getType();
+      return TrtTypeToTfType(trt_type, tf_type);
+    }
+    case TRT_ArgumentType::WEIGHTS:
+      *tf_type = weights().GetTensor().dtype();
+      return OkStatus();
+    case TRT_ArgumentType::RESOURCE:
+      *tf_type = DataType::DT_RESOURCE;
+      return OkStatus();
+  }
+}
+
+string TRT_TensorOrWeights::DebugString() const {
+  string output = "TRT_TensorOrWeights(type=";
+  if (is_tensor()) {
+    absl::StrAppend(&output,
+                    "tensor=", tensorflow::tensorrt::DebugString(tensor()),
+                    ", batch_size=", batch_size_);
+  } else {
+    absl::StrAppend(&output, "weights=", weights_.DebugString());
+  }
+  absl::StrAppend(&output, ")");
+  return output;
+}
+
+StatusOr<TRT_ShapedWeights> TrtWeightStore::GetTempWeights(
+    nvinfer1::DataType trt_dtype, const DimsAdapter& dims) {
+  DataType tf_dtype;
+  TF_RETURN_IF_ERROR(TrtTypeToTfType(trt_dtype, &tf_dtype));
+  TensorShape shape;
+  TF_RETURN_IF_ERROR(dims.TensorShape(&shape));
+  // TODO(jie): check weights size_bytes. 0 means type error
+  Tensor tensor(tf_dtype, shape);
+  StatusOr<TRT_ShapedWeights> weights =
+      TRT_ShapedWeights::CreateWithTensor(trt_dtype, dims, tensor);
+  TRT_ENSURE_OK(weights);
+  store_.emplace_back(std::move(tensor));
+  return weights;
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/weights.h

@@ -0,0 +1,295 @@
+/* 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_CONVERT_WEIGHTS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_WEIGHTS_H_
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <vector>
+
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_tensor_proxy.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+// Class to convert TF compile-time constants (e.g. Const nodes) to TRT weight.
+class TRT_ShapedWeights {
+ public:
+  explicit TRT_ShapedWeights(
+      nvinfer1::DataType type = nvinfer1::DataType::kFLOAT);
+
+  // Constructs a weights from another weights.
+  //
+  // NOTE: this does not copy the underlying buffer but only increase its
+  // reference count.
+  TRT_ShapedWeights(const TRT_ShapedWeights& rhs) = default;
+
+  nvinfer1::Weights GetTrtWeights() const;
+
+  const Tensor& GetTensor() const { return tensor_; }
+
+  // Returns a pointer of type const T to the underlying buffer of the tensor.
+  template <typename T>
+  const T* GetPointer() const {
+    int64 num_elem =
+        (tensor_.NumElements() * DataTypeSize(tensor_.dtype())) / sizeof(T);
+    return tensor_.bit_casted_shaped<T, 1>({num_elem}).data();
+  }
+
+  // Returns a pointer of type T to the underlying buffer of the tensor.
+  template <typename T>
+  T* GetPointer() {
+    int64 num_elem =
+        (tensor_.NumElements() * DataTypeSize(tensor_.dtype())) / sizeof(T);
+    return tensor_.bit_casted_shaped<T, 1>({num_elem}).data();
+  }
+
+  // Fills all the weight values with value.
+  template <typename T>
+  Status SetValues(T value) {
+    switch (type_) {
+      case nvinfer1::DataType::kFLOAT: {
+        float* ptr = tensor_.flat<float>().data();
+        std::fill(ptr, ptr + volume_, value);
+        break;
+      }
+      case nvinfer1::DataType::kHALF: {
+        Eigen::half* ptr = tensor_.flat<Eigen::half>().data();
+        std::fill(ptr, ptr + volume_, Eigen::half(value));
+        break;
+      }
+      case nvinfer1::DataType::kINT32: {
+        int32* ptr = tensor_.flat<int32>().data();
+        std::fill(ptr, ptr + volume_, value);
+        break;
+      }
+      default:
+        return errors::InvalidArgument(
+            "Unsupported data type ", tensorflow::tensorrt::DebugString(type_));
+    }
+    return OkStatus();
+  }
+
+  Status SetShape(DimsAdapter dims);
+  void SetShapeUnsafe(DimsAdapter dims) { shape_ = std::move(dims); }
+
+  // Returns total number of elements. Returning 0 means either some dim is 0
+  // or the number of dims is 0. Note that a TF scalar constant is marked as
+  // Dims{0, {1}}, and has a count() == 1.
+  int64_t count() const { return volume_; }
+
+  size_t size_bytes() const;
+
+  string DebugString() const;
+
+  template <typename T>
+  absl::Span<const T> GetSpan() const {
+    return absl::Span<const T>(tensor_.flat<T>().data(), volume_);
+  }
+
+  template <typename T>
+  std::vector<T> ToVector() const {
+    auto span = GetSpan<T>();
+    return std::vector<T>(span.data(), span.data() + span.size());
+  }
+
+  nvinfer1::DataType TrtDType() const { return type_; }
+
+  const DimsAdapter& Shape() const { return shape_; }
+  DimsAdapter& Shape() { return shape_; }
+
+ private:
+  // The shape of the weights. Defaults to the empty shape.
+  DimsAdapter shape_;
+
+  // This creation method is only used by TrtWeightStore, which creates the
+  // underlying buffer.
+  static StatusOr<TRT_ShapedWeights> CreateWithTensor(nvinfer1::DataType type,
+                                                      DimsAdapter dims,
+                                                      Tensor tensor);
+
+  nvinfer1::DataType type_;
+
+  // All weights should be stored inside TrtWeightStore to make sure lifetime of
+  // all the underlying tensors are available until the engine is built. For
+  // this reason, tensor_ should never be reassigned to a different value that
+  // is not already present in the TrtWeightStore.
+  Tensor tensor_;
+  // Contains the volume of the weight's shape.
+  int64_t volume_;
+
+  friend class TrtWeightStore;
+};
+
+// Container for TRT_ShapedWeights. We need this container because TRT does not
+// manage the lifetime of the weights buffer, it only keeps a pointer to it and
+// requires that the data referenced by the pointer be available until the
+// building of engine is complete. For more information see
+// https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/classnvinfer1_1_1_weights.html
+//
+// TODO(laigd): consider adding garbage collection to the unused weights.
+class TrtWeightStore {
+ public:
+  // Gets a TRT_ShapedWeights with 'type' and 'dims'.
+  StatusOr<TRT_ShapedWeights> GetTempWeights(nvinfer1::DataType trt_type,
+                                             const DimsAdapter& dims);
+
+  // Gets a TRT_ShapedWeights with the same data type and dimensions as
+  // 'weights'.
+  StatusOr<TRT_ShapedWeights> GetTempWeights(const TRT_ShapedWeights& weights) {
+    return GetTempWeights(weights.TrtDType(), weights.Shape());
+  }
+
+ private:
+  // The backend storage of the TRT_ShapedWeights.
+  std::vector<Tensor> store_;
+};
+
+// Enumerates the possible types of arguments of a converter. This determines
+// what object is contained in TRT_TensorOrWeights, and converters can require
+// a specific type for each of their arguments.
+enum class TRT_ArgumentType {
+  TENSOR = 0,
+  WEIGHTS = 1,
+  RESOURCE = 2,
+};
+
+struct OpConverterParams;
+
+// Represents a TRT-style input to a TF node, it can be either a
+// ITensorProxyPtr (representing nvinfer1::ITensor* or SimpleITensor),
+// or TRT_ShapedWeights which is compile-time constant.
+//
+// TODO(laigd): maybe rename it to TrtArgument, or mimic XlaCompiler::Argument.
+class TRT_TensorOrWeights {
+ public:
+  TRT_TensorOrWeights() {}
+  TRT_TensorOrWeights(ITensorProxyPtr);
+  TRT_TensorOrWeights(ITensorProxyPtr tensor, int batch_size);
+
+  // Constructs a wrapper for the given ITensor.
+  // This is used by Converter when building the TRT network, where the ITensor
+  // is owned by the TRT network being built. See comment for 'trt_tensor_'
+  // in trt_proxy_tensor.h.
+  explicit TRT_TensorOrWeights(nvinfer1::ITensor* tensor, int batch_size = -1);
+
+  // Creates a SimpleITensor for trt_dtype and trt_dims and takes ownership of
+  // the object. Constructs a wrapper for the SimpleITensor. This is used by
+  // TrtNodeValidator to encapsulate the type and shape information for
+  // validation of graph nodes, and the created ITensor is fake and temporary,
+  // and should not be used to build any TRT network. See comment for
+  // 'simple_tensor_' in trt_proxy_tensor.h.
+  explicit TRT_TensorOrWeights(nvinfer1::DataType trt_dtype,
+                               const nvinfer1::Dims& trt_dims, int batch_size);
+
+  // Constructs a wrapper for the given weights.
+  explicit TRT_TensorOrWeights(const TRT_ShapedWeights& weights);
+
+  // Constructs a wrapper for the given resource handle.
+  explicit TRT_TensorOrWeights(const ResourceHandle& resource);
+
+  TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs);
+
+  void operator=(const TRT_TensorOrWeights& rhs);
+
+  bool is_tensor() const {
+    return initialized_ && arg_type_ == TRT_ArgumentType::TENSOR;
+  }
+  bool is_weights() const {
+    return initialized_ && arg_type_ == TRT_ArgumentType::WEIGHTS;
+  }
+  bool is_resource() const {
+    return initialized_ && arg_type_ == TRT_ArgumentType::RESOURCE;
+  }
+
+  ITensorProxyPtr tensor() const;
+
+  ResourceHandle resource() const;
+
+  ITensorProxyPtr as_tensor(const OpConverterParams* params);
+
+  TRT_ShapedWeights& weights() {
+    DCHECK(is_weights());
+    return weights_;
+  }
+
+  const TRT_ShapedWeights& weights() const {
+    DCHECK(is_weights());
+    return weights_;
+  }
+
+  nvinfer1::Dims GetTrtDims() const;
+
+  Status GetTfType(DataType* tf_type) const;
+
+  int batch_size() const { return batch_size_; }
+
+  string DebugString() const;
+
+  nvinfer1::DataType TrtDType() const {
+    if (arg_type_ == TRT_ArgumentType::RESOURCE) {
+      VLOG(0) << "Calling TrtDType() with a RESOURCE argument is undefined "
+                 "behavior.";
+    }
+    return arg_type_ == TRT_ArgumentType::TENSOR ? tensor_proxy_ptr_->getType()
+                                                 : weights_.TrtDType();
+  }
+
+ private:
+  void set_batch_size(int batch_size) { batch_size_ = batch_size; }
+
+  // First dimension of the TF tensor (NOT tensor_) that is represented by
+  // tensor_ is treated as the "batch dimension" by TRT, and tensor_'s
+  // dimensions (obtained via tensor_->getDimensions()) do not contain the batch
+  // dimension. For example, when a TF tensor with shape (A,B,C) is represented
+  // in TRT, tensor_->getDimensions() will be (B,C) and batch_size_ will be A.
+  //
+  // This requires that all tensors in the subgraph that is converted to a TRT
+  // engine have the same batch size are represented by the first dimension of
+  // their shape, and Converter will verify this during conversion. The drawback
+  // is that currently it cannot convert a graph that doesn't have the batch
+  // size represented in the shapes or the batch sizes are different. See
+  // b/118387490 for more details.
+  //
+  // If use_implicit_batch is false, batch_size_ is unused and
+  // tensor_->getDimensions() will contain the entire shape (A,B,C).
+  //
+  // tensor_proxy_ptr_ is used when arg_type_ == TENSOR.
+  ITensorProxyPtr tensor_proxy_ptr_ = nullptr;
+  int batch_size_ = -1;
+
+  // For DT_RESOURCE arguments (there is no corresponding type in TRT).
+  // resource_ is used when arg_type_ == RESOURCE.
+  ResourceHandle resource_;
+
+  // weights_ is used when arg_type_ == WEIGHTS.
+  TRT_ShapedWeights weights_;
+  bool initialized_ = false;
+  TRT_ArgumentType arg_type_ = TRT_ArgumentType::WEIGHTS;
+
+  friend class Converter;
+};
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_WEIGHTS_H_