Home
last modified time | relevance | path

Searched refs:canUse32BitIndexMath (Results 1 – 23 of 23) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DSorting.cu253 cuda::detail::canUse32BitIndexMath(self) && in launch_kthvalue_kernel()
254 cuda::detail::canUse32BitIndexMath(values) && in launch_kthvalue_kernel()
255 cuda::detail::canUse32BitIndexMath(indices) ? ScalarType::Int : ScalarType::Long, in launch_kthvalue_kernel()
268 if (cuda::detail::canUse32BitIndexMath(vals) && in launch_median_kernel()
269 cuda::detail::canUse32BitIndexMath(inds) && in launch_median_kernel()
270 cuda::detail::canUse32BitIndexMath(self)) { in launch_median_kernel()
H A DReflectionPad.cu31 using at::cuda::detail::canUse32BitIndexMath;
274 TORCH_CHECK(canUse32BitIndexMath(input_), in reflection_pad2d_out_template()
364 TORCH_CHECK(canUse32BitIndexMath(input), in reflection_pad2d_backward_out_template()
366 TORCH_CHECK(canUse32BitIndexMath(grad_output_), in reflection_pad2d_backward_out_template()
433 canUse32BitIndexMath(input_), in TORCH_IMPL_FUNC()
491 TORCH_CHECK(canUse32BitIndexMath(input), in TORCH_IMPL_FUNC()
494 TORCH_CHECK(canUse32BitIndexMath(grad_output_), in TORCH_IMPL_FUNC()
574 canUse32BitIndexMath(input_), in TORCH_IMPL_FUNC()
628 TORCH_CHECK(canUse32BitIndexMath(input), "input tensor must fit into 32-bit index math"); in TORCH_IMPL_FUNC()
629 TORCH_CHECK(canUse32BitIndexMath(grad_output), "input tensor must fit into 32-bit index math"); in TORCH_IMPL_FUNC()
H A DGridSampler.cu766 if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) && in launch_grid_sampler_2d_forward_kernel()
767 canUse32BitIndexMath(output)) { in launch_grid_sampler_2d_forward_kernel()
811 if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) && in launch_grid_sampler_3d_forward_kernel()
812 canUse32BitIndexMath(output)) { in launch_grid_sampler_3d_forward_kernel()
866 if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) && in launch_grid_sampler_2d_backward_kernel()
867 canUse32BitIndexMath(grad_output)) { in launch_grid_sampler_2d_backward_kernel()
925 if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) && in launch_grid_sampler_3d_backward_kernel()
926 canUse32BitIndexMath(grad_output)) { in launch_grid_sampler_3d_backward_kernel()
H A DNormalization.cu73 if (!at::cuda::detail::canUse32BitIndexMath(self)) { in batch_norm_choose_impl()
583 cuda::detail::canUse32BitIndexMath(input) && in batch_norm_backward_cuda()
584 cuda::detail::canUse32BitIndexMath(grad_out)) { in batch_norm_backward_cuda()
659 if (cuda::detail::canUse32BitIndexMath(self)) { in batch_norm_stats_cuda()
717 if (cuda::detail::canUse32BitIndexMath(self)) { in batch_norm_gather_stats_with_counts_cuda()
730 if (at::cuda::detail::canUse32BitIndexMath(grad_output) && in batch_norm_backward_reduce_cuda()
746 if (cuda::detail::canUse32BitIndexMath(grad_output)) { in batch_norm_backward_reduce_cuda()
767 if (at::cuda::detail::canUse32BitIndexMath(self) && in batch_norm_backward_elemt_cuda()
780 if (cuda::detail::canUse32BitIndexMath(self)) { in batch_norm_backward_elemt_cuda()
H A DDepthwiseConv2d.cu407 TORCH_CHECK(canUse32BitIndexMath(input) && canUse32BitIndexMath(output)); in conv_depthwise2d_forward_out()
494 TORCH_CHECK(canUse32BitIndexMath(grad_input) && in conv_depthwise2d_backward_out()
495 canUse32BitIndexMath(grad_output)); in conv_depthwise2d_backward_out()
622 TORCH_CHECK(canUse32BitIndexMath(input) && in conv_depthwise2d_grad_weight_out()
623 canUse32BitIndexMath(grad_output)); in conv_depthwise2d_grad_weight_out()
H A DReplicationPadding.cu235 TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(input), in replication_pad2d_backward_out_cuda_template()
237 TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(gradOutput), in replication_pad2d_backward_out_cuda_template()
310 TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(input), in shapeAndGradOutputCheck3d()
344 TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(gradOutput), in shapeAndGradOutputCheck3d()
537 TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(input), in TORCH_IMPL_FUNC()
H A DIndexing.cu918 if (cuda::detail::canUse32BitIndexMath(result) && in index_add_cuda_impl()
919 cuda::detail::canUse32BitIndexMath(source) && in index_add_cuda_impl()
920 cuda::detail::canUse32BitIndexMath(index)) { in index_add_cuda_impl()
1091 if (cuda::detail::canUse32BitIndexMath(result) && in index_reduce_func_cuda_impl()
1092 cuda::detail::canUse32BitIndexMath(source) && in index_reduce_func_cuda_impl()
1093 cuda::detail::canUse32BitIndexMath(index)) { in index_reduce_func_cuda_impl()
1387 if (cuda::detail::canUse32BitIndexMath(out) && in index_select_out_cuda_impl()
1388 cuda::detail::canUse32BitIndexMath(self) && in index_select_out_cuda_impl()
1389 cuda::detail::canUse32BitIndexMath(index)) { in index_select_out_cuda_impl()
H A DTriangularOps.cu125 if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) { in triu_tril_cuda_template()
H A DRNN.cu29 using at::cuda::detail::canUse32BitIndexMath;
539 if (canUse32BitIndexMath(workspace)) { // See Note [64-bit index math check elision] in _thnn_fused_lstm_cell_cuda()
590 if (canUse32BitIndexMath(workspace)) { // See Note [64-bit index math check elision] in _thnn_fused_lstm_cell_backward_impl_cuda()
624 if (canUse32BitIndexMath(workspace)) { // See Note [64-bit index math check elision] in _thnn_fused_gru_cell_cuda()
653 if (canUse32BitIndexMath(workspace)) { // See Note [64-bit index math check elision] in _thnn_fused_gru_cell_backward_cuda()
H A DShape.cu478 return at::cuda::detail::canUse32BitIndexMath(t); in TORCH_IMPL_FUNC()
493 at::cuda::detail::canUse32BitIndexMath(result) && in TORCH_IMPL_FUNC()
510 at::cuda::detail::canUse32BitIndexMath(result) && in TORCH_IMPL_FUNC()
H A DTensorModeKernel.cpp79 canUse32BitIndexMath(self)) { in mode_kernel_impl()
H A DTensorTopK.cu894 if (at::cuda::detail::canUse32BitIndexMath(input) && in launch_gather_topk_kernel()
895 at::cuda::detail::canUse32BitIndexMath(values) && in launch_gather_topk_kernel()
896 at::cuda::detail::canUse32BitIndexMath(indices)) { in launch_gather_topk_kernel()
H A DIndexKernel.cu326 …AT_DISPATCH_INDEX_TYPES(cuda::detail::canUse32BitIndexMath(output) ? ScalarType::Int : ScalarType:… in put_kernel()
351 …AT_DISPATCH_INDEX_TYPES(cuda::detail::canUse32BitIndexMath(input) ? ScalarType::Int : ScalarType::… in take_kernel()
H A DSort.cu311 if (at::cuda::detail::canUse32BitIndexMath(key)) { in sortCommon()
H A DDropout.cu358 if (cuda::detail::canUse32BitIndexMath(self)){ in dropout_cuda()
H A DNLLLoss2d.cu325 at::native::canUse32BitIndexMath(input_, INT_MAX) ? ScalarType::Int : ScalarType::Long, in nll_loss2d_forward_out_cuda_template()
H A DAveragePool2d.cu417 at::native::canUse32BitIndexMath(input, INT_MAX) ? ScalarType::Int : ScalarType::Long, in TORCH_IMPL_FUNC()
H A DSoftMax.cu929 at::native::canUse32BitIndexMath(input, INT_MAX) ? ScalarType::Int : ScalarType::Long, in host_softmax()
/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DGridSamplerUtils.h99 at::native::canUse32BitIndexMath(input) && in cond_cudnn_grid_sampler()
100 at::native::canUse32BitIndexMath(grid) && in cond_cudnn_grid_sampler()
H A DCanUse32BitIndexMath.h11 TORCH_API bool canUse32BitIndexMath(const at::TensorBase &t, int64_t max_elem=std::numeric_limits<i…
H A DIndexingUtils.cpp6 bool canUse32BitIndexMath(const TensorBase& t, int64_t max_elem) { in canUse32BitIndexMath() function
/aosp_15_r20/external/pytorch/aten/src/ATen/cuda/
H A DCUDAApplyUtils.cuh477 if (detail::canUse32BitIndexMath(a) && in CUDA_tensor_apply2()
478 detail::canUse32BitIndexMath(b)) { in CUDA_tensor_apply2()
/aosp_15_r20/external/pytorch/aten/src/ATen/cuda/detail/
H A DIndexUtils.cuh10 using at::native::canUse32BitIndexMath;