diff options
author | Manjunath Kudlur <keveman@gmail.com> | 2015-11-06 16:27:58 -0800 |
---|---|---|
committer | Manjunath Kudlur <keveman@gmail.com> | 2015-11-06 16:27:58 -0800 |
commit | f41959ccb2d9d4c722fe8fc3351401d53bcf4900 (patch) | |
tree | ef0ca22cb2a5ac4bdec9d080d8e0788a53ed496d /tensorflow/core/common_runtime/gpu/gpu_stream_util.h |
TensorFlow: Initial commit of TensorFlow library.
TensorFlow is an open source software library for numerical computation
using data flow graphs.
Base CL: 107276108
Diffstat (limited to 'tensorflow/core/common_runtime/gpu/gpu_stream_util.h')
-rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_stream_util.h | 30 |
1 files changed, 30 insertions, 0 deletions
diff --git a/tensorflow/core/common_runtime/gpu/gpu_stream_util.h b/tensorflow/core/common_runtime/gpu/gpu_stream_util.h new file mode 100644 index 0000000000..e1c623382c --- /dev/null +++ b/tensorflow/core/common_runtime/gpu/gpu_stream_util.h @@ -0,0 +1,30 @@ +#ifndef TENSORFLOW_COMMON_RUNTIME_GPU_GPU_STREAM_UTIL_H_ +#define TENSORFLOW_COMMON_RUNTIME_GPU_GPU_STREAM_UTIL_H_ + +#include <unordered_map> + +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/public/status.h" + +namespace tensorflow { +namespace gpu_stream_util { + +struct AssignStreamsOpts { + int32 max_streams = 1; + // The following options specify a stream to use for specific op + // types. The value -1 allows ops to be assigned to any stream. + int32 send_stream = -1; + int32 recv_stream = -1; + int32 const_stream = -1; + int32 compute_stream = -1; +}; + +// Given the input graph, assigns every node in the graph with a +// stream_id that should be used. +Status AssignStreams(const Graph* graph, const AssignStreamsOpts& opts, + std::unordered_map<int, int>* node_to_stream_id); + +} // namespace gpu_stream_util +} // namespace tensorflow + +#endif // TENSORFLOW_COMMON_RUNTIME_GPU_GPU_STREAM_UTIL_H_ |