diff options
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_ |