diff options
| author |  A. Unique TensorFlower <nobody@tensorflow.org> | 2016-01-07 13:25:28 -0800 |
|---|---|---|
| committer |  Vijay Vasudevan <vrv@google.com> | 2016-01-07 13:25:28 -0800 |
| commit | f13442a8c6abadf61671ab70400c6ad287f4dcee (patch) | |
| tree | bf9a2ebdaac144780d435f50455b06f9ee766e44 /tensorflow/core/kernels/conv_ops_gpu.h | |
| parent | 68aab4180381054a6a87fb4b485569912674c54c (diff) | |
The default memory manager would delay 10s for any memory failure. This is not good for Cudnn scratch space allocation, which is only optional with performance gains.
Adding per-allocation attributes so BFC GPU allocator can avoid waiting for
memory failures for Cudnn scratch space allocation.
Change: 111617463
Diffstat (limited to 'tensorflow/core/kernels/conv_ops_gpu.h')
| -rw-r--r-- | tensorflow/core/kernels/conv_ops_gpu.h | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/conv_ops_gpu.h b/tensorflow/core/kernels/conv_ops_gpu.h index bcdc1c3510..8f131f7b81 100644 --- a/tensorflow/core/kernels/conv_ops_gpu.h +++ b/tensorflow/core/kernels/conv_ops_gpu.h @@ -57,8 +57,11 @@ class CudnnScratchAllocator : public perftools::gputools::ScratchAllocator { AllocateBytes(perftools::gputools::Stream* stream, int64 byte_size) override { Tensor temporary_memory; + AllocationAttributes allocation_attr; + allocation_attr.no_retry_on_failure = true; Status allocation_status(context_->allocate_temp( - DT_UINT8, TensorShape({byte_size}), &temporary_memory)); + DT_UINT8, TensorShape({byte_size}), &temporary_memory, + AllocatorAttributes(), allocation_attr)); if (!allocation_status.ok()) { LOG(WARNING) << allocation_status; return perftools::gputools::port::StatusOr< |