diff options
Diffstat (limited to 'tensorflow/stream_executor/cuda/cuda_dnn.h')
| -rw-r--r-- | tensorflow/stream_executor/cuda/cuda_dnn.h | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.h b/tensorflow/stream_executor/cuda/cuda_dnn.h index cfc7e29574..b280b73c70 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.h +++ b/tensorflow/stream_executor/cuda/cuda_dnn.h @@ -308,6 +308,13 @@ class CudnnSupport : public dnn::DnnSupport { bool DoPoolForward(Stream* stream, const dnn::PoolingDescriptor& pooling_dimensions, const dnn::BatchDescriptor& input_dimensions, + const DeviceMemory<double>& input_data, + const dnn::BatchDescriptor& output_dimensions, + DeviceMemory<double>* output_data) override; + + bool DoPoolForward(Stream* stream, + const dnn::PoolingDescriptor& pooling_dimensions, + const dnn::BatchDescriptor& input_dimensions, const DeviceMemory<float>& input_data, const dnn::BatchDescriptor& output_dimensions, DeviceMemory<float>* output_data) override; @@ -322,6 +329,15 @@ class CudnnSupport : public dnn::DnnSupport { bool DoPoolBackward(Stream* stream, const dnn::PoolingDescriptor& pooling_dimensions, const dnn::BatchDescriptor& input_dimensions, + const DeviceMemory<double>& input_data, + const dnn::BatchDescriptor& output_dimensions, + const DeviceMemory<double>& output_data, + const DeviceMemory<double>& input_diff_data, + DeviceMemory<double>* output_diff_data) override; + + bool DoPoolBackward(Stream* stream, + const dnn::PoolingDescriptor& pooling_dimensions, + const dnn::BatchDescriptor& input_dimensions, const DeviceMemory<float>& input_data, const dnn::BatchDescriptor& output_dimensions, const DeviceMemory<float>& output_data, |