Home
last modified time | relevance | path

Searched refs:scratch_allocator (Results 1 – 25 of 39) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_runner.cc77 se::ScratchAllocator* scratch_allocator, in RunGpuConvForward() argument
91 params.config.output_descriptor, &output_buf, scratch_allocator, in RunGpuConvForward()
97 se::ScratchAllocator* scratch_allocator, in RunGpuConvForwardActivation() argument
135 scratch_allocator, algorithm, options.profile_result); in RunGpuConvForwardActivation()
150 se::ScratchAllocator* scratch_allocator, in RunGpuConvInternalImpl() argument
158 return RunGpuConvForward(params, scratch_allocator, stream, options, in RunGpuConvInternalImpl()
169 params.config.input_descriptor, &input_buf, scratch_allocator, in RunGpuConvInternalImpl()
181 params.config.filter_descriptor, &filter_buf, scratch_allocator, in RunGpuConvInternalImpl()
186 params, scratch_allocator, stream, options, input_buf, filter_buf, in RunGpuConvInternalImpl()
198 se::ScratchAllocator* scratch_allocator, in RunGpuConvInternalImpl() argument
[all …]
Dgpu_conv_algorithm_picker.cc128 ScratchAllocator* scratch_allocator, se::Stream* stream) { in GetMIOpenAlgorithms() argument
146 params.config.conv_desc, scratch_allocator, &algorithms); in GetMIOpenAlgorithms()
463 se::RedzoneAllocator scratch_allocator( in PickBestAlgorithmNoCacheCuda() local
475 &scratch_allocator, stream, options); in PickBestAlgorithmNoCacheCuda()
491 scratch_allocator.TotalAllocatedBytesExcludingRedzones(); in PickBestAlgorithmNoCacheCuda()
507 CheckRedzones(scratch_allocator, stream, "scratch", instr, &result)); in PickBestAlgorithmNoCacheCuda()
689 ScratchAllocator scratch_allocator(device_ordinal, allocator); in PickBestAlgorithmNoCacheRocm() local
694 stream_exec_, &scratch_allocator, stream)); in PickBestAlgorithmNoCacheRocm()
729 &scratch_allocator, stream, options); in PickBestAlgorithmNoCacheRocm()
744 int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes(); in PickBestAlgorithmNoCacheRocm()
Dfft_thunk.cc126 FftScratchAllocator scratch_allocator(buffer_allocations.device_ordinal(), in ExecuteOnStream() local
173 kInPlaceFft, batch_size, &scratch_allocator); in ExecuteOnStream()
177 &stream, fft_plan.get(), &scratch_allocator); in ExecuteOnStream()
Dgpu_conv_runner.h119 se::ScratchAllocator* scratch_allocator, se::Stream* stream,
/external/tensorflow/tensorflow/stream_executor/cuda/
Dcuda_fft.cc81 fft::Type type, int batch_count, ScratchAllocator *scratch_allocator) { in Initialize() argument
86 scratch_allocator_ = scratch_allocator; in Initialize()
102 if (scratch_allocator == nullptr) { in Initialize()
190 return UpdateScratchAllocator(stream, scratch_allocator); in Initialize()
194 if (scratch_allocator == nullptr) { in Initialize()
229 return UpdateScratchAllocator(stream, scratch_allocator); in Initialize()
238 ScratchAllocator *scratch_allocator) { in Initialize() argument
243 /*output_distance=*/0, type, 1, scratch_allocator); in Initialize()
247 Stream *stream, ScratchAllocator *scratch_allocator) { in UpdateScratchAllocator() argument
248 scratch_allocator_ = scratch_allocator; in UpdateScratchAllocator()
[all …]
Dcuda_fft.h73 int batch_count, ScratchAllocator* scratch_allocator);
78 ScratchAllocator* scratch_allocator);
81 ScratchAllocator *scratch_allocator);
Dcuda_dnn.cc2540 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionForwardWorkspace() argument
2569 if (TF_PREDICT_FALSE(!scratch_allocator)) { in AllocateCudnnConvolutionForwardWorkspace()
2574 return scratch_allocator->AllocateBytes(size_in_bytes); in AllocateCudnnConvolutionForwardWorkspace()
2584 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionBackwardDataWorkspace() argument
2615 if (TF_PREDICT_FALSE(!scratch_allocator)) { in AllocateCudnnConvolutionBackwardDataWorkspace()
2620 return scratch_allocator->AllocateBytes(size_in_bytes); in AllocateCudnnConvolutionBackwardDataWorkspace()
2630 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionBackwardFilterWorkspace() argument
2661 if (TF_PREDICT_FALSE(!scratch_allocator)) { in AllocateCudnnConvolutionBackwardFilterWorkspace()
2666 return scratch_allocator->AllocateBytes(size_in_bytes); in AllocateCudnnConvolutionBackwardFilterWorkspace()
2693 const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, in GetCudnnConvolutionForwardAlgorithm() argument
[all …]
Dcuda_blas.h113 int batch_count, ScratchAllocator *scratch_allocator);
149 ScratchAllocator *scratch_allocator,
Dcuda_dnn.h290 DeviceMemory<double>* output_data, ScratchAllocator* scratch_allocator,
304 DeviceMemory<float>* output_data, ScratchAllocator* scratch_allocator,
320 ScratchAllocator* scratch_allocator,
334 DeviceMemory<int8>* output_data, ScratchAllocator* scratch_allocator,
348 DeviceMemory<float>* output_data, ScratchAllocator* scratch_allocator,
628 ScratchAllocator* scratch_allocator,
703 ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc,
713 ScratchAllocator* scratch_allocator, DeviceMemory<uint8>* scratch_memory,
/external/tensorflow/tensorflow/stream_executor/
Dfft.h110 ScratchAllocator *scratch_allocator) = 0;
115 ScratchAllocator *scratch_allocator) = 0;
120 bool in_place_fft, ScratchAllocator *scratch_allocator) = 0;
167 ScratchAllocator *scratch_allocator) = 0;
176 Stream *stream, Plan *plan, ScratchAllocator *scratch_allocator) = 0;
225 ScratchAllocator *scratch_allocator) override; \
228 bool in_place_fft, ScratchAllocator *scratch_allocator) override; \
231 fft::Type type, bool in_place_fft, ScratchAllocator *scratch_allocator) \
242 bool in_place_fft, int batch_count, ScratchAllocator *scratch_allocator) \
245 ScratchAllocator *scratch_allocator) \
Dstream.h333 DeviceMemory<OutputType> *output, ScratchAllocator *scratch_allocator, in ConvolveWithAlgorithm() argument
342 convolution_descriptor, algorithm_config, scratch_allocator, in ConvolveWithAlgorithm()
364 DeviceMemory<double> *output, ScratchAllocator *scratch_allocator,
378 DeviceMemory<float> *output, ScratchAllocator *scratch_allocator,
393 DeviceMemory<Eigen::half> *output, ScratchAllocator *scratch_allocator,
407 ScratchAllocator *scratch_allocator,
421 DeviceMemory<float> *output, ScratchAllocator *scratch_allocator,
444 ScratchAllocator *scratch_allocator, in ConvolveBackwardDataWithAlgorithm() argument
454 algorithm_config, scratch_allocator, &algorithm_desc, in ConvolveBackwardDataWithAlgorithm()
476 ScratchAllocator *scratch_allocator, in ConvolveBackwardFilterWithAlgorithm() argument
[all …]
/external/tensorflow/tensorflow/stream_executor/rocm/
Drocm_fft.cc161 fft::Type type, int batch_count, ScratchAllocator *scratch_allocator) { in Initialize() argument
180 if (scratch_allocator == nullptr) { in Initialize()
271 return UpdateScratchAllocator(stream, scratch_allocator); in Initialize()
275 if (scratch_allocator == nullptr) { in Initialize()
312 return UpdateScratchAllocator(stream, scratch_allocator); in Initialize()
321 ScratchAllocator *scratch_allocator) { in Initialize() argument
326 /*output_distance=*/0, type, 1, scratch_allocator); in Initialize()
330 Stream *stream, ScratchAllocator *scratch_allocator) { in UpdateScratchAllocator() argument
332 auto allocated = scratch_allocator->AllocateBytes(scratch_size_bytes_); in UpdateScratchAllocator()
389 ScratchAllocator *scratch_allocator) { in Create1dPlanWithScratchAllocator() argument
[all …]
Drocm_fft.h72 int batch_count, ScratchAllocator *scratch_allocator);
77 ScratchAllocator *scratch_allocator);
80 ScratchAllocator *scratch_allocator);
Drocm_blas.h130 ScratchAllocator *scratch_allocator, Stream *stream,
160 int ldc, int batch_count, ScratchAllocator *scratch_allocator);
Drocm_dnn.h242 ScratchAllocator* scratch_allocator,
328 DeviceMemory<double>* output_data, ScratchAllocator* scratch_allocator,
342 DeviceMemory<float>* output_data, ScratchAllocator* scratch_allocator,
358 ScratchAllocator* scratch_allocator,
372 DeviceMemory<int8>* output_data, ScratchAllocator* scratch_allocator,
842 ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc,
861 ScratchAllocator* scratch_allocator, DeviceMemory<uint8>* scratch_memory,
872 ScratchAllocator* scratch_allocator,
883 ScratchAllocator* scratch_allocator,
DBUILD184 "//tensorflow/stream_executor:scratch_allocator",
217 "//tensorflow/stream_executor:scratch_allocator",
260 "//tensorflow/stream_executor:scratch_allocator",
/external/tensorflow/tensorflow/core/kernels/
Dconv_grad_input_ops.cc359 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, ctx); in operator ()() local
410 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in operator ()() local
418 : static_cast<se::ScratchAllocator*>(&scratch_allocator); in operator ()()
433 : scratch_allocator.TotalByteSize()); in operator ()()
442 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, ctx); in operator ()()
450 &scratch_allocator, &algorithms), in operator ()()
475 input_desc, &in_backprop_ptr, &scratch_allocator, in operator ()()
485 result.set_scratch_bytes(scratch_allocator.TotalByteSize()); in operator ()()
502 input_desc, &in_backprop_ptr, &scratch_allocator, algorithm_config, in operator ()()
Dconv_ops_3d.cc531 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
538 : static_cast<se::ScratchAllocator*>(&scratch_allocator); in launch()
555 : scratch_allocator.TotalByteSize()); in launch()
564 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch()
572 output_ptr, conv_desc, &scratch_allocator, &algorithms), in launch()
596 output_desc, &output_ptr, &scratch_allocator, in launch()
607 result.set_scratch_bytes(scratch_allocator.TotalByteSize()); in launch()
624 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
627 &output_ptr, &scratch_allocator, algorithm_config, nullptr); in launch()
Dcudnn_pooling_gpu.cc107 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
111 &output_data, &scratch_allocator) in Compute()
249 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
254 &input_backprop_data, &scratch_allocator) in Compute()
Dconv_grad_ops_3d.cc1444 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in Compute() local
1452 : static_cast<se::ScratchAllocator*>(&scratch_allocator); in Compute()
1469 : scratch_allocator.TotalByteSize()); in Compute()
1480 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in Compute()
1487 &scratch_allocator, &algorithms)); in Compute()
1494 input_desc, &in_backprop_ptr, &scratch_allocator, in Compute()
1504 result.set_scratch_bytes(scratch_allocator.TotalByteSize()); in Compute()
1521 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in Compute() local
1525 input_desc, &in_backprop_ptr, &scratch_allocator, algorithm_config, in Compute()
1935 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, in Compute() local
[all …]
Dconv_ops_fused_impl.h395 DnnScratchAllocator scratch_allocator(ConvolveScratchSize(), context);
402 : static_cast<se::ScratchAllocator*>(&scratch_allocator);
418 : scratch_allocator.TotalByteSize());
670 se::ScratchAllocator* scratch_allocator,
682 scratch_allocator, algorithm_config, profile_result);
700 DnnScratchAllocator scratch_allocator(ConvolveScratchSize(), context);
701 Status cudnn_launch_status = launch(algorithm_config, &scratch_allocator,
Dconv_grad_filter_ops.cc1004 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, in operator ()() local
1012 : static_cast<se::ScratchAllocator*>(&scratch_allocator); in operator ()()
1029 : scratch_allocator.TotalByteSize()); in operator ()()
1038 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, in operator ()()
1048 conv_desc, &scratch_allocator, &algorithms), in operator ()()
1073 filter_desc, &filter_backprop_ptr, &scratch_allocator, in operator ()()
1083 result.set_scratch_bytes(scratch_allocator.TotalByteSize()); in operator ()()
1099 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, ctx); in operator ()() local
1102 filter_desc, &filter_backprop_ptr, &scratch_allocator, algorithm_config, in operator ()()
Dconv_ops.cc1020 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in operator ()() local
1024 : static_cast<se::ScratchAllocator*>(&scratch_allocator); in operator ()()
1041 : scratch_allocator.TotalByteSize()); in operator ()()
1051 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in operator ()()
1059 output_ptr, conv_desc, &scratch_allocator, &algorithms), in operator ()()
1085 output_desc, &output_ptr, &scratch_allocator, in operator ()()
1095 result.set_scratch_bytes(scratch_allocator.TotalByteSize()); in operator ()()
1115 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in operator ()() local
1118 &output_ptr, &scratch_allocator, algorithm_config, nullptr); in operator ()()
Dpooling_ops_common.cc395 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
399 &output_data, &scratch_allocator) in Compute()
713 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
718 &input_backprop_data, &scratch_allocator) in Compute()
/external/tensorflow/tensorflow/core/platform/default/build_config/
DBUILD41 "//tensorflow/stream_executor:scratch_allocator",
94 "//tensorflow/stream_executor:scratch_allocator",

12