chore: import upstream snapshot with attribution

tensorflow/compiler/tf2tensorrt/segment/segment.h

@@ -0,0 +1,94 @@
+/* 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_SEGMENT_SEGMENT_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_
+
+#include <set>
+#include <vector>
+
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/tf2tensorrt/segment/union_find.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"
+#include "tensorflow/core/platform/types.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace segment {
+
+constexpr char kTftrtOpMaxBatchSizeAttr[] = "_tftrt_op_max_batch_size";
+
+struct SegmentOptions {
+  // This struct holds per graph segmenting parameters.
+  // Segment must contain at least this many nodes.
+  int minimum_segment_size = 2;
+  bool use_implicit_batch = true;
+  // The maximum batch size used to build the engines in the graph, when
+  // use_implicit_batch is true.
+  std::optional<int> maximum_batch_size = std::nullopt;
+  // When use_implicit_batch is false or when we are building dynamic engines,
+  // we allow dynamic non-batch dimensions.
+  bool allow_dynamic_non_batch_dim = false;
+  // The name of the device to put the segment on.
+  std::set<string> exclude_node_list;
+};
+
+struct NodePtrCompare {
+  bool operator()(const Node* lhs, const Node* rhs) const {
+    return lhs->name() < rhs->name();
+  }
+};
+
+struct Segment {
+  Segment() {}
+  Segment(const ClusterProperty& property,
+          const std::set<const Node*, NodePtrCompare>& nodes)
+      : property(property), nodes(nodes) {}
+  ClusterProperty property;
+  std::set<const Node*, NodePtrCompare> nodes;
+};
+
+// Vector of segments, each entry contains a set of node pointers.
+using SegmentVector = std::vector<Segment>;
+
+// Get the subgraphs of a graph that can be handled by TensorRT.
+//
+// @param tf_graph Graph of the network.
+// @graph_properties is the static graph properties.
+// @param candidate_fn A function that returns OK for a Node* if
+// that node can be handled by TensorRT.
+// @param segments Returns the TensorRT segments/subgraphs. Each entry
+// in the vector describes a subgraph by giving a set of the names of
+// all the NodeDefs in that subgraph.
+// @return the status.
+Status SegmentGraph(const Graph* tf_graph,
+                    const grappler::GraphProperties* graph_properties,
+                    const std::function<Status(const Node*)>& candidate_fn,
+                    const std::function<bool(const Edge*)>& input_candidate_fn,
+                    const std::function<bool(const Edge*)>& output_candidate_fn,
+                    const SegmentOptions& options, SegmentVector* segments);
+
+}  // namespace segment
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_

tensorflow/compiler/tf2tensorrt/segment/segment_test.cc

@@ -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.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2tensorrt/segment/segment.h"
+
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/graph/testlib.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/session.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace segment {
+namespace test {
+
+class SegmentTest : public ::testing::Test {
+ protected:
+  std::function<Status(const Node*)> MakeCandidateFn(
+      const std::set<string>& node_names) {
+    return [node_names](const Node* node) -> Status {
+      if (node_names.find(node->name()) != node_names.end()) {
+        return OkStatus();
+      }
+      return errors::NotFound("Not a user specified candidate");
+    };
+  }
+
+  std::function<bool(const Edge*)> MakeInputEdgeCandidateFn(
+      const std::set<string>& node_names) {
+    return [node_names](const Edge* in_edge) -> bool {
+      return node_names.find(in_edge->dst()->name()) != node_names.end();
+    };
+  }
+
+  std::function<bool(const Edge*)> MakeOutputEdgeCandidateFn(
+      const std::set<string>& node_names) {
+    return [node_names](const Edge* out_edge) -> bool {
+      return node_names.find(out_edge->src()->name()) != node_names.end();
+    };
+  }
+
+  void RunTest(const Graph* graph,
+               const grappler::GraphProperties* graph_properties,
+               const std::set<string>& candidates,
+               const std::set<string>& input_candidates,
+               const std::set<string>& output_candidates,
+               const std::vector<std::set<string>>& expected_segments) {
+    SegmentVector segments;
+    TF_EXPECT_OK(SegmentGraph(graph, graph_properties,
+                              MakeCandidateFn(candidates),
+                              MakeInputEdgeCandidateFn(input_candidates),
+                              MakeOutputEdgeCandidateFn(output_candidates),
+                              segment_options_, &segments));
+    ValidateSegment(segments, expected_segments);
+  }
+
+  void RunTest(const Graph* graph, const std::set<string>& candidates,
+               const std::set<string>& input_candidates,
+               const std::set<string>& output_candidates,
+               const std::vector<std::set<string>>& expected_segments) {
+    RunTest(graph, nullptr, candidates, input_candidates, output_candidates,
+            expected_segments);
+  }
+
+  void ValidateSegment(const SegmentVector& segments,
+                       const std::vector<std::set<string>>& expected_segments) {
+    EXPECT_EQ(expected_segments.size(), segments.size());
+    for (int i = 0; i < segments.size(); ++i) {
+      std::set<string> segment_node_names;
+      for (const Node* node : segments[i].nodes) {
+        segment_node_names.insert(node->name());
+      }
+      const auto& expected = expected_segments[i];
+      for (const auto& name : expected) {
+        EXPECT_TRUE(segment_node_names.count(name))
+            << "Segment " << i << " is missing expected node: " << name;
+      }
+      if (segment_node_names.size() == expected.size()) continue;
+      for (const auto& name : segment_node_names) {
+        EXPECT_TRUE(expected.count(name))
+            << "Unexpected node found in segment " << i << ": " << name;
+      }
+    }
+  }
+
+  void DisableImplicitBatchMode() {
+    segment_options_.use_implicit_batch = false;
+    segment_options_.allow_dynamic_non_batch_dim = true;
+  }
+
+  void EnableImplicitBatchModeForStaticEngine(int maximum_batch_size = 1000) {
+    segment_options_.use_implicit_batch = true;
+    segment_options_.maximum_batch_size = maximum_batch_size;
+    segment_options_.allow_dynamic_non_batch_dim = false;
+  }
+
+  SegmentOptions segment_options_;
+};
+
+std::set<string> operator-(const std::set<string>& lhs, const string& rhs) {
+  std::set<string> result = lhs;
+  CHECK(result.erase(rhs));
+  return result;
+}
+
+TEST_F(SegmentTest, Empty) {
+  Scope s = Scope::NewRootScope();
+  Graph g(OpRegistry::Global());
+  TF_EXPECT_OK(s.ToGraph(&g));
+  // Expect no segments/subgraphs.
+  DisableImplicitBatchMode();
+  RunTest(&g, {}, {}, {}, {});
+}
+
+TEST_F(SegmentTest, Simple) {
+  //           feed
+  //          //  \\
+  //       add0    add1
+  //        | \    /
+  //        |  add2
+  //        | /   \\
+  //       add3    add4
+  //          \    /
+  //          <sink>
+  Scope s = Scope::NewRootScope();
+  auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT);
+  auto add0 = ops::Add(s.WithOpName("add0"), feed, feed);
+  auto add1 = ops::Add(s.WithOpName("add1"), feed, feed);
+  auto add2 = ops::Add(s.WithOpName("add2"), add0, add1);
+  auto add3 = ops::Add(s.WithOpName("add3"), add0, add2);
+  auto add4 = ops::Add(s.WithOpName("add4"), add2, add2);
+  Graph g(OpRegistry::Global());
+  TF_EXPECT_OK(s.ToGraph(&g));
+
+  // All Add operations are candidates, and we expect all of them to be
+  // collapsed into a single segment
+  const std::set<string> all_adds = {"add0", "add1", "add2", "add3", "add4"};
+  DisableImplicitBatchMode();
+  RunTest(&g, all_adds, all_adds, all_adds, {all_adds});
+
+  // Make add1 not a candidate, and we expect all other Add operations to be
+  // collapsed into a single segment
+  auto without_add1 = all_adds - "add1";
+  RunTest(&g, without_add1, without_add1, without_add1, {without_add1});
+
+  // Make add1 not a candidate and add2 not an input candidate, and we expect
+  // add0 and add2 are removed from the segment.
+  auto without_add2 = all_adds - "add2";
+  RunTest(&g, without_add1, without_add2, without_add1, {{"add3", "add4"}});
+
+  // Making add2 not an input candidate itself won't affect anything.
+  RunTest(&g, all_adds, without_add2, all_adds, {all_adds});
+
+  // Making add1 not an input candidate.
+  RunTest(&g, all_adds, without_add1, all_adds, {without_add1});
+
+  // Making add3 not an output candidate doesn't affect anything, since it's
+  // output is sink.
+  auto without_add3 = all_adds - "add3";
+  RunTest(&g, all_adds, all_adds, without_add3, {all_adds});
+}
+
+TEST_F(SegmentTest, WithDeviceAssignments) {
+  //           feed
+  //          //  \\
+  //       add0    add1
+  //        | \    /
+  //        |  add2
+  //        | /   \\
+  //       add3    add4
+  //          \    /
+  //          <sink>
+  Scope s = Scope::NewRootScope();
+  auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT);
+  auto add0 = ops::Add(s.WithOpName("add0"), feed, feed);
+  auto add1 = ops::Add(s.WithOpName("add1"), feed, feed);
+  auto add2 = ops::Add(s.WithOpName("add2"), add0, add1);
+  auto add3 = ops::Add(s.WithOpName("add3"), add0, add2);
+  auto add4 = ops::Add(s.WithOpName("add4"), add2, add2);
+
+  const std::set<string> all_adds = {"add0", "add1", "add2", "add3", "add4"};
+  DisableImplicitBatchMode();
+
+  {
+    Graph g(OpRegistry::Global());
+    TF_EXPECT_OK(s.ToGraph(&g));
+    RunTest(&g, all_adds, all_adds, all_adds, {all_adds});
+  }
+
+  {
+    // Assigning add1 to CPU to exclude it from the cluster.
+    add1.node()->set_assigned_device_name("/device:CPU:0");
+    Graph g(OpRegistry::Global());
+    TF_EXPECT_OK(s.ToGraph(&g));
+    RunTest(&g, all_adds, all_adds, all_adds, {all_adds - "add1"});
+    add1.node()->set_assigned_device_name("");
+  }
+
+  {
+    // Assigning operations add3 and add4 to another GPU to exclude the
+    // operation from the cluster.
+    constexpr char kGpu0[] = "/device:GPU:0";
+    add0.node()->set_assigned_device_name(kGpu0);
+    add1.node()->set_assigned_device_name(kGpu0);
+    add2.node()->set_assigned_device_name(kGpu0);
+    constexpr char kGpu1[] = "/device:GPU:1";
+    add3.node()->set_assigned_device_name(kGpu1);
+    add4.node()->set_assigned_device_name(kGpu1);
+    Graph g(OpRegistry::Global());
+    TF_EXPECT_OK(s.ToGraph(&g));
+    RunTest(&g, all_adds, all_adds, all_adds, {{"add0", "add1", "add2"}});
+  }
+
+  {
+    // Assigning the operations to two compatible GPU devices resulting in
+    // one cluster with all operations.
+    constexpr char kGpuAny[] = "/device:GPU:*";
+    add3.node()->set_assigned_device_name(kGpuAny);
+    add4.node()->set_assigned_device_name(kGpuAny);
+    Graph g(OpRegistry::Global());
+    TF_EXPECT_OK(s.ToGraph(&g));
+    RunTest(&g, all_adds, all_adds, all_adds, {all_adds});
+  }
+}
+
+TEST_F(SegmentTest, AvoidCycle) {
+  //           feed
+  //          //  \\
+  //       add0    add1
+  //        | \    /
+  //        |  add2
+  //        |  /  \\
+  //       add3    add4
+  //          \    /
+  //          <sink>
+  Scope s = Scope::NewRootScope();
+  auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT);
+  auto add0 = ops::Add(s.WithOpName("add0"), feed, feed);
+  auto add1 = ops::Add(s.WithOpName("add1"), feed, feed);
+  auto add2 = ops::Add(s.WithOpName("add2"), add0, add1);
+  auto add3 = ops::Add(s.WithOpName("add3"), add0, add2);
+  auto add4 = ops::Add(s.WithOpName("add4"), add2, add2);
+  Graph g(OpRegistry::Global());
+  TF_EXPECT_OK(s.ToGraph(&g));
+
+  // add2 is not a TRT candidate so there should be no segments generated.
+  const std::set<string> without_add2 = {"add0", "add1", "add3", "add4"};
+  DisableImplicitBatchMode();
+  RunTest(&g, without_add2, without_add2, without_add2, {});
+}
+
+TEST_F(SegmentTest, Multiple) {
+  //              feed
+  //           //  ||  \\
+  //        add0  add1  add7
+  //        |  \  /     / \\
+  //        |  add2    /   \\
+  //        |   || \   |   ||
+  //        |   ||  add5  add8
+  //        |  /  \ /  \   /
+  //        add3  add4  add6
+  //           \   |   /
+  //             <sink>
+  Scope s = Scope::NewRootScope();
+  auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT);
+  auto add0 = ops::Add(s.WithOpName("add0"), feed, feed);
+  auto add1 = ops::Add(s.WithOpName("add1"), feed, feed);
+  auto add7 = ops::Add(s.WithOpName("add7"), feed, feed);
+  auto add2 = ops::Add(s.WithOpName("add2"), add0, add1);
+  auto add5 = ops::Add(s.WithOpName("add5"), add2, add7);
+  auto add8 = ops::Add(s.WithOpName("add8"), add7, add7);
+  auto add3 = ops::Add(s.WithOpName("add3"), add0, add2);
+  auto add4 = ops::Add(s.WithOpName("add4"), add2, add5);
+  auto add6 = ops::Add(s.WithOpName("add6"), add5, add8);
+  Graph g(OpRegistry::Global());
+  TF_EXPECT_OK(s.ToGraph(&g));
+
+  const std::set<string> all_adds = {"add0", "add1", "add2", "add3", "add4",
+                                     "add5", "add6", "add7", "add8"};
+  // Make add5 not a TRT candidate, and we expect two segments.
+  auto without_add5 = all_adds - "add5";
+  DisableImplicitBatchMode();
+  RunTest(&g, without_add5, without_add5, without_add5,
+          {{"add0", "add1", "add2", "add3"}, {"add6", "add8"}});
+
+  // Make add8 not a candidate and add6 not an input candidate, then all direct
+  // and indirect inputs of add6 will be removed from the segment.
+  auto without_add8 = all_adds - "add8";
+  auto without_add6 = all_adds - "add6";
+  RunTest(&g, without_add8, without_add6, all_adds, {{"add3", "add4"}});
+
+  // Make add3 not a candidate and add0 not an output candidate, then all
+  // direct and indirect outputs of add0 will be removed from the segment.
+  auto without_add3 = all_adds - "add3";
+  auto without_add0 = all_adds - "add0";
+  RunTest(&g, without_add3, all_adds, without_add0, {{"add1", "add7", "add8"}});
+}
+
+TEST_F(SegmentTest, BigIfElse) {
+  //           feed
+  //            ||
+  //           add0
+  //         //    \\
+  //       add1    add4
+  //        ||      ||
+  //       add2    add5
+  //        ||      ||
+  //       add3    add6
+  //         \\    //
+  //           add7
+  //            ||
+  //          <sink>
+  Scope s = Scope::NewRootScope();
+  auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT);
+  auto add0 = ops::Add(s.WithOpName("add0"), feed, feed);
+  auto add1 = ops::Add(s.WithOpName("add1"), add0, add0);
+  auto add2 = ops::Add(s.WithOpName("add2"), add1, add1);
+  auto add3 = ops::Add(s.WithOpName("add3"), add2, add2);
+  auto add4 = ops::Add(s.WithOpName("add4"), add0, add0);
+  auto add5 = ops::Add(s.WithOpName("add5"), add4, add4);
+  auto add6 = ops::Add(s.WithOpName("add6"), add5, add5);
+  auto add7 = ops::Add(s.WithOpName("add7"), add3, add6);
+  Graph g(OpRegistry::Global());
+  TF_EXPECT_OK(s.ToGraph(&g));
+
+  // Make add2 not a TRT candidate, and we expect 2 segments.
+  const std::set<string> all_adds = {"add0", "add1", "add2", "add3",
+                                     "add4", "add5", "add6", "add7"};
+  DisableImplicitBatchMode();
+  RunTest(&g, all_adds - "add2", all_adds, all_adds,
+          {{"add0", "add1"}, {"add3", "add4", "add5", "add6", "add7"}});
+}
+
+TEST_F(SegmentTest, IdentityOps) {
+  Scope s = Scope::NewRootScope();
+  auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT);
+  auto identity0 = ops::Identity(s.WithOpName("identity0"), feed);
+  auto identity1 = ops::Identity(s.WithOpName("identity1"), identity0);
+  auto identity2 = ops::Identity(s.WithOpName("identity2"), identity1);
+  auto identity3 = ops::Identity(s.WithOpName("identity3"), identity2);
+  Graph g(OpRegistry::Global());
+  TF_EXPECT_OK(s.ToGraph(&g));
+
+  const std::set<string> all_identities = {"identity0", "identity1",
+                                           "identity2", "identity3"};
+  // Identity ops are not counted as effective ops in the segment, so no segment
+  // will be formed in this case.
+  DisableImplicitBatchMode();
+  RunTest(&g, all_identities, all_identities, all_identities, {});
+}
+
+// Testing implicit batch mode segmentation: it excludes the add-2 operation
+// with a dynamic non-batch dimension.
+TEST_F(SegmentTest, ExcludeAddWithDynamicNonBatchDimension) {
+  Scope s = Scope::NewRootScope();
+  auto feed_0_shape = ops::Placeholder::Shape(PartialTensorShape({-1, 2, 3}));
+  auto feed_1_shape = ops::Placeholder::Shape(PartialTensorShape({-1, -1, 3}));
+  auto const_val = ops::Const<float>(s, {1.0}, {});
+  auto feed_0 =
+      ops::Placeholder(s.WithOpName("feed-1"), DT_FLOAT, feed_0_shape);
+  auto feed_1 =
+      ops::Placeholder(s.WithOpName("feed-2"), DT_FLOAT, feed_1_shape);
+  auto add_0 = ops::Add(s.WithOpName("add-0"), feed_0, const_val);
+  auto add_1 = ops::Add(s.WithOpName("add-1"), add_0, feed_0);
+  auto add_2 = ops::Add(s.WithOpName("add-2"), const_val, feed_1);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("add-2");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"add-0", "add-1", "add-2"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes,
+          {all_nodes - "add-2"});
+}
+
+// Testing implicit batch mode segmentation: It excludes the reshape operation
+// with a dynamic non-batch output dimension.
+// TODO(bixia): hoist the check for reshape should not change batch size from
+// the converter to the segmenter and add another test case for excluding
+// a reshape without dynamic dimensions involved.
+TEST_F(SegmentTest, ExcludeReshapeWithDynamicNonBatchDimensionInOutput) {
+  Scope s = Scope::NewRootScope();
+  auto feed_0_shape = ops::Placeholder::Shape(PartialTensorShape({-1, 2, 3}));
+  auto const_val = ops::Const<float>(s, {1.0}, {});
+  auto feed_0 =
+      ops::Placeholder(s.WithOpName("feed-1"), DT_FLOAT, feed_0_shape);
+  auto add_0 = ops::Add(s.WithOpName("add-0"), feed_0, const_val);
+  auto reshape = ops::Reshape(s.WithOpName("reshape"), add_0, Input({6, -1}));
+  auto add_1 = ops::Add(s.WithOpName("add-1"), reshape, const_val);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("add-1");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"add-0", "reshape", "add-1"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes, {});
+}
+
+TEST_F(SegmentTest, RankOneCannotUseImplicitBatch) {
+  Scope s = Scope::NewRootScope();
+  auto input_0_shape = ops::Placeholder::Shape(TensorShape({3}));
+  auto input_1_shape = ops::Placeholder::Shape(TensorShape({3}));
+  auto input_0 =
+      ops::Placeholder(s.WithOpName("input-0"), DT_FLOAT, input_0_shape);
+  auto input_1 =
+      ops::Placeholder(s.WithOpName("input-1"), DT_FLOAT, input_1_shape);
+  auto const_val = ops::Const(s.WithOpName("const-scalar"), 1.0f, {});
+  auto output_0 = ops::Add(s.WithOpName("output-0"), input_0, const_val);
+  auto output_1 = ops::Add(s.WithOpName("output-1"), input_1, const_val);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("output-0");
+  item.fetch.push_back("output-1");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"const-scalar", "output-0", "output-1"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes, {});
+}
+
+TEST_F(SegmentTest, TwoChainsDiffBatchSizes) {
+  Scope s = Scope::NewRootScope();
+  auto input_0_shape = ops::Placeholder::Shape(TensorShape({2, 3}));
+  auto input_1_shape = ops::Placeholder::Shape(TensorShape({5, 3}));
+  auto input_0 =
+      ops::Placeholder(s.WithOpName("input-0"), DT_FLOAT, input_0_shape);
+  auto input_1 =
+      ops::Placeholder(s.WithOpName("input-1"), DT_FLOAT, input_1_shape);
+  auto const_val = ops::Const(s.WithOpName("const-scalar"), 1.0f, {});
+  auto output_0 = ops::Add(s.WithOpName("output-0"), input_0, const_val);
+  auto output_1 = ops::Add(s.WithOpName("output-1"), input_1, const_val);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("output-0");
+  item.fetch.push_back("output-1");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"const-scalar", "output-0", "output-1"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes,
+          /*expected_segments=*/{{"output-0", "const-scalar"}});
+
+  // Converter will create engines based on the static batch size
+  EnableImplicitBatchModeForStaticEngine(1);
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes,
+          /*expected_segments=*/{{"output-0", "const-scalar"}});
+}
+
+TEST_F(SegmentTest, SameRankImplicitBroadcastingStaticBatchSize) {
+  Scope s = Scope::NewRootScope();
+  auto input_0_shape = ops::Placeholder::Shape(TensorShape({2, 3, 1}));
+  auto input_1_shape = ops::Placeholder::Shape(TensorShape({1, 3, 4}));
+  auto input_2_shape = ops::Placeholder::Shape(TensorShape({2, 3, 4}));
+  auto input_0 =
+      ops::Placeholder(s.WithOpName("input-0"), DT_FLOAT, input_0_shape);
+  auto input_1 =
+      ops::Placeholder(s.WithOpName("input-1"), DT_FLOAT, input_1_shape);
+  auto input_2 =
+      ops::Placeholder(s.WithOpName("input-2"), DT_FLOAT, input_2_shape);
+  auto multiple = ops::Mul(s.WithOpName("multiple"), input_2, input_2);
+  auto output_0 = ops::Add(s.WithOpName("output-0"), input_0, multiple);
+  auto output_1 = ops::Add(s.WithOpName("output-1"), input_1, multiple);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("output-0");
+  item.fetch.push_back("output-1");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"multiple", "output-0", "output-1"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes,
+          {all_nodes});
+}
+
+TEST_F(SegmentTest, SameRankImplicitBroadcastingDynamicBatchSize) {
+  Scope s = Scope::NewRootScope();
+  auto input_0_shape = ops::Placeholder::Shape(PartialTensorShape({-1, 2}));
+  auto input_1_shape = ops::Placeholder::Shape(TensorShape({1, 2}));
+  auto input_0 =
+      ops::Placeholder(s.WithOpName("input-0"), DT_FLOAT, input_0_shape);
+  auto input_1 =
+      ops::Placeholder(s.WithOpName("input-1"), DT_FLOAT, input_1_shape);
+  auto const_val = ops::Const(s.WithOpName("const-val"), 1.0f, {1, 1});
+  auto add_0 = ops::Add(s.WithOpName("add-0"), input_0, const_val);
+  auto output_0 = ops::Add(s.WithOpName("output-0"), input_0, add_0);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("output-0");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"const-val", "add-0", "output-0"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes,
+          {{"const-val", "add-0", "output-0"}});
+}
+
+TEST_F(SegmentTest, IncompatibleBatchSizes) {
+  Scope s = Scope::NewRootScope();
+  auto input_0_shape = ops::Placeholder::Shape(PartialTensorShape({-1, 2}));
+  auto input_1_shape = ops::Placeholder::Shape(TensorShape({2, 2}));
+  auto input_0 =
+      ops::Placeholder(s.WithOpName("input-0"), DT_FLOAT, input_0_shape);
+  auto input_1 =
+      ops::Placeholder(s.WithOpName("input-1"), DT_FLOAT, input_1_shape);
+  auto const_val = ops::Const(s.WithOpName("const-val"), 1.0f, {2, 2});
+  auto add_0 = ops::Add(s.WithOpName("add-0"), input_0, const_val);
+  auto output_0 = ops::Add(s.WithOpName("output-0"), input_0, add_0);
+
+  grappler::GrapplerItem item;
+  item.fetch.push_back("output-0");
+  TF_EXPECT_OK(s.ToGraphDef(&item.graph));
+
+  grappler::GraphProperties static_graph_properties(item);
+  TF_EXPECT_OK(static_graph_properties.InferStatically(true));
+
+  Graph g(OpRegistry::Global());
+  TF_CHECK_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), item.graph, &g));
+
+  const std::set<string> all_nodes = {"const-val", "add-0", "output-0"};
+  EnableImplicitBatchModeForStaticEngine();
+  RunTest(&g, &static_graph_properties, all_nodes, all_nodes, all_nodes, {});
+}
+}  // namespace test
+}  // namespace segment
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/segment/union_find.cc

