1 /* Copyright 2018 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/compiler/xla/service/stream_pool.h"
17 
18 #include <memory>
19 
20 #include "tensorflow/compiler/xla/test_helpers.h"
21 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
22 
23 namespace xla {
24 namespace {
25 
26 class StreamPoolTest : public ::testing::Test {
27  protected:
NewStreamExecutor()28   std::unique_ptr<se::StreamExecutor> NewStreamExecutor() {
29     se::Platform* platform =
30         se::MultiPlatformManager::PlatformWithName("Host").ConsumeValueOrDie();
31     se::StreamExecutorConfig config(/*ordinal=*/0);
32     return platform->GetUncachedExecutor(config).ConsumeValueOrDie();
33   }
34 };
35 
TEST_F(StreamPoolTest,EmptyPool)36 TEST_F(StreamPoolTest, EmptyPool) { StreamPool pool; }
37 
TEST_F(StreamPoolTest,OneStreamPool)38 TEST_F(StreamPoolTest, OneStreamPool) {
39   std::unique_ptr<se::StreamExecutor> executor = NewStreamExecutor();
40   StreamPool pool;
41 
42   // Borrow and return a stream.
43   StreamPool::Ptr stream1 = pool.BorrowStream(executor.get());
44   se::Stream* stream1_ptr = stream1.get();
45   EXPECT_TRUE(stream1->ok());
46   stream1 = nullptr;
47 
48   // Borrow and return another stream.
49   StreamPool::Ptr stream2 = pool.BorrowStream(executor.get());
50   se::Stream* stream2_ptr = stream2.get();
51   EXPECT_TRUE(stream2->ok());
52   stream2 = nullptr;
53 
54   // The underlying streams should be the same, since stream1 was the
55   // only stream available in the pool when stream2 was borrowed.
56   EXPECT_EQ(stream1_ptr, stream2_ptr);
57 }
58 
TEST_F(StreamPoolTest,TwoStreamPool)59 TEST_F(StreamPoolTest, TwoStreamPool) {
60   std::unique_ptr<se::StreamExecutor> executor = NewStreamExecutor();
61   StreamPool pool;
62 
63   // Borrow two streams.
64   StreamPool::Ptr stream1 = pool.BorrowStream(executor.get());
65   se::Stream* stream1_ptr = stream1.get();
66   EXPECT_TRUE(stream1->ok());
67   StreamPool::Ptr stream2 = pool.BorrowStream(executor.get());
68   se::Stream* stream2_ptr = stream2.get();
69   EXPECT_TRUE(stream2->ok());
70 
71   // The underlying streams should be different, since we haven't
72   // returned either of them yet.
73   EXPECT_NE(stream1_ptr, stream2_ptr);
74 
75   // Return stream1 and borrow stream3.
76   stream1 = nullptr;
77   StreamPool::Ptr stream3 = pool.BorrowStream(executor.get());
78   se::Stream* stream3_ptr = stream3.get();
79   EXPECT_TRUE(stream3->ok());
80 
81   // stream1 and stream3 should be the same.
82   EXPECT_EQ(stream1_ptr, stream3_ptr);
83   EXPECT_NE(stream2_ptr, stream3_ptr);
84 
85   // Return stream2, and borrow stream4.
86   stream2 = nullptr;
87   StreamPool::Ptr stream4 = pool.BorrowStream(executor.get());
88   se::Stream* stream4_ptr = stream4.get();
89   EXPECT_TRUE(stream4->ok());
90 
91   // Stream2 and stream4 should be the same.
92   EXPECT_EQ(stream2_ptr, stream4_ptr);
93   EXPECT_NE(stream3_ptr, stream4_ptr);
94 }
95 
TEST_F(StreamPoolTest,BadStreamDiscarded)96 TEST_F(StreamPoolTest, BadStreamDiscarded) {
97   std::unique_ptr<se::StreamExecutor> executor = NewStreamExecutor();
98   StreamPool pool;
99 
100   // Borrow a stream.
101   StreamPool::Ptr stream1 = pool.BorrowStream(executor.get());
102   EXPECT_TRUE(stream1->ok());
103 
104   // Force an error on the stream; here we call a method that requires
105   // DNN support, which we know the Host platform doesn't support.
106   stream1->ThenDepthConcatenate({}, {}, nullptr);
107   EXPECT_FALSE(stream1->ok());
108 
109   // Return stream1 and borrow stream2.
110   stream1 = nullptr;
111   StreamPool::Ptr stream2 = pool.BorrowStream(executor.get());
112   se::Stream* stream2_ptr = stream2.get();
113   EXPECT_TRUE(stream2->ok());
114 
115   // The underlying streams should be different. They would have been
116   // the same, but since we forced an error on stream1, it cannot be
117   // put back into the pool. Sadly we can't just check:
118   //    EXPECT_NE(stream1_ptr, stream2_ptr);
119   //
120   // The above should hold logically, but it may fail if the new
121   // stream instance allocated for stream2 happens to reside in the
122   // same memory address as stream1, which has been deleted.
123   //
124   // The check that stream2->ok() serves as a good-enough check.
125 
126   // Return stream2 and borrow stream3. The previous error on stream1
127   // has no effect on these streams, and they are the same.
128   stream2 = nullptr;
129   StreamPool::Ptr stream3 = pool.BorrowStream(executor.get());
130   se::Stream* stream3_ptr = stream3.get();
131   EXPECT_TRUE(stream3->ok());
132   EXPECT_EQ(stream2_ptr, stream3_ptr);
133 }
134 
TEST_F(StreamPoolTest,BadStreamAfterReturnDiscarded)135 TEST_F(StreamPoolTest, BadStreamAfterReturnDiscarded) {
136   std::unique_ptr<se::StreamExecutor> executor = NewStreamExecutor();
137   StreamPool pool;
138 
139   // Borrow a stream.
140   StreamPool::Ptr stream1 = pool.BorrowStream(executor.get());
141   EXPECT_TRUE(stream1->ok());
142 
143   // Return the stream, but hold a handle to it.
144   se::Stream* stream1_ptr = stream1.get();
145   stream1 = nullptr;
146 
147   // Now stream1 is back in the pool, force an error on the stream. Here we call
148   // a method that requires DNN support, which we know the Host platform doesn't
149   // support.
150   stream1_ptr->ThenDepthConcatenate({}, {}, nullptr);
151   EXPECT_FALSE(stream1_ptr->ok());
152 
153   // Borrow stream2.
154   StreamPool::Ptr stream2 = pool.BorrowStream(executor.get());
155   EXPECT_TRUE(stream2->ok());
156 
157   // The underlying streams should be different. They would have been
158   // the same, but since we forced an error on stream1, it cannot be
159   // put back into the pool. Sadly we can't just check:
160   //    EXPECT_NE(stream1_ptr, stream2_ptr);
161   //
162   // The above should hold logically, but it may fail if the new
163   // stream instance allocated for stream2 happens to reside in the
164   // same memory address as stream1, which has been deleted.
165   //
166   // The check that stream2->ok() serves as a good-enough check.
167 }
168 
169 }  // namespace
170 }  // namespace xla
171