Home
last modified time | relevance | path

Searched refs:device_allocator (Results 1 – 25 of 29) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dcompiler.h107 se::DeviceMemoryAllocator* device_allocator() const { in device_allocator() function
110 void set_device_allocator(se::DeviceMemoryAllocator* device_allocator) { in set_device_allocator() argument
111 device_allocator_ = device_allocator; in set_device_allocator()
230 se::DeviceMemoryAllocator* device_allocator = nullptr; member
248 se::DeviceMemoryAllocator* device_allocator) { in RunHloPasses() argument
250 CompileOptions{device_allocator}); in RunHloPasses()
275 se::DeviceMemoryAllocator* device_allocator) { in RunBackend() argument
277 CompileOptions{device_allocator}); in RunBackend()
300 se::DeviceMemoryAllocator* device_allocator) { in Compile() argument
302 CompileOptions{device_allocator}); in Compile()
H A Dlocal_service.cc175 {build_options.device_allocator(), in CompileExecutables()
192 Compiler::CompileOptions{build_options.device_allocator(), in CompileExecutables()
221 Compiler::CompileOptions{build_options.device_allocator(), in CompileAotResults()
H A Dllvm_compiler.cc47 options.device_allocator)); in Compile()
50 options.device_allocator)); in Compile()
/aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/pluggable_device/
H A Dpluggable_device_process_state.cc119 Allocator* device_allocator = nullptr; in GetPluggableDeviceAllocator() local
126 device_allocator = new PluggableDeviceBFCAllocator( in GetPluggableDeviceAllocator()
131 device_allocator = new PluggableDeviceSimpleAllocator(sub_allocator); in GetPluggableDeviceAllocator()
134 allocator_parts = {std::unique_ptr<Allocator>(device_allocator), in GetPluggableDeviceAllocator()
135 device_allocator, sub_allocator}; in GetPluggableDeviceAllocator()
H A Dpluggable_device_factory.cc249 Allocator* device_allocator = process_state->GetPluggableDeviceAllocator( in CreatePluggableDevice() local
251 if (device_allocator == nullptr) { in CreatePluggableDevice()
256 absl::optional<AllocatorStats> stats = device_allocator->GetStats(); in CreatePluggableDevice()
269 GetShortDeviceDescription(platform_device_id, *desc), device_allocator, in CreatePluggableDevice()
H A Dpluggable_device.cc143 const std::string& physical_device_desc, Allocator* device_allocator, in PluggableDevice() argument
148 device_allocator_(device_allocator), in PluggableDevice()
H A Dpluggable_device_process_state.h108 Allocator* device_allocator; member
H A Dpluggable_device.h51 Allocator* device_allocator, Allocator* cpu_allocator,
/aosp_15_r20/external/pytorch/c10/cuda/
H A DCUDACachingAllocator.cpp3183 std::vector<std::unique_ptr<DeviceCachingAllocator>> device_allocator; member in c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator
3200 const auto size = static_cast<int64_t>(device_allocator.size()); in init()
3202 device_allocator.resize(device_count); in init()
3204 device_allocator[i] = std::make_unique<DeviceCachingAllocator>(); in init()
3210 return !device_allocator.empty(); in initialized()
3220 0 <= device && static_cast<size_t>(device) < device_allocator.size(), in malloc()
3224 Block* block = device_allocator[device]->malloc(device, size, stream); in malloc()
3247 device_allocator[block->device]->free(block); in free()
3252 0 <= device && static_cast<size_t>(device) < device_allocator.size(), in setMemoryFraction()
3262 device_allocator[device]->setMemoryFraction(fraction); in setMemoryFraction()
[all …]
H A DCUDAMallocAsyncAllocator.cpp895 CudaMallocAsyncAllocator device_allocator; variable
901 return &device_allocator; in allocator()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Damdgpu_compiler.cc78 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloConvolutionCanonicalization() argument
110 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloPostLayoutAssignment() argument
112 hlo_module, stream_exec, device_allocator)); in OptimizeHloPostLayoutAssignment()
H A Dgpu_compiler.h123 se::DeviceMemoryAllocator* device_allocator);
128 se::DeviceMemoryAllocator* device_allocator);
132 se::DeviceMemoryAllocator* device_allocator) = 0;
H A Dnvptx_compiler.cc73 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloConvolutionCanonicalization() argument
121 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloPostLayoutAssignment() argument
147 hlo_module, stream_exec, device_allocator)); in OptimizeHloPostLayoutAssignment()
155 post_pipeline.AddPass<GemmAlgorithmPicker>(stream_exec, device_allocator); in OptimizeHloPostLayoutAssignment()
H A Damdgpu_compiler.h38 se::DeviceMemoryAllocator* device_allocator) override;
42 se::DeviceMemoryAllocator* device_allocator) override;
H A Dnvptx_compiler.h42 se::DeviceMemoryAllocator* device_allocator) override;
46 se::DeviceMemoryAllocator* device_allocator) override;
H A Dgpu_compiler.cc338 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloModule() argument
584 hlo_module, stream_exec, device_allocator)); in OptimizeHloModule()
606 device_allocator)); in OptimizeHloModule()
724 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloPostLayoutAssignment() argument
813 pipeline.AddPass<GpuConvAlgorithmPicker>(stream_exec, device_allocator); in OptimizeHloPostLayoutAssignment()
834 OptimizeHloModule(module.get(), stream_exec, options.device_allocator)); in RunHloPasses()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/
H A Dtpu_on_demand_compiler.cc65 auto allocator = ApiConverter::ToC(options.device_allocator); in RunHloPasses()
94 auto allocator = ApiConverter::ToC(options.device_allocator); in RunBackend()
146 ApiConverter::ToC(options.device_allocator); in Compile()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dxla_platform_info.h35 std::shared_ptr<se::DeviceMemoryAllocator> device_allocator) in XlaPlatformInfo() argument
39 device_allocator_(device_allocator) {} in XlaPlatformInfo()
H A Dxla_platform_info.cc182 options.device_allocator = GetAllocator(device, stream, platform_info); in GenerateCompilerOptions()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Dllvm_compiler_test.cc48 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloConvolutionCanonicalization() argument
54 se::DeviceMemoryAllocator* device_allocator) { in OptimizeHloPostLayoutAssignment() argument
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/linalg/
H A DBatchLinearAlgebraLib.cpp168 auto& device_allocator = *at::cuda::getCUDADeviceAllocator(); in apply_ldl_solve_cusolver()
169 auto workdata_device = device_allocator.allocate(worksize_device); in apply_ldl_solve_cusolver()
739 auto& device_allocator = *at::cuda::getCUDADeviceAllocator(); in apply_cholesky_cusolver_potrf_looped() local
740 auto workdata_device = device_allocator.allocate(worksize_device * batch_size); in apply_cholesky_cusolver_potrf_looped()
1014 auto& device_allocator = *at::cuda::getCUDADeviceAllocator(); in apply_geqrf() local
1015 auto work_device_data = device_allocator.allocate(worksize_device); in apply_geqrf()
1264 auto& device_allocator = *at::cuda::getCUDADeviceAllocator(); in apply_syevd() local
1265 auto work_device_data = device_allocator.allocate(worksize_device); in apply_syevd()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/
H A Dconvert_graph.cc697 std::pair<int, Allocator*> device_allocator = in CreateStaticEngine() local
701 if (device_allocator.first >= 0) { in CreateStaticEngine()
702 cuda_device_id = device_allocator.first; in CreateStaticEngine()
703 trt_allocator.reset(new TRTDeviceAllocator(device_allocator.second)); in CreateStaticEngine()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dexecutable_build_options.h74 se::DeviceMemoryAllocator* device_allocator() const;
H A Dexecutable_build_options.cc31 se::DeviceMemoryAllocator* ExecutableBuildOptions::device_allocator() const { in device_allocator() function in xla::ExecutableBuildOptions
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dxla_compiler.h194 std::shared_ptr<se::DeviceMemoryAllocator> device_allocator; member

12