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/core/platform/thread_annotations.h" 23 #include "tensorflow/stream_executor/gpu/gpu_driver.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(); SetPriority(int priority)51 void SetPriority(int priority) { priority_ = priority; } priority()52 int priority() const { return priority_; } 53 54 // Explicitly destroy the CUDA resources associated with this stream, used by 55 // StreamExecutor::DeallocateStream(). 56 void Destroy(); 57 58 // Returns true if no work is pending or executing on the stream. 59 bool IsIdle() const; 60 61 // Retrieves an event which indicates that all work enqueued into the stream 62 // has completed. Ownership of the event is not transferred to the caller, the 63 // event is owned by this stream. completed_event()64 GpuEventHandle* completed_event() { return &completed_event_; } 65 66 // Returns the GpuStreamHandle value for passing to the CUDA API. 67 // 68 // Precond: this GpuStream has been allocated (otherwise passing a nullptr 69 // into the NVIDIA library causes difficult-to-understand faults). gpu_stream()70 GpuStreamHandle gpu_stream() const { 71 DCHECK(gpu_stream_ != nullptr); 72 return const_cast<GpuStreamHandle>(gpu_stream_); 73 } 74 75 // TODO(timshen): Migrate away and remove this function. cuda_stream()76 GpuStreamHandle cuda_stream() const { return gpu_stream(); } 77 parent()78 GpuExecutor* parent() const { return parent_; } 79 80 private: 81 GpuExecutor* parent_; // Executor that spawned this stream. 82 GpuStreamHandle gpu_stream_; // Wrapped CUDA stream handle. 83 int priority_ = 0; 84 85 // Event that indicates this stream has completed. 86 GpuEventHandle completed_event_ = nullptr; 87 }; 88 89 // Helper functions to simplify extremely common flows. 90 // Converts a Stream to the underlying GpuStream implementation. 91 GpuStream* AsGpuStream(Stream* stream); 92 93 // Extracts a GpuStreamHandle from a GpuStream-backed Stream object. 94 GpuStreamHandle AsGpuStreamValue(Stream* stream); 95 96 } // namespace gpu 97 } // namespace stream_executor 98 99 #endif // TENSORFLOW_STREAM_EXECUTOR_GPU_GPU_STREAM_H_ 100