• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2019 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 // Defines the GpuStream type - the CUDA-specific implementation of the generic
17 // StreamExecutor Stream interface.
18 
19 #ifndef TENSORFLOW_STREAM_EXECUTOR_GPU_GPU_STREAM_H_
20 #define TENSORFLOW_STREAM_EXECUTOR_GPU_GPU_STREAM_H_
21 
22 #include "tensorflow/stream_executor/gpu/gpu_driver.h"
23 #include "tensorflow/stream_executor/platform/thread_annotations.h"
24 #include "tensorflow/stream_executor/stream_executor_internal.h"
25 
26 namespace stream_executor {
27 namespace gpu {
28 
29 class GpuExecutor;
30 
31 // Wraps a GpuStreamHandle in order to satisfy the platform-independent
32 // StreamInterface.
33 //
34 // Thread-safe post-initialization.
35 class GpuStream : public internal::StreamInterface {
36  public:
GpuStream(GpuExecutor * parent)37   explicit GpuStream(GpuExecutor* parent)
38       : parent_(parent), gpu_stream_(nullptr), completed_event_(nullptr) {}
39 
40   // Note: teardown is handled by a parent's call to DeallocateStream.
~GpuStream()41   ~GpuStream() override {}
42 
GpuStreamHack()43   void* GpuStreamHack() override { return gpu_stream_; }
GpuStreamMemberHack()44   void** GpuStreamMemberHack() override {
45     return reinterpret_cast<void**>(&gpu_stream_);
46   }
47 
48   // Explicitly initialize the CUDA resources associated with this stream, used
49   // by StreamExecutor::AllocateStream().
50   bool Init();
51 
52   // Explicitly destroy the CUDA resources associated with this stream, used by
53   // StreamExecutor::DeallocateStream().
54   void Destroy();
55 
56   // Returns true if no work is pending or executing on the stream.
57   bool IsIdle() const;
58 
59   // Retrieves an event which indicates that all work enqueued into the stream
60   // has completed. Ownership of the event is not transferred to the caller, the
61   // event is owned by this stream.
completed_event()62   GpuEventHandle* completed_event() { return &completed_event_; }
63 
64   // Returns the GpuStreamHandle value for passing to the CUDA API.
65   //
66   // Precond: this GpuStream has been allocated (otherwise passing a nullptr
67   // into the NVIDIA library causes difficult-to-understand faults).
gpu_stream()68   GpuStreamHandle gpu_stream() const {
69     DCHECK(gpu_stream_ != nullptr);
70     return const_cast<GpuStreamHandle>(gpu_stream_);
71   }
72 
73   // TODO(timshen): Migrate away and remove this function.
cuda_stream()74   GpuStreamHandle cuda_stream() const { return gpu_stream(); }
75 
parent()76   GpuExecutor* parent() const { return parent_; }
77 
78  private:
79   GpuExecutor* parent_;         // Executor that spawned this stream.
80   GpuStreamHandle gpu_stream_;  // Wrapped CUDA stream handle.
81 
82   // Event that indicates this stream has completed.
83   GpuEventHandle completed_event_ = nullptr;
84 };
85 
86 // Helper functions to simplify extremely common flows.
87 // Converts a Stream to the underlying GpuStream implementation.
88 GpuStream* AsGpuStream(Stream* stream);
89 
90 // Extracts a GpuStreamHandle from a GpuStream-backed Stream object.
91 GpuStreamHandle AsGpuStreamValue(Stream* stream);
92 
93 }  // namespace gpu
94 }  // namespace stream_executor
95 
96 #endif  // TENSORFLOW_STREAM_EXECUTOR_GPU_GPU_STREAM_H_
97