1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include "tensorflow/core/common_runtime/gpu/gpu_stream_util.h" 17 18 #include "tensorflow/cc/ops/sendrecv_ops.h" 19 #include "tensorflow/cc/ops/standard_ops.h" 20 #include "tensorflow/core/framework/op.h" 21 #include "tensorflow/core/framework/types.pb.h" 22 #include "tensorflow/core/graph/graph_def_builder.h" 23 #include "tensorflow/core/graph/node_builder.h" 24 #include "tensorflow/core/kernels/ops_testutil.h" 25 #include "tensorflow/core/kernels/ops_util.h" 26 #include "tensorflow/core/lib/core/status_test_util.h" 27 #include "tensorflow/core/platform/test.h" 28 29 namespace tensorflow { 30 namespace { 31 32 class GpuStreamUtilTest : public OpsTestBase { 33 protected: 34 }; 35 36 TEST_F(GpuStreamUtilTest, BogusOpts) { 37 auto root = Scope::NewRootScope().ExitOnError(); 38 Graph g(OpRegistry::Global()); 39 TF_ASSERT_OK(root.ToGraph(&g)); 40 std::unordered_map<int, int> node_to_stream_id; 41 gpu_stream_util::AssignStreamsOpts opts; 42 Status status; 43 status = gpu_stream_util::AssignStreams(nullptr, opts, &node_to_stream_id); 44 EXPECT_FALSE(status.ok()); 45 status = gpu_stream_util::AssignStreams(&g, opts, nullptr); 46 EXPECT_FALSE(status.ok()); 47 opts.max_streams = 0; 48 status = gpu_stream_util::AssignStreams(&g, opts, &node_to_stream_id); 49 EXPECT_FALSE(status.ok()); 50 opts.max_streams = 1; 51 opts.compute_stream = 5; 52 status = gpu_stream_util::AssignStreams(&g, opts, &node_to_stream_id); 53 EXPECT_FALSE(status.ok()); 54 } 55 56 TEST_F(GpuStreamUtilTest, EmptyGraph) { 57 auto root = Scope::NewRootScope().ExitOnError(); 58 Graph g(OpRegistry::Global()); 59 TF_ASSERT_OK(root.ToGraph(&g)); 60 std::unordered_map<int, int> node_to_stream_id; 61 gpu_stream_util::AssignStreamsOpts opts; 62 TF_ASSERT_OK(gpu_stream_util::AssignStreams(&g, opts, &node_to_stream_id)); 63 EXPECT_EQ(2, node_to_stream_id.size()); // _SOURCE and _SINK 64 } 65 66 TEST_F(GpuStreamUtilTest, SimpleGraphOneStream) { 67 auto root = Scope::DisabledShapeInferenceScope().ExitOnError(); 68 ops::MatMul(root, {}, {}); 69 Graph g(OpRegistry::Global()); 70 TF_ASSERT_OK(root.ToGraph(&g)); 71 72 std::unordered_map<int, int> node_to_stream_id; 73 gpu_stream_util::AssignStreamsOpts opts; 74 TF_ASSERT_OK(gpu_stream_util::AssignStreams(&g, opts, &node_to_stream_id)); 75 76 // There should be 5 nodes assigned. 77 EXPECT_EQ(5, node_to_stream_id.size()); 78 79 // All of them should have stream 0. 80 for (const auto& it : node_to_stream_id) { 81 EXPECT_EQ(0, it.second); 82 } 83 } 84 85 TEST_F(GpuStreamUtilTest, SimpleGraphManyStreams) { 86 auto root = Scope::DisabledShapeInferenceScope().ExitOnError(); 87 ops::MatMul(root, {}, {}); 88 Graph g(OpRegistry::Global()); 89 TF_ASSERT_OK(root.ToGraph(&g)); 90 91 std::unordered_map<int, int> node_to_stream_id; 92 gpu_stream_util::AssignStreamsOpts opts; 93 opts.max_streams = 3; 94 TF_ASSERT_OK(gpu_stream_util::AssignStreams(&g, opts, &node_to_stream_id)); 95 96 // There should be 5 nodes assigned. 97 EXPECT_EQ(5, node_to_stream_id.size()); 98 99 // All of them should have a stream in the range [0..max_streams). 100 for (const auto& it : node_to_stream_id) { 101 EXPECT_GE(it.second, 0); 102 EXPECT_LT(it.second, opts.max_streams); 103 } 104 } 105 106 TEST_F(GpuStreamUtilTest, StreamOverrides) { 107 auto root = Scope::DisabledShapeInferenceScope().ExitOnError(); 108 ops::_Recv(root.WithOpName("input"), DT_FLOAT, "input", "/cpu:0", 0, 109 "/device:GPU:0"); 110 Output n = ops::MatMul(root, {}, {}); 111 ops::_Send(root.WithOpName("output"), n, "output", "/device:GPU:0", 0, 112 "/cpu:0"); 113 Graph g(OpRegistry::Global()); 114 TF_ASSERT_OK(root.ToGraph(&g)); 115 116 // Perform stream assignment using a large number of streams, but with 117 // op types constrained to specific streams. 118 std::unordered_map<int, int> node_to_stream_id; 119 gpu_stream_util::AssignStreamsOpts opts; 120 opts.max_streams = 100; 121 opts.const_stream = 90; 122 opts.send_stream = 91; 123 opts.recv_stream = 92; 124 opts.compute_stream = 93; 125 TF_ASSERT_OK(gpu_stream_util::AssignStreams(&g, opts, &node_to_stream_id)); 126 127 // There should be 7 nodes assigned. 128 EXPECT_EQ(7, node_to_stream_id.size()); // including _SOURCE and _SINK 129 130 // Nodes should be assigned to streams by op type. 131 for (const auto& it : node_to_stream_id) { 132 Node* n = g.FindNodeId(it.first); 133 const string& op = n->type_string(); 134 const int stream = it.second; 135 if (op == "Const") { 136 EXPECT_EQ(stream, 90); 137 } else if (op == "_Send") { 138 EXPECT_EQ(stream, 91); 139 } else if (op == "_Recv") { 140 EXPECT_EQ(stream, 92); 141 } else { // Compute. 142 EXPECT_EQ(stream, 93); 143 } 144 } 145 } 146 147 } // namespace 148 } // namespace tensorflow 149
