Home
last modified time | relevance | path

Searched defs:device_guard (Results 1 – 25 of 47) sorted by relevance

12

/aosp_15_r20/external/pytorch/aten/src/ATen/test/
H A Dcuda_stream_test.cpp136 at::cuda::CUDAGuard device_guard(1); in TEST() local
207 at::cuda::CUDAGuard device_guard{1}; in TEST() local
303 at::cuda::CUDAGuard device_guard(0); in TEST() local
330 at::cuda::CUDAGuard device_guard(0); in TEST() local
334 at::cuda::CUDAGuard device_guard(1); in TEST() local
356 at::cuda::CUDAGuard device_guard(0); in TEST() local
398 at::cuda::CUDAGuard device_guard(0); in TEST() local
414 at::cuda::CUDAGuard device_guard(0); in TEST() local
422 at::cuda::CUDAGuard device_guard(0); in TEST() local
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DTensorIndexing.cpp69 OptionalDeviceGuard device_guard(device_of(*this)); in index() local
78 OptionalDeviceGuard device_guard(device_of(*this)); in index_put_() local
84 OptionalDeviceGuard device_guard(device_of(*this)); in index_put_() local
/aosp_15_r20/external/pytorch/torch/csrc/autograd/
H A Dpython_torch_functions_manual.cpp55 OptionalDeviceGuard device_guard(device_of(result)); in dispatch_range() local
66 DeviceGuard device_guard(options.device()); in dispatch_range() local
152 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_nonzero() local
158 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_nonzero() local
164 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_nonzero_numpy() local
H A Dinput_buffer.cpp199 c10::OptionalDeviceGuard device_guard{opt_sync_stream->device()}; in add() local
219 c10::OptionalDeviceGuard device_guard{device_of(var)}; in add() local
H A Dpython_variable_indexing.cpp342 OptionalDeviceGuard device_guard(device_of(self_)); in THPVariable_getitem() local
448 OptionalDeviceGuard device_guard(device_of(self_)); in THPVariable_setitem() local
/aosp_15_r20/external/pytorch/torch/csrc/utils/
H A Dtensor_new.cpp647 at::OptionalDeviceGuard device_guard(deviceOptional); in legacy_sparse_tensor_generic_ctor_new() local
659 at::OptionalDeviceGuard device_guard(deviceOptional); in legacy_sparse_tensor_generic_ctor_new() local
740 at::OptionalDeviceGuard device_guard(deviceOptional); in legacy_tensor_generic_ctor_new() local
955 at::OptionalDeviceGuard device_guard(r.deviceOptional(ARG_DEVICE)); in sparse_compressed_tensor_ctor_worker() local
1010 at::OptionalDeviceGuard device_guard(r.deviceOptional(ARG_DEVICE1)); in sparse_compressed_tensor_ctor_worker() local
1182 at::OptionalDeviceGuard device_guard(r.deviceOptional(ARG_DEVICE)); in sparse_coo_tensor_ctor() local
1216 at::OptionalDeviceGuard device_guard(r.deviceOptional(ARG_DEVICE1)); in sparse_coo_tensor_ctor() local
1249 at::OptionalDeviceGuard device_guard(r.deviceOptional(ARG_DEVICE2)); in sparse_coo_tensor_ctor() local
/aosp_15_r20/external/pytorch/tools/autograd/templates/
H A Dpython_variable_methods.cpp250 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_contiguous() local
294 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_copy_() local
320 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_to() local
383 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_invert() local
448 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_nonzero() local
454 OptionalDeviceGuard device_guard(device_of(self)); in dispatch_nonzero_numpy() local
953 OptionalDeviceGuard device_guard(device_of(self_)); in THPVariable_new() local
965 OptionalDeviceGuard device_guard(device_of(self_)); in THPVariable_new_tensor() local
/aosp_15_r20/external/pytorch/torch/csrc/cuda/
H A Dnccl.cpp588 at::cuda::OptionalCUDAGuard device_guard; in broadcast() local
640 at::cuda::OptionalCUDAGuard device_guard; in reduce() local
694 at::cuda::OptionalCUDAGuard device_guard; in all_reduce() local
736 at::cuda::OptionalCUDAGuard device_guard; in reduce_scatter() local
777 at::cuda::OptionalCUDAGuard device_guard; in all_gather() local
/aosp_15_r20/external/pytorch/c10/xpu/test/impl/
H A DXPUGuardTest.cpp19 const c10::DeviceGuard device_guard(device); in TEST() local
38 const c10::DeviceGuard device_guard(device); in TEST() local
/aosp_15_r20/external/pytorch/aten/src/ATen/cuda/
H A DEmptyTensor.cpp16 const DeviceGuard device_guard(device); in empty_cuda() local
56 const DeviceGuard device_guard(device); in empty_strided_cuda() local
H A DCUDAGraph.cpp227 c10::OptionalDeviceGuard device_guard{capture_stream_.device()}; in replay() local
/aosp_15_r20/external/pytorch/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/
H A DOpenRegMem.cpp73 const c10::DeviceGuard device_guard(device); in empty_openreg() local
91 const c10::DeviceGuard device_guard(device); in empty_strided_openreg() local
/aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/quantization/
H A Dquantization_gpu.cu60 at::cuda::OptionalCUDAGuard device_guard; in _float_to_bfloat16_cuda() local
111 at::cuda::OptionalCUDAGuard device_guard; in _bfloat16_to_float_cuda() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/
H A Dflash_api.cpp450 at::cuda::CUDAGuard device_guard{(char)q.get_device()}; in mha_fwd() local
689 at::cuda::CUDAGuard device_guard{(char)q.get_device()}; in mha_varlen_fwd() local
931 at::cuda::CUDAGuard device_guard{(char)q.get_device()}; in mha_bwd() local
1155 at::cuda::CUDAGuard device_guard{(char)q.get_device()}; in mha_varlen_bwd() local
1385 at::cuda::CUDAGuard device_guard{(char)q.get_device()}; in mha_fwd_kvcache() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DCopy.cu188 CUDAGuard device_guard(src_device); in copy_device_to_device() local
345 cuda::OptionalCUDAGuard device_guard; in copy_kernel_cuda() local
H A DLoops.cuh185 const OptionalDeviceGuard device_guard(iter.device(1)); in opmath_gpu_kernel_with_scalars() local
210 OptionalDeviceGuard device_guard; in opmath_symmetric_gpu_kernel_with_scalars() local
H A DForeachReduceOp.cu188 const at::cuda::OptionalCUDAGuard device_guard( in foreach_tensor_max_cuda() local
498 const at::cuda::OptionalCUDAGuard device_guard( in foreach_tensor_norm_cuda() local
H A DJitLoops.cuh174 const OptionalDeviceGuard device_guard(iter.device(1)); in opmath_jitted_gpu_kernel_with_scalars() local
/aosp_15_r20/external/pytorch/aten/src/ATen/cuda/detail/
H A DCUDAHooks.cpp132 at::OptionalDeviceGuard device_guard; in isPinnedPtr() local
458 at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index)); in deviceSynchronize() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/quantized/cuda/
H A DFusedObsFakeQuant.cu140 at::cuda::OptionalCUDAGuard device_guard; in _calculate_moving_average() local
196 at::cuda::OptionalCUDAGuard device_guard; in _calc_moving_avg_qparams_helper() local
H A DEmbeddingBag.cu214 at::cuda::OptionalCUDAGuard device_guard; in embedding_bag_byte_impl() local
396 at::cuda::OptionalCUDAGuard device_guard; in embedding_bag_4bit_impl() local
/aosp_15_r20/external/pytorch/torch/csrc/
H A DStorageSharing.cpp301 at::DeviceGuard device_guard(storage.device()); in THPStorage_shareCuda() local
462 at::cuda::CUDAGuard device_guard(device); in THPStorage_newSharedCuda() local
/aosp_15_r20/external/pytorch/torch/_inductor/codegen/xpu/
H A Ddevice_op_overrides.py15 def device_guard(self, device_idx): member in XPUDeviceOpOverrides
/aosp_15_r20/external/pytorch/torch/_inductor/codegen/cuda/
H A Ddevice_op_overrides.py15 def device_guard(self, device_idx): member in CUDADeviceOpOverrides
/aosp_15_r20/external/pytorch/aten/src/ATen/mps/
H A DEmptyTensor.cpp105 const DeviceGuard device_guard(device); in empty_strided_mps() local

12