Home
last modified time | relevance | path

Searched defs:workspace_allocator (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/
H A Ddnn.h1175 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 Dstream.cc340 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 Dcuda_dnn.cc1859 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 Drocm_dnn.cc2105 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 Dcudnn_rnn_ops.cc800 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 Dctc_loss_op.cc356 DnnScratchAllocator workspace_allocator(1LL << 32, ctx); in Compute() local
H A Dfused_batch_norm_op.cc935 workspace_allocator; in operator ()() local
1130 workspace_allocator; in operator ()() local