@@ -0,0 +1,154 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2tensorrt/segment/union_find.h"
+
+#include "absl/strings/str_format.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace segment {
+
+namespace {
+template <typename T>
+inline bool CheckIfCompatible(const std::optional<T>& a,
+                              const std::optional<T>& b) {
+  if (a.has_value() && b.has_value()) {
+    return *a == *b;
+  }
+  return true;
+}
+
+template <typename T>
+inline bool UnifyValues(std::optional<T>& a, std::optional<T>& b) {
+  if (a.has_value()) {
+    b = a;
+  } else {
+    a = b;
+  }
+  return true;
+}
+
+template <typename T>
+inline std::optional<T> MergeCompatible(const std::optional<T>& a,
+                                        const std::optional<T>& b) {
+  DCHECK(CheckIfCompatible(a, b));
+  return a.has_value() ? a : b;
+}
+
+}  // namespace
+
+ClusterBatchSize::ClusterBatchSize()
+    : batch_size_(std::nullopt), max_batch_size_(std::nullopt) {}
+
+bool ClusterBatchSize::operator==(const ClusterBatchSize& other) {
+  return batch_size_ == other.batch_size_ &&
+         max_batch_size_ == other.max_batch_size_;
+}
+
+ClusterBatchSize& ClusterBatchSize::SetBatchSize(int batch_size) {
+  SetBatchSize(static_cast<std::optional<int>>(batch_size));
+  return *this;
+}
+
+ClusterBatchSize& ClusterBatchSize::SetBatchSize(
+    const std::optional<int>& batch_size) {
+  batch_size_ = MergeCompatible<int>(batch_size_, batch_size);
+  if (batch_size_.has_value() && batch_size_.value() >= 0) {
+    SetMaxBatchSize(batch_size_);
+  }
+  return *this;
+}
+
+bool ClusterBatchSize::HasBatchSize() const { return batch_size_.has_value(); }
+
+int ClusterBatchSize::GetBatchSize() const {
+  DCHECK(HasBatchSize());
+  return batch_size_.value();
+}
+
+ClusterBatchSize& ClusterBatchSize::SetMaxBatchSize(int max_batch_size) {
+  SetBatchSize(static_cast<std::optional<int>>(max_batch_size));
+  return *this;
+}
+
+ClusterBatchSize& ClusterBatchSize::SetMaxBatchSize(
+    const std::optional<int>& max_batch_size) {
+  max_batch_size_ = MergeCompatible<int>(max_batch_size_, max_batch_size);
+  return *this;
+}
+
+std::optional<int> ClusterBatchSize::GetOptionalMaxBatchSize() const {
+  return max_batch_size_;
+}
+
+bool ClusterBatchSize::MergeIfCompatible(const ClusterBatchSize& other) {
+  if (!CheckIfCompatible(batch_size_, other.batch_size_) ||
+      !CheckIfCompatible(max_batch_size_, other.max_batch_size_)) {
+    return false;
+  }
+
+  SetBatchSize(other.batch_size_);
+  SetMaxBatchSize(other.max_batch_size_);
+  return true;
+}
+
+string ClusterBatchSize::ToString() const {
+  string s;
+  const auto append_optional_num = [&](const std::optional<int>& num) {
+    if (num.has_value()) {
+      absl::StrAppendFormat(&s, "%d", num.value());
+    } else {
+      absl::StrAppendFormat(&s, "?");
+    }
+  };
+  absl::StrAppendFormat(&s, "batch_size=");
+  append_optional_num(batch_size_);
+  absl::StrAppendFormat(&s, ", max_batch_size=");
+  append_optional_num(max_batch_size_);
+  return s;
+}
+
+ClusterProperty::ClusterProperty(const ClusterBatchSize& batch_size,
+                                 const DeviceNameUtils::ParsedName& device_name)
+    : batch_size_(batch_size), device_name_(device_name) {}
+
+Status ClusterProperty::Merge(const ClusterProperty& other) {
+  ClusterBatchSize merged_batch_size(batch_size_);
+  if (!merged_batch_size.MergeIfCompatible(other.batch_size_)) {
+    return errors::Internal(
+        "trying to merge clusters with incompatible batch sizes.");
+  }
+
+  std::optional<DeviceNameUtils::ParsedName> merged_device_name =
+      MergeIfCompatible(device_name_, other.device_name_);
+  if (!merged_device_name.has_value()) {
+    return errors::Internal(
+        "trying to merge clusters with incompatible device assignment.");
+  }
+
+  batch_size_ = std::move(merged_batch_size);
+  device_name_ = std::move(merged_device_name.value());
+  return OkStatus();
+}
+
+}  // namespace segment
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/segment/union_find.h

