Searched defs:workspace_allocator (Results 1 – 7 of 7) sorted by relevance
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/ |
H A D | dnn.h | 1175 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() 1195 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() 1227 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationBackward() 1246 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationBackward() 1672 ScratchAllocator* workspace_allocator) { in DoNormalizeBackwardWithDimensions() 2130 ScratchAllocator* workspace_allocator, in DoRnnForward() 2152 ScratchAllocator* workspace_allocator, in DoRnnForward() 2174 ScratchAllocator* workspace_allocator, in DoRnnForward() 2219 virtual bool DoRnnBackward( in DoRnnBackward() 2248 virtual bool DoRnnBackward( in DoRnnBackward() [all …]
|
H A D | stream.cc | 340 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationForward() 365 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationBackward() 394 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationForward() 420 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationBackward() 631 ScratchAllocator *workspace_allocator) { in ThenNormalizeBackwardWithDimensions() 2047 ScratchAllocator *workspace_allocator, in ThenRnnForward() 2082 ScratchAllocator *workspace_allocator, in ThenRnnForward() 2118 ScratchAllocator *workspace_allocator, in ThenRnnForward() 2137 Stream &Stream::ThenRnnBackward( in ThenRnnBackward() 2206 ScratchAllocator *workspace_allocator, in ThenRnnBackward() [all …]
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/cuda/ |
H A D | cuda_dnn.cc | 1859 ScratchAllocator* workspace_allocator) { in CreateRnnWorkspace() 1880 ScratchAllocator* workspace_allocator) { in CreateBatchNormForwardWorkspace() 1904 ScratchAllocator* workspace_allocator) { in CreateBatchNormBackwardWorkspace() 1945 ScratchAllocator* workspace_allocator, in DoRnnForwardImpl() 2165 ScratchAllocator* workspace_allocator, in DoRnnBackwardImpl() 2486 ScratchAllocator* workspace_allocator, in DoRnnForward() 2529 ScratchAllocator* workspace_allocator, in DoRnnForward() 2573 ScratchAllocator* workspace_allocator, in DoRnnForward() 2600 bool CudnnSupport::DoRnnBackward( in DoRnnBackward() 2676 ScratchAllocator* workspace_allocator, in DoRnnBackward() [all …]
|
/aosp_15_r20/external/tensorflow/tensorflow/stream_executor/rocm/ |
H A D | rocm_dnn.cc | 2105 ScratchAllocator* workspace_allocator, in CreateRnnWorkspace() 2151 ScratchAllocator* workspace_allocator) { in DoRnnForwardImpl() 2273 ScratchAllocator* workspace_allocator) { in DoRnnBackwardImpl() 2648 ScratchAllocator* workspace_allocator, in DoRnnForward() 2691 ScratchAllocator* workspace_allocator, in DoRnnForward() 2735 ScratchAllocator* workspace_allocator, in DoRnnForward() 2741 bool MIOpenSupport::DoRnnBackward( in DoRnnBackward() 2817 ScratchAllocator* workspace_allocator, in DoRnnBackward() 2846 bool MIOpenSupport::DoRnnBackward( in DoRnnBackward() 3649 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() [all …]
|
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
H A D | cudnn_rnn_ops.cc | 800 ScratchAllocator* workspace_allocator, in DoForward() 883 ScratchAllocator* workspace_allocator, in DoBackward() 1537 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in ComputeAndReturnAlgorithm() local 1765 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in MaybeAutotune() local 1951 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in ComputeImpl() local
|
H A D | ctc_loss_op.cc | 356 DnnScratchAllocator workspace_allocator(1LL << 32, ctx); in Compute() local
|
H A D | fused_batch_norm_op.cc | 935 workspace_allocator; in operator ()() local 1130 workspace_allocator; in operator ()() local
|