| /external/tensorflow/tensorflow/stream_executor/cuda/ |
| D | cuda_dnn.h | 89 ScratchAllocator* workspace_allocator,
107 ScratchAllocator* workspace_allocator,
125 ScratchAllocator* workspace_allocator,
150 ScratchAllocator* workspace_allocator,
175 ScratchAllocator* workspace_allocator,
200 ScratchAllocator* workspace_allocator,
230 ScratchAllocator* workspace_allocator) override;
244 ScratchAllocator* workspace_allocator) override;
255 ScratchAllocator* workspace_allocator) override;
266 ScratchAllocator* workspace_allocator) override;
[all …]
|
| D | cuda_dnn.cc | 1678 ScratchAllocator* workspace_allocator) { in CreateRnnWorkspace() argument
1689 return workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateRnnWorkspace()
1699 ScratchAllocator* workspace_allocator) { in CreateBatchNormForwardWorkspace() argument
1714 return workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateBatchNormForwardWorkspace()
1722 ScratchAllocator* workspace_allocator) { in CreateBatchNormBackwardWorkspace() argument
1738 return workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateBatchNormBackwardWorkspace()
1761 ScratchAllocator* workspace_allocator, in DoRnnForwardImpl() argument
1775 workspace_allocator)) in DoRnnForwardImpl()
1908 ScratchAllocator* workspace_allocator, in DoRnnBackwardImpl() argument
1922 workspace_allocator)); in DoRnnBackwardImpl()
[all …]
|
| /external/tensorflow/tensorflow/stream_executor/rocm/ |
| D | rocm_dnn.h | 116 ScratchAllocator* workspace_allocator,
134 ScratchAllocator* workspace_allocator,
152 ScratchAllocator* workspace_allocator,
177 ScratchAllocator* workspace_allocator,
202 ScratchAllocator* workspace_allocator,
227 ScratchAllocator* workspace_allocator,
268 ScratchAllocator* workspace_allocator) override;
282 ScratchAllocator* workspace_allocator) override;
293 ScratchAllocator* workspace_allocator) override;
304 ScratchAllocator* workspace_allocator) override;
[all …]
|
| D | rocm_dnn.cc | 2051 ScratchAllocator* workspace_allocator, in CreateRnnWorkspace() argument
2066 workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateRnnWorkspace()
2097 ScratchAllocator* workspace_allocator) { in DoRnnForwardImpl() argument
2122 workspace_allocator, &workspace)) { in DoRnnForwardImpl()
2219 ScratchAllocator* workspace_allocator) { in DoRnnBackwardImpl() argument
2243 workspace_allocator, &workspace)) { in DoRnnBackwardImpl()
2593 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
2617 reserve_space_allocator, workspace_allocator); in DoRnnForward()
2635 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
2659 reserve_space_allocator, workspace_allocator); in DoRnnForward()
[all …]
|
| /external/tensorflow/tensorflow/stream_executor/ |
| D | dnn.h | 1024 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() argument
1043 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() argument
1073 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationBackward() argument
1089 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationBackward() argument
1642 ScratchAllocator* workspace_allocator) = 0;
1650 ScratchAllocator* workspace_allocator) { in DoPoolForward() argument
1661 ScratchAllocator* workspace_allocator) { in DoPoolForward() argument
1672 ScratchAllocator* workspace_allocator) { in DoPoolForward() argument
1686 ScratchAllocator* workspace_allocator) { in DoPoolBackward() argument
1699 ScratchAllocator* workspace_allocator) { in DoPoolBackward() argument
[all …]
|
| D | stream.h | 259 ScratchAllocator *workspace_allocator);
269 ScratchAllocator *workspace_allocator);
284 ScratchAllocator *workspace_allocator);
295 ScratchAllocator *workspace_allocator);
545 ScratchAllocator *workspace_allocator = nullptr);
552 ScratchAllocator *workspace_allocator = nullptr);
559 ScratchAllocator *workspace_allocator = nullptr);
566 ScratchAllocator *workspace_allocator = nullptr);
575 ScratchAllocator *workspace_allocator = nullptr);
584 ScratchAllocator *workspace_allocator = nullptr);
[all …]
|
| D | stream.cc | 355 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationForward() argument
363 is_training, reserve_space_allocator, workspace_allocator)); in ThenBatchNormalizationForward()
378 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationBackward() argument
386 reserve_space_data, workspace_allocator)); in ThenBatchNormalizationBackward()
406 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationForward() argument
414 is_training, reserve_space_allocator, workspace_allocator)); in ThenBatchNormalizationForward()
430 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationBackward() argument
438 reserve_space_data, workspace_allocator)); in ThenBatchNormalizationBackward()
824 DeviceMemory<double> *output_data, ScratchAllocator *workspace_allocator) { in ThenPoolForward() argument
827 PARAM(workspace_allocator)); in ThenPoolForward()
[all …]
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | cudnn_rnn_ops.cc | 797 ScratchAllocator* workspace_allocator, in DoForward() argument
836 workspace_allocator, output_profile_result) in DoForward()
860 ScratchAllocator* workspace_allocator, in DoBackward() argument
917 ¶ms_backprop_data, &reserve_space_uint8, workspace_allocator, in DoBackward()
1495 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in ComputeAndReturnAlgorithm() local
1519 &workspace_allocator, /*output_profile_result=*/nullptr); in ComputeAndReturnAlgorithm()
1721 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in MaybeAutoTune() local
1725 &reserve_space_allocator, &workspace_allocator, in MaybeAutoTune()
1739 ¶ms_backprop, nullptr, true, &workspace_allocator, in MaybeAutoTune()
1907 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in ComputeImpl() local
[all …]
|
| D | ctc_loss_op.cc | 347 DnnScratchAllocator workspace_allocator(1LL << 32, ctx); in Compute() local
354 *grads_desc, &grads_data, &workspace_allocator) in Compute()
|
| D | fused_batch_norm_op.cc | 925 workspace_allocator; in operator ()() local
929 workspace_allocator.reset( in operator ()()
964 workspace_allocator.get()) in operator ()()
1107 workspace_allocator; in operator ()() local
1113 workspace_allocator.reset( in operator ()()
1132 reserve_space_data_ptr, workspace_allocator.get()) in operator ()()
|