diff options
| author |  Wen-Heng (Jack) Chung <whchung@gmail.com> | 2018-07-11 18:35:43 +0000 |
|---|---|---|
| committer | Wen-Heng (Jack) Chung <whchung@gmail.com> | 2018-07-11 20:34:36 +0000 |
| commit | 456aaa2fdbf821296a31f5493955f4653ae119dd (patch) | |
| tree | 68f2f260cb66dd135a9c207006d721910d272e36 /tensorflow/stream_executor/stream.h | |
| parent | 135e419e780423a888ddd45e479129493336c52b (diff) | |
[ROCm] Interface changes for pooling APIs in StreamExecutor
Due to the design of MIOpen, the DNN library on ROCm platform, an instance of
ScratchAllocator has to be passed into pooling routines. This commit address
such interface changes and the implementation in CUDA StreamExecutor.
Diffstat (limited to 'tensorflow/stream_executor/stream.h')
| -rw-r--r-- | tensorflow/stream_executor/stream.h | 21 |
1 files changed, 14 insertions, 7 deletions
diff --git a/tensorflow/stream_executor/stream.h b/tensorflow/stream_executor/stream.h index e8885e1eb6..63d64947c8 100644 --- a/tensorflow/stream_executor/stream.h +++ b/tensorflow/stream_executor/stream.h @@ -629,19 +629,22 @@ class Stream { const dnn::BatchDescriptor &input_dimensions, const DeviceMemory<double> &input_data, const dnn::BatchDescriptor &output_dimensions, - DeviceMemory<double> *output_data); + DeviceMemory<double> *output_data, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenPoolForward(const dnn::PoolingDescriptor &pooling_dimensions, const dnn::BatchDescriptor &input_dimensions, const DeviceMemory<float> &input_data, const dnn::BatchDescriptor &output_dimensions, - DeviceMemory<float> *output_data); + DeviceMemory<float> *output_data, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenPoolForward(const dnn::PoolingDescriptor &pooling_dimensions, const dnn::BatchDescriptor &input_dimensions, const DeviceMemory<Eigen::half> &input_data, const dnn::BatchDescriptor &output_dimensions, - DeviceMemory<Eigen::half> *output_data); + DeviceMemory<Eigen::half> *output_data, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenPoolBackward(const dnn::PoolingDescriptor &pooling_dimensions, const dnn::BatchDescriptor &input_dimensions, @@ -649,7 +652,8 @@ class Stream { const dnn::BatchDescriptor &output_dimensions, const DeviceMemory<double> &output_data, const DeviceMemory<double> &input_diff_data, - DeviceMemory<double> *output_diff_data); + DeviceMemory<double> *output_diff_data, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenPoolBackward(const dnn::PoolingDescriptor &pooling_dimensions, const dnn::BatchDescriptor &input_dimensions, @@ -657,7 +661,8 @@ class Stream { const dnn::BatchDescriptor &output_dimensions, const DeviceMemory<float> &output_data, const DeviceMemory<float> &input_diff_data, - DeviceMemory<float> *output_diff_data); + DeviceMemory<float> *output_diff_data, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenPoolBackward(const dnn::PoolingDescriptor &pooling_dimensions, const dnn::BatchDescriptor &input_dimensions, @@ -665,7 +670,8 @@ class Stream { const dnn::BatchDescriptor &output_dimensions, const DeviceMemory<Eigen::half> &output_data, const DeviceMemory<Eigen::half> &input_diff_data, - DeviceMemory<Eigen::half> *output_diff_data); + DeviceMemory<Eigen::half> *output_diff_data, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenNormalize(const dnn::NormalizeDescriptor &normalize_descriptor, const DeviceMemory<float> &input_data, @@ -684,7 +690,8 @@ class Stream { const DeviceMemory<float> &raw_data, const DeviceMemory<float> &normalized_data, const DeviceMemory<float> &normalized_variable_gradient, - DeviceMemory<float> *raw_variable_gradient); + DeviceMemory<float> *raw_variable_gradient, + ScratchAllocator *workspace_allocator = nullptr); Stream &ThenActivate(dnn::ActivationMode activation_mode, const dnn::BatchDescriptor &dimensions, |