Home
last modified time | relevance | path

Searched refs:workspace_allocator (Results 1 – 10 of 10) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/cuda/
H A Dcuda_dnn.h97 ScratchAllocator* workspace_allocator,
116 ScratchAllocator* workspace_allocator,
135 ScratchAllocator* workspace_allocator,
161 ScratchAllocator* workspace_allocator,
187 ScratchAllocator* workspace_allocator,
213 ScratchAllocator* workspace_allocator,
291 ScratchAllocator* workspace_allocator) override;
306 ScratchAllocator* workspace_allocator) override;
319 ScratchAllocator* workspace_allocator) override;
333 ScratchAllocator* workspace_allocator) override;
[all …]
H A Dcuda_dnn.cc1859 ScratchAllocator* workspace_allocator) { in CreateRnnWorkspace() argument
1870 return workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateRnnWorkspace()
1880 ScratchAllocator* workspace_allocator) { in CreateBatchNormForwardWorkspace() argument
1895 return workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateBatchNormForwardWorkspace()
1904 ScratchAllocator* workspace_allocator) { in CreateBatchNormBackwardWorkspace() argument
1921 return workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateBatchNormBackwardWorkspace()
1945 ScratchAllocator* workspace_allocator, in DoRnnForwardImpl() argument
1981 TF_ASSIGN_OR_RETURN(workspace, workspace_allocator->AllocateBytes( in DoRnnForwardImpl()
2032 workspace_allocator)); in DoRnnForwardImpl()
2165 ScratchAllocator* workspace_allocator, in DoRnnBackwardImpl() argument
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/stream_executor/rocm/
H A Drocm_dnn.h117 ScratchAllocator* workspace_allocator,
136 ScratchAllocator* workspace_allocator,
155 ScratchAllocator* workspace_allocator,
181 ScratchAllocator* workspace_allocator,
207 ScratchAllocator* workspace_allocator,
233 ScratchAllocator* workspace_allocator,
295 ScratchAllocator* workspace_allocator) override;
310 ScratchAllocator* workspace_allocator) override;
323 ScratchAllocator* workspace_allocator) override;
337 ScratchAllocator* workspace_allocator) override;
[all …]
H A Drocm_dnn.cc2105 ScratchAllocator* workspace_allocator, in CreateRnnWorkspace() argument
2120 workspace_allocator->AllocateBytes(workspace_size_in_bytes); in CreateRnnWorkspace()
2151 ScratchAllocator* workspace_allocator) { in DoRnnForwardImpl() argument
2176 workspace_allocator, &workspace)) { in DoRnnForwardImpl()
2273 ScratchAllocator* workspace_allocator) { in DoRnnBackwardImpl() argument
2297 workspace_allocator, &workspace)) { in DoRnnBackwardImpl()
2648 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
2672 reserve_space_allocator, workspace_allocator); in DoRnnForward()
2691 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
2715 reserve_space_allocator, workspace_allocator); in DoRnnForward()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/
H A Ddnn.h1175 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() argument
1195 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationForward() argument
1227 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationBackward() argument
1246 ScratchAllocator* workspace_allocator) { in DoBatchNormalizationBackward() argument
1629 DeviceMemoryBase output_data, ScratchAllocator* workspace_allocator) = 0;
1639 ScratchAllocator* workspace_allocator) = 0;
1672 ScratchAllocator* workspace_allocator) { in DoNormalizeBackwardWithDimensions() argument
2130 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
2152 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
2174 ScratchAllocator* workspace_allocator, in DoRnnForward() argument
[all …]
H A Dstream.h264 ScratchAllocator *workspace_allocator);
276 ScratchAllocator *workspace_allocator);
291 ScratchAllocator *workspace_allocator);
305 ScratchAllocator *workspace_allocator);
474 ScratchAllocator *workspace_allocator = nullptr) {
479 workspace_allocator);
493 ScratchAllocator *workspace_allocator = nullptr) {
498 input_diff_data, *output_diff_data, workspace_allocator);
515 ScratchAllocator *workspace_allocator = nullptr);
1243 ScratchAllocator *workspace_allocator,
[all …]
H A Dstream.cc340 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationForward() argument
348 is_training, reserve_space_allocator, workspace_allocator)); in ThenBatchNormalizationForward()
365 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationBackward() argument
374 workspace_allocator)); in ThenBatchNormalizationBackward()
394 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationForward() argument
402 is_training, reserve_space_allocator, workspace_allocator)); in ThenBatchNormalizationForward()
420 ScratchAllocator *workspace_allocator) { in ThenBatchNormalizationBackward() argument
429 workspace_allocator)); in ThenBatchNormalizationBackward()
631 ScratchAllocator *workspace_allocator) { in ThenNormalizeBackwardWithDimensions() argument
634 PARAM(raw_variable_gradient), PARAM(workspace_allocator)); in ThenNormalizeBackwardWithDimensions()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dcudnn_rnn_ops.cc800 ScratchAllocator* workspace_allocator, in DoForward() argument
858 reserve_space_allocator, workspace_allocator, in DoForward()
883 ScratchAllocator* workspace_allocator, in DoBackward() argument
960 &reserve_space_uint8, workspace_allocator, output_profile_result) in DoBackward()
1537 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in ComputeAndReturnAlgorithm() local
1562 &workspace_allocator, /*output_profile_result=*/nullptr); in ComputeAndReturnAlgorithm()
1765 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in MaybeAutotune() local
1769 &reserve_space_allocator, &workspace_allocator, in MaybeAutotune()
1783 &params_backprop, nullptr, true, &workspace_allocator, in MaybeAutotune()
1951 CudnnRnnAllocatorInTemp<uint8> workspace_allocator(context); in ComputeImpl() local
[all …]
H A Dctc_loss_op.cc356 DnnScratchAllocator workspace_allocator(1LL << 32, ctx); in Compute() local
363 *grads_desc, &grads_data, &workspace_allocator) in Compute()
H A Dfused_batch_norm_op.cc935 workspace_allocator; in operator ()() local
939 workspace_allocator.reset( in operator ()()
974 workspace_allocator.get()) in operator ()()
1130 workspace_allocator; in operator ()() local
1136 workspace_allocator.reset( in operator ()()
1158 workspace_allocator.get()) in operator ()()