Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/stream_executor/cuda/
Dcuda_fft.cc80 fft::Type type, int batch_count, ScratchAllocator *scratch_allocator) { in Initialize() argument
100 if (scratch_allocator == nullptr) { in Initialize()
188 return UpdateScratchAllocator(stream, scratch_allocator); in Initialize()
192 if (scratch_allocator == nullptr) { in Initialize()
227 return UpdateScratchAllocator(stream, scratch_allocator); in Initialize()
236 ScratchAllocator *scratch_allocator) { in Initialize() argument
241 /*output_distance=*/0, type, 1, scratch_allocator); in Initialize()
245 Stream *stream, ScratchAllocator *scratch_allocator) { in UpdateScratchAllocator() argument
248 scratch_allocator->AllocateBytes(stream, scratch_size_bytes_); in UpdateScratchAllocator()
310 ScratchAllocator *scratch_allocator) { in Create1dPlanWithScratchAllocator() argument
[all …]
Dcuda_fft.h72 int batch_count, ScratchAllocator* scratch_allocator);
77 ScratchAllocator* scratch_allocator);
80 ScratchAllocator *scratch_allocator);
Dcuda_dnn.cc2215 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionForwardWorkspace() argument
2243 if (TF_PREDICT_FALSE(!scratch_allocator)) { in AllocateCudnnConvolutionForwardWorkspace()
2248 return scratch_allocator->AllocateBytes(stream, size_in_bytes); in AllocateCudnnConvolutionForwardWorkspace()
2258 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionBackwardDataWorkspace() argument
2288 if (TF_PREDICT_FALSE(!scratch_allocator)) { in AllocateCudnnConvolutionBackwardDataWorkspace()
2293 return scratch_allocator->AllocateBytes(stream, size_in_bytes); in AllocateCudnnConvolutionBackwardDataWorkspace()
2303 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionBackwardFilterWorkspace() argument
2333 if (TF_PREDICT_FALSE(!scratch_allocator)) { in AllocateCudnnConvolutionBackwardFilterWorkspace()
2338 return scratch_allocator->AllocateBytes(stream, size_in_bytes); in AllocateCudnnConvolutionBackwardFilterWorkspace()
2346 const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, in GetCudnnConvolutionForwardAlgorithm() argument
[all …]
Dcuda_dnn.h283 DeviceMemory<double>* output_data, ScratchAllocator* scratch_allocator,
297 DeviceMemory<float>* output_data, ScratchAllocator* scratch_allocator,
315 ScratchAllocator* scratch_allocator,
329 DeviceMemory<int8>* output_data, ScratchAllocator* scratch_allocator,
594 ScratchAllocator* scratch_allocator,
660 ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc,
DBUILD208 "//tensorflow/stream_executor:scratch_allocator",
246 "//tensorflow/stream_executor:scratch_allocator",
291 "//tensorflow/stream_executor:scratch_allocator",
/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.h311 DeviceMemory<double> *output, ScratchAllocator *scratch_allocator,
322 DeviceMemory<float> *output, ScratchAllocator *scratch_allocator,
333 DeviceMemory<Eigen::half> *output, ScratchAllocator *scratch_allocator,
347 DeviceMemory<double> *output, ScratchAllocator *scratch_allocator,
361 DeviceMemory<float> *output, ScratchAllocator *scratch_allocator,
376 DeviceMemory<Eigen::half> *output, ScratchAllocator *scratch_allocator,
390 ScratchAllocator *scratch_allocator,
412 ScratchAllocator *scratch_allocator,
424 ScratchAllocator *scratch_allocator,
436 ScratchAllocator *scratch_allocator,
[all …]
DBUILD447 name = "scratch_allocator",
448 srcs = ["scratch_allocator.cc"],
449 hdrs = ["scratch_allocator.h"],
Dstream.cc450 ScratchAllocator *scratch_allocator, in ThenFusedConvolveWithAlgorithm() argument
466 activation_mode, output_descriptor, output, scratch_allocator, in ThenFusedConvolveWithAlgorithm()
488 ScratchAllocator *scratch_allocator, in ThenFusedConvolveWithAlgorithm() argument
504 activation_mode, output_descriptor, output, scratch_allocator, in ThenFusedConvolveWithAlgorithm()
527 DeviceMemory<Eigen::half> *output, ScratchAllocator *scratch_allocator, in ThenFusedConvolveWithAlgorithm() argument
543 activation_mode, output_descriptor, output, scratch_allocator, in ThenFusedConvolveWithAlgorithm()
565 ScratchAllocator *scratch_allocator, in ThenFusedConvolveWithAlgorithm() argument
581 activation_mode, output_descriptor, output, scratch_allocator, in ThenFusedConvolveWithAlgorithm()
600 ScratchAllocator *scratch_allocator, in ThenConvolveWithAlgorithm() argument
617 scratch_allocator, &algorithm_desc, &scratch_memory) in ThenConvolveWithAlgorithm()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcudnn_conv_runner.cc114 se::ScratchAllocator* scratch_allocator, in RunCudnnConvImpl() argument
248 scratch_allocator, algorithm, profile_result); in RunCudnnConvImpl()
259 scratch_allocator, algorithm, profile_result); in RunCudnnConvImpl()
270 scratch_allocator, algorithm, profile_result); in RunCudnnConvImpl()
303 output_descriptor, &output_buf, scratch_allocator, algorithm, in RunCudnnConvImpl()
399 ScratchBufAllocator scratch_allocator(scratch_buf); in RunCudnnConv() local
400 return RunCudnnConv(conv, operand_buffers, result_buffer, &scratch_allocator, in RunCudnnConv()
407 se::ScratchAllocator* scratch_allocator, se::Stream* stream, in RunCudnnConv() argument
420 return RunCudnnConvImpl<Eigen::half>(params, scratch_allocator, stream, in RunCudnnConv()
423 return RunCudnnConvImpl<float>(params, scratch_allocator, stream, in RunCudnnConv()
[all …]
Dfft_thunk.cc117 FftScratchAllocator scratch_allocator(buffer_allocations.device_ordinal(), in ExecuteOnStream() local
150 fft_type_, kInPlaceFft, batch_size, &scratch_allocator); in ExecuteOnStream()
154 stream, fft_plan_.get(), &scratch_allocator); in ExecuteOnStream()
Dcusolver_rewriter.cc187 ScratchAllocator scratch_allocator(device_ordinal, allocator); in RunOnComputation() local
196 RunOnInstruction(&context, &scratch_allocator, instruction)); in RunOnComputation()
Dcudnn_conv_algorithm_picker.cc236 ScratchAllocator scratch_allocator(device_ordinal, allocator); in PickBestAlgorithm() local
247 &scratch_allocator, &stream, options); in PickBestAlgorithm()
262 int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes(); in PickBestAlgorithm()
Dcudnn_conv_runner.h62 se::ScratchAllocator* scratch_allocator, se::Stream* stream,
/external/tensorflow/tensorflow/stream_executor/rocm/
Drocm_fft.cc149 fft::Type type, int batch_count, ScratchAllocator *scratch_allocator) { in Initialize() argument
168 if (scratch_allocator == nullptr) { in Initialize()
264 scratch_allocator->AllocateBytes(stream, size_in_bytes); in Initialize()
281 if (scratch_allocator == nullptr) { in Initialize()
320 scratch_allocator->AllocateBytes(stream, size_in_bytes); in Initialize()
341 ScratchAllocator *scratch_allocator) { in Initialize() argument
346 /*output_distance=*/0, type, 1, scratch_allocator); in Initialize()
390 ScratchAllocator *scratch_allocator) { in Create1dPlanWithScratchAllocator() argument
394 type, scratch_allocator); in Create1dPlanWithScratchAllocator()
419 bool in_place_fft, ScratchAllocator *scratch_allocator) { in Create2dPlanWithScratchAllocator() argument
[all …]
Drocm_fft.h71 int batch_count, ScratchAllocator *scratch_allocator);
76 ScratchAllocator *scratch_allocator);
DBUILD172 "//tensorflow/stream_executor:scratch_allocator",
199 "//tensorflow/stream_executor:scratch_allocator",
236 "//tensorflow/stream_executor:scratch_allocator",
Drocm_dnn.h274 DeviceMemory<double>* output_data, ScratchAllocator* scratch_allocator,
288 DeviceMemory<float>* output_data, ScratchAllocator* scratch_allocator,
306 ScratchAllocator* scratch_allocator,
320 DeviceMemory<int8>* output_data, ScratchAllocator* scratch_allocator,
767 ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc,
/external/tensorflow/tensorflow/core/kernels/
Dconv_ops_3d.cc454 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
460 output_desc, &output_ptr, &scratch_allocator, in launch()
471 scratch_allocator.TotalByteSize()); in launch()
484 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
489 &output_ptr, &scratch_allocator, in launch()
Dconv_ops_fused_impl.h526 DnnScratchAllocator scratch_allocator(ConvolveScratchSize(), context);
530 launch(se::dnn::AlgorithmConfig(profile_algorithm), &scratch_allocator,
540 scratch_allocator.TotalByteSize());
777 DnnScratchAllocator* scratch_allocator,
789 scratch_allocator, algorithm_config, profile_result)
806 DnnScratchAllocator scratch_allocator(ConvolveScratchSize(), context);
807 bool cudnn_launch_status = launch(algorithm_config, &scratch_allocator,
Dconv_grad_ops_3d.cc1370 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in Compute() local
1377 conv_desc, input_desc, &in_backprop_ptr, &scratch_allocator, in Compute()
1386 if (scratch_allocator.TotalByteSize() == 0 && in Compute()
1407 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in Compute() local
1413 conv_desc, input_desc, &in_backprop_ptr, &scratch_allocator, in Compute()
1776 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, in Compute() local
1784 &scratch_allocator, AlgorithmConfig(profile_algorithm), in Compute()
1793 if (scratch_allocator.TotalByteSize() == 0 && in Compute()
1814 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, in Compute() local
1820 filter_desc, &filter_backprop_ptr, &scratch_allocator, in Compute()
Dbatch_matmul_op_impl.h385 CublasScratchAllocator scratch_allocator(context);
392 batch_size, &scratch_allocator)
477 CublasScratchAllocator scratch_allocator(context);
484 batch_size, &scratch_allocator)
Dconv_ops.cc864 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in operator ()() local
870 output_desc, &output_ptr, &scratch_allocator, in operator ()()
881 scratch_allocator.TotalByteSize()); in operator ()()
894 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in operator ()() local
899 &output_ptr, &scratch_allocator, in operator ()()
/external/tensorflow/tensorflow/contrib/fused_conv/kernels/
Dfused_conv2d_bias_activation_op.cc718 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
726 output_desc, &output_ptr, &scratch_allocator, in launch()
737 scratch_allocator.TotalByteSize()); in launch()
753 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
760 &output_ptr, &scratch_allocator, algorithm_config, in launch()
/external/tensorflow/tensorflow/core/platform/default/build_config/
DBUILD46 "//tensorflow/stream_executor:scratch_allocator",
96 "//tensorflow/stream_executor:scratch_allocator",

12