@@ -0,0 +1,218 @@
+/* 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_SEGMENT_UNION_FIND_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_
+
+#include "absl/types/optional.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace segment {
+
+// ClusterBatchSize is a data structure to record the batch size we have seen
+// for a cluster during segmentation.
+//
+// With the help of shape inference, all the dynamic batch sizes are converted
+// to a negative integer number.
+// If the number is -1, then nothing is known about the dynamic batch size.
+// Ideally, we should not put nodes with -1 batch size into the same cluster,
+// as they will likely have different batch sizes at runtime. However, we
+// currently treat -1 as an equivalent class for simple implementation. We may
+// need to revise this if it causes performance issues.
+// If the number is strictly less than -1, then it represents a equivalent
+// class. It is inferred that all the nodes with the same equivalent class
+// (strictly less than -1) shall have the same batch size at runtime.
+//
+// When constructing clusters for implicit batch mode, we support both
+// dynamic batch sizes and static batch sizes. As all the nodes inside the same
+// cluster shall have the same batch size at runtime, we restrict nodes inside a
+// cluster to either have the same dynamic batch size equivalent class or the
+// same static batch size value.
+//
+// Besides, all the nodes with an annotated max batch size inside the same
+// cluster shall have the same annotated max batch size. (It is allowed if
+// part or all the nodes inside the cluster doesn't have annotated max batch
+// size). Static batch sizes are treated as max batch size annotations. The
+// converter max batch size is used for an OP with a dynamic batch size and no
+// annotated max batch size.
+//
+// cluster:  a = a1[1,3] + a1[1,3]
+// ClusterBatchSize: batch_size_ = 1
+//                   max_batch_size_ = 1
+//
+// cluster:  b = b1[-1,3] + b2[-1, 3]
+// ClusterBatchSize: batch_size_ = -1
+//                   max_batch_size_ = null
+//
+// cluster:  c = c1[-2,3] + c2[-2, 3](max_batch_size=100)
+// ClusterBatchSize: batch_size_ = -2
+//                   max_batch_size_ = 100
+//
+// When constructing cluster for explicit batch mode, all ClusterBatchSize is
+// irrelevant.
+//
+
+class ClusterBatchSize {
+ public:
+  ClusterBatchSize();
+
+  bool operator==(const ClusterBatchSize& other);
+  bool operator!=(const ClusterBatchSize& other) { return !(*this == other); }
+
+  // Sets the batch size assuming that the object doesn't have a batch size yet:
+  //   A non-negative input representing a static batch size value.
+  //   A negative input representing a dynamic batch size equivalent class.
+  ClusterBatchSize& SetBatchSize(int batch_size);
+  bool HasBatchSize() const;
+  int GetBatchSize() const;
+
+  // Sets the max batch size assuming that the object doesn't have a max batch
+  // size yet.
+  ClusterBatchSize& SetMaxBatchSize(int max_batch_size);
+  std::optional<int> GetOptionalMaxBatchSize() const;
+
+  // Merge `other` into the current ClusterBatchSize if the two are not
+  // conflicting. Two ClusterBatchSizes are conflicting iff they both have a
+  // value and their values are different.
+  bool MergeIfCompatible(const ClusterBatchSize& other);
+
+  // Returns a string for the batch size and the annotated max batch size.
+  // For the batch size:
+  //   If the object has a static batch size, return a string representing a
+  //     non-negative integer.
+  //   If the object has a dynamic batch size, return a string representing a
+  //     negative integer as an equivalent class.
+  //   If the object doesn't have a batch size yet, return "?".
+  // For the annotated max batch size:
+  //   If the cluster has annotated max batch size in at least one of the nodes,
+  //     return a string representing the annotated max batch size. Otherwise,
+  //     return "?".
+  std::string ToString() const;
+
+ private:
+  ClusterBatchSize& SetBatchSize(const std::optional<int>& batch_size);
+  ClusterBatchSize& SetMaxBatchSize(const std::optional<int>& batch_size);
+
+  std::optional<int> batch_size_;
+  std::optional<int> max_batch_size_;
+};
+
+inline std::ostream& operator<<(std::ostream& os,
+                                const ClusterBatchSize& batch_size) {
+  return os << batch_size.ToString();
+}
+
+// Represents the accumulated properties of a cluster during segmentation,
+// including information about batch size and device assignment. Clusters shall
+// have compatible properties in order to be merged together.
+class ClusterProperty {
+ public:
+  ClusterProperty() {}
+  ClusterProperty(const ClusterBatchSize& batch_size,
+                  const DeviceNameUtils::ParsedName& device_name);
+
+  // Returns the batch size of the cluster and compresses the path from this
+  // object to the root object.
+  const ClusterBatchSize& BatchSize() const { return batch_size_; }
+
+  // Returns the device name of the cluster and compresses the path from this
+  // object to the root object.
+  const DeviceNameUtils::ParsedName& DeviceName() const { return device_name_; }
+
+  Status Merge(const ClusterProperty& other);
+
+ private:
+  ClusterBatchSize batch_size_;
+  DeviceNameUtils::ParsedName device_name_;
+};
+
+// Represents a disjoint set of copyable value with type T and accumulated
+// property of the values with type P. Most of the methods in this class are
+// side-effecting as they also compress the path from the object to the parent
+// of its containing set.
+template <typename T, typename P = ClusterProperty>
+class UnionFind {
+ public:
+  UnionFind() : size_(1), parent_(nullptr) {}
+  UnionFind(const T& v, const P& p)
+      : size_(1), parent_(nullptr), value_(v), property_(p) {}
+  UnionFind(const T& v, P&& p)
+      : size_(1), parent_(nullptr), value_(v), property_(p) {}
+
+  // Returns the number of elements in the set and compresses the path from
+  // this object to the root of the set.
+  int Size() { return FindRoot()->size_; }
+
+  // Returns the accumulated property of all the elements in the set and
+  // compresses the path from this object to the root of the set.
+  const P& Property() { return FindRoot()->property_; }
+
+  // Merges this set with 'other'. This updates the size_ and property_ of the
+  // set. The size_ and property_ of 'other' becomes inaccessible as only the
+  // size_ and property_ of the root of the set is accessible.
+  Status Merge(UnionFind* other);
+
+  // Retrieves the value for the root of the set.
+  const T& ParentValue() { return FindRoot()->value_; }
+
+  // Returns the value for the object.
+  const T& Value() const { return value_; }
+
+ private:
+  // Returns the root object for the set and compresses the path from this
+  // object to the root object.
+  UnionFind* FindRoot();
+
+  int size_;
+  UnionFind* parent_;
+  T value_;
+  P property_;
+};
+
+template <typename T, typename P>
+Status UnionFind<T, P>::Merge(UnionFind* other) {
+  UnionFind<T>* a = FindRoot();
+  UnionFind<T>* b = other->FindRoot();
+  if (a == b) return OkStatus();
+
+  P merged_property(a->property_);
+  TF_RETURN_IF_ERROR(merged_property.Merge(b->property_));
+  b->parent_ = a;
+  a->size_ += b->size_;
+  a->property_ = std::move(merged_property);
+  return OkStatus();
+}
+
+template <typename T, typename P>
+UnionFind<T, P>* UnionFind<T, P>::FindRoot() {
+  if (!parent_) return this;
+  // Path compression: update intermediate nodes to point to the root of the
+  // equivalence class.
+  parent_ = parent_->FindRoot();
+  return parent_;
+}
+
+}  // namespace segment
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_