Home
last modified time | relevance | path

Searched defs:scratch_allocator (Results 1 – 25 of 29) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/cuda/
H A Dcuda_fft.cc100 fft::Type type, int batch_count, ScratchAllocator *scratch_allocator) { in Initialize()
259 ScratchAllocator *scratch_allocator) { in Initialize()
268 Stream *stream, ScratchAllocator *scratch_allocator) { in UpdateScratchAllocator()
332 ScratchAllocator *scratch_allocator) { in Create1dPlanWithScratchAllocator()
365 bool in_place_fft, ScratchAllocator *scratch_allocator) { in Create2dPlanWithScratchAllocator()
400 bool in_place_fft, ScratchAllocator *scratch_allocator) { in Create3dPlanWithScratchAllocator()
448 bool in_place_fft, int batch_count, ScratchAllocator *scratch_allocator) { in CreateBatchedPlanWithScratchAllocator()
473 Stream *stream, fft::Plan *plan, ScratchAllocator *scratch_allocator) { in UpdatePlanWithScratchAllocator()
H A Dcuda_blas.cc1193 ScratchAllocator *scratch_allocator) { in DoBlasInternalImpl()
1330 ScratchAllocator *scratch_allocator) { in DoBlasInternalImpl()
1347 ScratchAllocator *scratch_allocator) { in DoBlasInternalImpl()
1363 ScratchAllocator *scratch_allocator) { in DoBlasInternalImpl()
1380 int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasInternalImpl()
1397 int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasInternalImpl()
H A Dcuda_dnn.cc2892 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionForwardWorkspace()
2941 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionBackwardDataWorkspace()
2991 ScratchAllocator* scratch_allocator) { in AllocateCudnnConvolutionBackwardFilterWorkspace()
3058 const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, in GetCudnnConvolutionForwardAlgorithm()
3121 const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, in GetCudnnConvolutionBackwardDataAlgorithm()
3182 const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, in GetCudnnConvolutionBackwardFilterAlgorithm()
3964 ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc, in DoPrepareForConvolution()
5765 ScratchAllocator* scratch_allocator, in DoFusedConvolve()
5843 ScratchAllocator* scratch_allocator, DeviceMemory<uint8>* scratch_memory, in DoPrepareForCtcLoss()
H A Dcuda_blas_lt.cc262 ScratchAllocator& scratch_allocator, in DoMatmul()
/aosp_15_r20/external/tensorflow/tensorflow/stream_executor/rocm/
H A Drocm_fft.cc161 fft::Type type, int batch_count, ScratchAllocator *scratch_allocator) { in Initialize()
322 ScratchAllocator *scratch_allocator) { in Initialize()
331 Stream *stream, ScratchAllocator *scratch_allocator) { in UpdateScratchAllocator()
391 ScratchAllocator *scratch_allocator) { in Create1dPlanWithScratchAllocator()
420 bool in_place_fft, ScratchAllocator *scratch_allocator) { in Create2dPlanWithScratchAllocator()
450 bool in_place_fft, ScratchAllocator *scratch_allocator) { in Create3dPlanWithScratchAllocator()
485 bool in_place_fft, int batch_count, ScratchAllocator *scratch_allocator) { in CreateBatchedPlanWithScratchAllocator()
500 Stream *stream, fft::Plan *plan, ScratchAllocator *scratch_allocator) { in UpdatePlanWithScratchAllocator()
H A Drocm_blas.cc712 int batch_count, uint64_t batch_stride, ScratchAllocator *scratch_allocator, in AllocateStridedBuffer()
771 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasGemmBatchedInternal()
865 int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasGemmBatched()
887 int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasGemmBatched()
905 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasGemmBatched()
925 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasGemmBatched()
945 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in DoBlasGemmBatched()
H A Drocm_dnn.cc2451 ScratchAllocator* scratch_allocator, DeviceMemory<uint8>* scratch_memory, in DoPrepareForCtcLoss()
2911 ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc, in DoPrepareForConvolution()
3154 ScratchAllocator* scratch_allocator, in GetConvolveRunners()
3219 ScratchAllocator* scratch_allocator, in GetMIOpenConvolveAlgorithms()
3241 ScratchAllocator* scratch_allocator, in GetMIOpenConvolveAlgorithmsImmediateMode()
3450 ScratchAllocator* scratch_allocator, in GetMIOpenConvolveAlgorithmsFindMode()
3807 ScratchAllocator* scratch_allocator, in DoFusedConvolve()
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dmatmul_op_fused.cc204 BlasScratchAllocator scratch_allocator(context); in AutotuneMatmul() local
304 se::blas::ProfileResult* profile_result) { in operator ()()
318 BlasScratchAllocator scratch_allocator(context); in operator ()() local
H A Dconv_ops_gpu.h133 AllocateScratchOrFallback(se::ScratchAllocator* scratch_allocator, in AllocateScratchOrFallback()
162 DnnScratchAllocator* scratch_allocator, in LaunchAutotunedConv()
H A Dcudnn_pooling_gpu.cc108 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
244 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
H A Dconv_ops_gpu.cc55 DnnScratchAllocator scratch_allocator(scratch_size_limit, ctx); in AutotuneConvImpl() local
384 DnnScratchAllocator scratch_allocator(scratch_size_limit, ctx); in AutotuneUnfusedConv() local
H A Dpooling_ops_common.cc398 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
722 DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); in Compute() local
H A Dconv_grad_input_ops.cc386 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, ctx); in operator ()() local
H A Dconv_ops_3d.cc513 DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); in launch() local
H A Dconv_grad_ops_3d.cc1529 DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, in Compute() local
1928 DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, in Compute() local
H A Dfft_ops.cc479 CufftScratchAllocator scratch_allocator(CufftScratchSize, ctx); in DoFFT() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_conv_algorithm_picker.cc190 ScratchAllocator* scratch_allocator, se::Stream* stream) { in GetMIOpenAlgorithms()
489 se::RedzoneAllocator scratch_allocator( in AutotuneOneConvRunner() local
828 ScratchAllocator scratch_allocator(device_ordinal, allocator); in PickBestAlgorithmNoCacheRocm() local
H A Dmatmul_utils.cc693 se::ScratchAllocator& scratch_allocator, in DoMatmul()
722 se::ScratchAllocator& scratch_allocator, in ExecuteOnStream()
H A Dcublas_lt_matmul_thunk.cc61 se::OwningScratchAllocator<> scratch_allocator(allocs.device_ordinal(), in ExecuteOnStream() local
H A Dfft_thunk.cc104 se::OwningScratchAllocator<2> scratch_allocator(device_ordinal, in RunFft() local
H A Dgemm_algorithm_picker.cc296 se::OwningScratchAllocator<> scratch_allocator( in DoGemmAutotune() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/
H A Dstream.h344 ScratchAllocator *scratch_allocator, in ConvolveWithAlgorithm()
377 DeviceMemory<OutputT> *output, ScratchAllocator *scratch_allocator, in FusedConvolveWithAlgorithm()
H A Ddnn.h1312 DeviceMemoryBase output_data, ScratchAllocator* scratch_allocator, in DoFusedConvolve()
1330 ScratchAllocator* scratch_allocator, AlgorithmDesc* algorithm_desc, in PrepareForConvolution()
2626 ScratchAllocator* scratch_allocator, AlgorithmDesc* algorithm_desc, in DoPrepareForConvolution()
2640 ScratchAllocator* scratch_allocator, DeviceMemory<uint8>* scratch_memory, in DoPrepareForCtcLoss()
H A Dstream_executor_pimpl.cc301 ScratchAllocator* scratch_allocator, in GetConvolveRunners()
344 ScratchAllocator* scratch_allocator, in GetMIOpenConvolveAlgorithms()
H A Dstream.cc1716 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in ThenBlasGemmBatchedWithScratch()
1755 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in ThenBlasGemmBatchedWithScratch()
1792 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in ThenBlasGemmBatchedWithScratch()
1838 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in ThenBlasGemmBatchedWithScratch()
1874 int ldc, int batch_count, ScratchAllocator *scratch_allocator) { in ThenBlasGemmBatchedWithScratch()

12